Qualcomm AI Engine Direct - Cast an int64 get_attr right before its first user

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",

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):