diff --git a/backends/qualcomm/_passes/BUCK b/backends/qualcomm/_passes/BUCK index f5e635c739a..6d46f1972e1 100644 --- a/backends/qualcomm/_passes/BUCK +++ b/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", diff --git a/backends/qualcomm/_passes/qnn_pass_manager.py b/backends/qualcomm/_passes/qnn_pass_manager.py index e8745524512..4c8d7e0cd8e 100644 --- a/backends/qualcomm/_passes/qnn_pass_manager.py +++ b/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): diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 58ce89a6e30..2c1818f8a17 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/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" diff --git a/backends/samsung/_passes/enn_pass_manager.py b/backends/samsung/_passes/enn_pass_manager.py index 874aa7115c9..d148c0a57d5 100644 --- a/backends/samsung/_passes/enn_pass_manager.py +++ b/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 diff --git a/backends/transforms/utils.py b/backends/transforms/utils.py index 92794364cb0..314dd8100d5 100644 --- a/backends/transforms/utils.py +++ b/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