diff --git a/backends/qualcomm/_passes/i64_to_i32.py b/backends/qualcomm/_passes/i64_to_i32.py index a26956b5569..d3ab8e99284 100644 --- a/backends/qualcomm/_passes/i64_to_i32.py +++ b/backends/qualcomm/_passes/i64_to_i32.py @@ -143,7 +143,18 @@ def _cast_constant_to_int32(self, graph_module: torch.fx.GraphModule): # If the scalar operand has an int64 data type, it should be explicitly cast to int32. node_val = n.meta["val"] if self._is_tensor_of_dtype(node_val, torch.int64): - with graph_module.graph.inserting_after(n): + # A get_attr may sit above the graph inputs (create_const_node puts it + # first). Once lifted it becomes a placeholder after them, so cast + # right before its first user rather than right after it. + first_user = next( + (u for u in graph_module.graph.nodes if u in n.users), None + ) + insert_point = ( + graph_module.graph.inserting_before(first_user) + if n.op == "get_attr" and first_user is not None + else graph_module.graph.inserting_after(n) + ) + with insert_point: args = (n,) to_dst_node = graph_module.graph.create_node( "call_function", diff --git a/backends/qualcomm/tests/rework/src/pattern.py b/backends/qualcomm/tests/rework/src/pattern.py index 7b90fde61e2..f47a3581aa4 100644 --- a/backends/qualcomm/tests/rework/src/pattern.py +++ b/backends/qualcomm/tests/rework/src/pattern.py @@ -4209,9 +4209,15 @@ def forward(self, x, y): squeeze_out = view_out.squeeze(-1) return squeeze_out, conv_out + class _RemainderByScalar(torch.nn.Module): + # DecomposeRemainder turns the scalar into an int64 get_attr at the graph head + def forward(self, x, idx): + return x + (idx % 80).float() + @staticmethod @unpack_pass_fixtures def test( + subtests, quantizer, compile_spec, backend_type: QnnExecuTorchBackendType, @@ -4237,6 +4243,18 @@ def test( gm, exir_ops.edge.aten._to_copy.default, expected_casts ) + # The int32 cast of a constant must still follow it once the constant is lifted + with subtests.test(msg="constant_cast_order"): + gm = pass_pipeline.lower_edge_ep( + module=I64toI32._RemainderByScalar(), + sample_input=(torch.randn(1, 3), torch.tensor([[3, 170, 85]])), + backend_type=backend_type, + compile_spec=compile_spec, + target_pass=target_pass, + quantizer=quantizer, + ).graph_module + gm.graph.lint() + class InsertIOQDQ: class _Basic(torch.nn.Module):