Qualcomm AI Engine Direct - Keep node names unique across QNN passes

backends/qualcomm/_passes/qnn_pass_manager.py

@@ -100,6 +100,12 @@
 class QnnPassManager(PassManager):
 
     def _transform(self, graph_module: GraphModule):
+        # TODO: Remove once torch.export records the node names it assigns.
+        # Until then a pass could reuse one, and since torch.fx cannot reserve
+        # a name publicly, this writes the namespace's private name set.
+        namespace = graph_module.graph._graph_namespace
+        for node in graph_module.graph.nodes:
+            namespace._used_names.add(node.name)
         return self(graph_module).graph_module
 
     def _reset(self):

backends/qualcomm/tests/test_passes.py

@@ -188,6 +188,55 @@ def forward(self, x):
                 )
                 torch.testing.assert_close(gm(x), model(x))
 
+    def test_export_pipeline_keeps_node_names_unique(self):
+        # torch.export names the inputs t_1 and mul, which the autocast and
+        # attention decompositions would otherwise give to new nodes too.
+        class Autocast(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.w = torch.nn.Parameter(torch.randn(8, 8))
+
+            def forward(self, x, t_1):
+                with torch.autocast("cpu", dtype=torch.bfloat16):
+                    y = x @ self.w.t() @ self.w.t()
+                return y.float() + t_1
+
+        class Attention(torch.nn.Module):
+            def forward(self, q, k, v, mul):
+                attention = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+                return attention + mul
+
+        cases = (
+            (Autocast(), (torch.randn(4, 8), torch.randn(4, 8)), "t_1"),
+            (Attention(), tuple(torch.randn(1, 2, 4, 8) for _ in range(4)), "mul"),
+        )
+        for model, inputs, name in cases:
+            with self.subTest(model=type(model).__name__):
+                ep = torch.export.export(model.eval(), inputs)
+                # Guard against a vacuous test: the colliding name has to be
+                # an input's, and torch.export must leave it unrecorded.
+                self.assertIn(
+                    name, [n.name for n in ep.graph.nodes if n.op == "placeholder"]
+                )
+                self.assertNotIn(
+                    name,
+                    ep.graph._graph_namespace._used_names,
+                    "torch.export now records the names it assigns, so the "
+                    "reservation in QnnPassManager._transform can go",
+                )
+                ep = get_qnn_pass_manager_cls()().transform_for_export_pipeline(ep)
+                # The decompositions that add the colliding nodes have run.
+                sdpa = torch.ops.aten.scaled_dot_product_attention.default
+                for n in ep.graph.nodes:
+                    self.assertNotIsInstance(
+                        n.target, torch._higher_order_ops.wrap.WrapWithAutocast
+                    )
+                    self.assertNotEqual(n.target, sdpa)
+                # ep.module() copies the graph and renames duplicates, so lint
+                # the graph the passes produced.
+                ep.graph_module.graph.lint()
+                torch.testing.assert_close(ep.module()(*inputs), model(*inputs))
+
     def test_build_op_wrappers_returns_context_binary(self):
         op_name = "ctx_loader_build"
         ctx_bin = b"qnn_context_binary"
@@ -967,6 +1016,15 @@ def forward(self, x):
             if "Failed to initialize QNN backend" in str(e):
                 self.skipTest(f"QNN SDK not available: {e}")
             raise
+        top_level_ops = {
+            node.target
+            for node in edge.exported_program().graph.nodes
+            if node.op == "call_function"
+        }
+        self.assertIn(exir_ops.edge.aten.mul.Tensor, top_level_ops)
+        self.assertIn(exir_ops.edge.aten.sub.Tensor, top_level_ops)
+        self.assertNotIn(exir_ops.edge.aten.add.Tensor, top_level_ops)
+        self.assertIn(torch.ops.higher_order.executorch_call_delegate, top_level_ops)
         edge.to_executorch()
 
     def test_partitioner_falls_back_on_bfloat16(self):