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

backends/qualcomm/_passes/BUCK

@@ -14,6 +14,7 @@ fbcode_target(_kind = runtime.python_library,
         "//executorch/backends/transforms:addmm_mm_to_linear",
         "//executorch/backends/transforms:decompose_sdpa",
         "//executorch/backends/transforms:fuse_batch_norm_with_conv",
+        "//executorch/backends/transforms:utils",
         "//executorch/exir/backend:backend_details",
         "//executorch/exir/backend:compile_spec_schema",
         "//executorch/backends/qualcomm/quantizer:quantizer",

backends/qualcomm/_passes/qnn_pass_manager.py

@@ -87,6 +87,7 @@
 from executorch.backends.transforms.decompose_sdpa import (
     DecomposeScaledDotProductAttention,
 )
+from executorch.backends.transforms.utils import reserve_node_names
 from executorch.exir import ExportedProgram
 from executorch.exir.pass_manager import PassManager
 from executorch.exir.program._program import (
@@ -100,6 +101,7 @@
 class QnnPassManager(PassManager):
 
     def _transform(self, graph_module: GraphModule):
+        reserve_node_names(graph_module)
         return self(graph_module).graph_module
 
     def _reset(self):

backends/qualcomm/tests/test_passes.py

@@ -188,6 +188,81 @@ def forward(self, x):
                 )
                 torch.testing.assert_close(gm(x), model(x))
 
+    def test_export_pipeline_keeps_node_names_unique(self):
+        # torch.export leaves the input name t_1 unrecorded, and renames
+        # input * 2 to mul_1 to make room for the input mul without recording
+        # that either. The decompositions would otherwise reuse those names.
+        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):
+                    # The second transpose is named t_1 inside the region, and
+                    # the decomposition copies it next to the input t_1.
+                    y = x @ self.w.t() @ self.w.t()
+                return y.float() + t_1
+
+        class Attention(torch.nn.Module):
+            # input is a Python builtin name. It has to survive the reservation
+            # as it is, because the graph signature refers to inputs by name.
+            def forward(self, input, k, v, mul):
+                return (
+                    torch.nn.functional.scaled_dot_product_attention(input * 2, k, v)
+                    + mul
+                )
+
+        cases = (
+            (
+                Autocast(),
+                (torch.randn(4, 8), torch.randn(4, 8)),
+                torch._higher_order_ops.wrap.wrap_with_autocast,
+            ),
+            (
+                Attention(),
+                tuple(torch.randn(1, 2, 4, 8) for _ in range(4)),
+                torch.ops.aten.scaled_dot_product_attention.default,
+            ),
+        )
+        no_collision = []
+        for model, inputs, decomposed_op in cases:
+            with self.subTest(model=type(model).__name__):
+                ep = torch.export.export(model.eval(), inputs)
+                self.assertIn(decomposed_op, [n.target for n in ep.graph.nodes])
+                ep = get_qnn_pass_manager_cls()().transform_for_export_pipeline(ep)
+                self.assertNotIn(decomposed_op, [n.target for n in ep.graph.nodes])
+                # ep.module() copies the graph and renames duplicates, so lint
+                # the graph the passes produced.
+                ep.graph.lint()
+                self.assertEqual(
+                    [n.name for n in ep.graph.nodes if n.op == "placeholder"],
+                    [spec.arg.name for spec in ep.graph_signature.input_specs],
+                )
+                torch.testing.assert_close(ep.module()(*inputs), model(*inputs))
+                ep = torch.export.export(model.eval(), inputs)
+                with patch(
+                    "executorch.backends.qualcomm._passes.qnn_pass_manager."
+                    "reserve_node_names"
+                ) as reserve:
+                    try:
+                        get_qnn_pass_manager_cls()().transform_for_export_pipeline(
+                            ep
+                        ).graph.lint()
+                    except Exception as e:
+                        # The pass manager wraps the lint error in its own.
+                        self.assertIn("Node redefined name", str(e.__cause__ or e))
+                    else:
+                        no_collision.append(type(model).__name__)
+                    reserve.assert_called_once()
+        # A skip, not a failure: a torch update that removes the collision is fine.
+        if no_collision:
+            self.skipTest(
+                f"{no_collision} no longer collide without reserve_node_names, so "
+                "those cases no longer test it; see "
+                "https://github.com/pytorch/pytorch/issues/202927"
+            )
+
     def test_build_op_wrappers_returns_context_binary(self):
         op_name = "ctx_loader_build"
         ctx_bin = b"qnn_context_binary"

backends/samsung/_passes/enn_pass_manager.py

@@ -33,6 +33,7 @@
 )
 from executorch.backends.transforms.remove_clone_ops import RemoveCloneOpsTransform
 from executorch.backends.transforms.remove_getitem_op import RemoveGetItemPass
+from executorch.backends.transforms.utils import reserve_node_names
 
 from executorch.exir import ExportedProgram
 from executorch.exir.pass_manager import PassManager
@@ -44,6 +45,7 @@ def __init__(self) -> None:
         super().__init__()
 
     def _transform(self, graph_module: GraphModule):
+        reserve_node_names(graph_module)
         return self(graph_module).graph_module
 
     # before annotator, for quant models

backends/transforms/utils.py

@@ -59,6 +59,19 @@ def _get_fake_tensor_mode(graph: torch.fx.Graph, data: torch.Tensor) -> torch.Te
     return FakeTensorConverter().from_real_tensor(example_fake_tensor.fake_mode, t=data)
 
 
+def reserve_node_names(graph_module: torch.fx.GraphModule) -> None:
+    """
+    Records every node name of the top-level graph in its namespace, so that
+    nodes a pass adds there later get names that are not already taken.
+    """
+    # TODO: Remove once torch.export records the names it gives inputs and the
+    # ops it renames for them (https://github.com/pytorch/pytorch/issues/202927).
+    # torch.fx has no public way to reserve a name.
+    namespace = graph_module.graph._graph_namespace
+    for node in graph_module.graph.nodes:
+        namespace._used_names.add(node.name)
+
+
 def is_get_attr_node(node: torch.fx.Node) -> bool:
     """
     Returns true if the given node is a get attr node for a tensor of the model