Qualcomm AI Engine Direct - Keep node names unique across the QNN and Samsung passes, and decompose constant-only ops on HTP

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/partition/qnn_partitioner.py

@@ -68,6 +68,15 @@
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
+
+def _is_runtime_input(
+    node: torch.fx.Node, edge_program: torch.export.ExportedProgram
+) -> bool:
+    return not is_parameter(node, edge_program) or is_mutable_buffer_input(
+        node, edge_program
+    )
+
+
 # dequantize ops with a CPU out variant; others cannot fall back
 CPU_DQ_OPS = {
     exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
@@ -319,17 +328,11 @@ def _drop_partitions_without_runtime_inputs(
         only by constants once a CPU op splits it from the rest. Such a
         partition stays on CPU.
         """
-
-        def is_runtime_input(node: torch.fx.Node) -> bool:
-            return not is_parameter(node, edge_program) or is_mutable_buffer_input(
-                node, edge_program
-            )
-
         return [
             partition
             for partition in partitions
             if any(
-                is_runtime_input(arg)
+                _is_runtime_input(arg, edge_program)
                 for node in partition.nodes
                 for arg in node.all_input_nodes
                 if arg not in partition.nodes
@@ -396,7 +399,37 @@ def ops_to_not_decompose(
         """
         Determines which op should not be decomposed during partitioning.
         The list of operators is obtained from `get_skip_decomp_table()`.
-        The filter function (`filter_fn`) can be used to further refine which nodes are not decomposed. (advanced use case)
+        The filter is `filter_fn`. On HTP, ops that read only constants are also
+        decomposed, except quantize and dequantize ops.
         """
         do_not_decompose = get_skip_decomp_table()
-        return (do_not_decompose, filter_fn)
+        # On the lowered program this filter tells exir which kept-whole ops
+        # may stay on CPU, and none may, so it stays plain filter_fn there.
+        if self.backend != QnnExecuTorchBackendType.kHtpBackend or ep.dialect == "EDGE":
+            return (do_not_decompose, filter_fn)
+        # An op that reads only constants can end up in a partition that
+        # _drop_partitions_without_runtime_inputs keeps on CPU after
+        # decomposition, so every such op is decomposed. Partitions are not
+        # known yet, so this also decomposes ones whose partition stays
+        # delegated; folding constant-only subgraphs first would avoid that.
+        constant_only = set()
+        for node in ep.graph.nodes:
+            if node.op == "placeholder":
+                if not _is_runtime_input(node, ep):
+                    constant_only.add(node)
+            elif node.op == "call_function" and all(
+                arg in constant_only for arg in node.all_input_nodes
+            ):
+                constant_only.add(node)
+        # Quantize and dequantize ops stay whole: QNN builds weight encodings
+        # from them.
+        quant_ops = {
+            op
+            for op in do_not_decompose
+            if op.namespace in ("quantized_decomposed", "torchao")
+        }
+        return (
+            do_not_decompose,
+            lambda node: (node not in constant_only or node.target in quant_ops)
+            and filter_fn(node),
+        )

backends/qualcomm/tests/test_passes.py

@@ -31,7 +31,10 @@
 )
 from executorch.backends.qualcomm.builders.qnn_constants import OpContextLoader
 from executorch.backends.qualcomm.builders.utils import is_parameter
-from executorch.backends.qualcomm.partition.qnn_partitioner import QnnOperatorSupport
+from executorch.backends.qualcomm.partition.qnn_partitioner import (
+    QnnOperatorSupport,
+    QnnPartitioner,
+)
 from executorch.backends.qualcomm.qnn_preprocess import QnnBackend
 from executorch.backends.qualcomm.quantizer.quantizer import QnnQuantizer, QuantDtype
 from executorch.backends.qualcomm.serialization.qc_schema import (
@@ -188,6 +191,74 @@ 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 q * 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):
+            def forward(self, q, k, v, mul):
+                return (
+                    torch.nn.functional.scaled_dot_product_attention(q * 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()
+                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"
@@ -946,16 +1017,19 @@ class ConstantBranch(torch.nn.Module):
             def __init__(self):
                 super().__init__()
                 self.register_buffer("table", torch.randn(4, 16))
+                self.register_buffer("weight", torch.randn(16, 16))
 
             def forward(self, x):
-                return x + torch.relu(self.table * 2 - 1)
+                table = torch.nn.functional.linear(self.table, self.weight)
+                return x + torch.relu(table * 2 - 1)
 
         sample_input = (torch.randn(4, 16),)
         compiler_specs = generate_qnn_executorch_compiler_spec(
             soc_model=QcomChipset.SM8650,
             backend_options=generate_htp_compiler_spec(use_fp16=True),
         )
-        # With relu on CPU, table * 2 - 1 reads nothing at runtime.
+        # With relu on CPU, the branch before it reads nothing at runtime. It
+        # also holds linear, which QNN otherwise keeps from decomposing.
         try:
             edge = to_edge_transform_and_lower_to_qnn(
                 ConstantBranch().eval(),
@@ -967,8 +1041,86 @@ 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.assertNotIn(exir_ops.edge.aten.linear.default, top_level_ops)
+        self.assertIn(torch.ops.higher_order.executorch_call_delegate, top_level_ops)
         edge.to_executorch()
 
+    def test_htp_keeps_whole_only_ops_that_read_a_runtime_input(self):
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("table", torch.randn(4, 16))
+                self.register_buffer("weight", torch.randn(16, 16))
+                self.register_buffer(
+                    "qweight", torch.randint(-8, 8, (16, 32), dtype=torch.int8)
+                )
+                self.register_buffer("scales", torch.rand(16, 2) + 0.1)
+                self.register_buffer("zero_points", torch.zeros(16, 2).long())
+
+            def forward(self, x, y):
+                weight = torch.ops.quantized_decomposed.dequantize_per_channel_group(
+                    self.qweight,
+                    self.scales,
+                    self.zero_points,
+                    -8,
+                    7,
+                    torch.int8,
+                    16,
+                    torch.float32,
+                )
+                table = torch.nn.functional.linear(self.table, self.weight)
+                return torch.nn.functional.linear(x, self.weight) + table, (
+                    torch.nn.functional.linear(y, weight)
+                )
+
+        ep = torch.export.export(Model(), (torch.randn(4, 16), torch.randn(4, 32)))
+        compiler_specs = generate_qnn_executorch_compiler_spec(
+            soc_model=QcomChipset.SM8650,
+            backend_options=generate_htp_compiler_spec(use_fp16=True),
+        )
+        _, keep_whole = QnnPartitioner(compiler_specs).ops_to_not_decompose(ep)
+        linears = {
+            node.args[0].name: node
+            for node in ep.graph.nodes
+            if node.target == torch.ops.aten.linear.default
+        }
+        placeholders = {b: p for p, b in ep.graph_signature.inputs_to_buffers.items()}
+        self.assertFalse(keep_whole(linears[placeholders["table"]]))
+        self.assertTrue(keep_whole(linears["x"]))
+        (dequantize,) = (
+            node
+            for node in ep.graph.nodes
+            if node.target
+            == torch.ops.quantized_decomposed.dequantize_per_channel_group.default
+        )
+        self.assertTrue(keep_whole(dequantize))
+        # After lowering, exir uses the filter to report a kept-whole op left
+        # on CPU, so there it must not let a constant-only one through.
+        edge = to_edge(
+            ep,
+            compile_config=EdgeCompileConfig(
+                _check_ir_validity=False,
+                preserve_ops=[torch.ops.aten.linear.default],
+            ),
+        ).exported_program()
+        _, keep_whole = QnnPartitioner(compiler_specs).ops_to_not_decompose(edge)
+        placeholders = {b: p for p, b in edge.graph_signature.inputs_to_buffers.items()}
+        (constant_linear,) = (
+            node
+            for node in edge.graph.nodes
+            if node.target == exir_ops.edge.aten.linear.default
+            and node.args[0].name == placeholders["table"]
+        )
+        self.assertTrue(keep_whole(constant_linear))
+
     def test_partitioner_falls_back_on_bfloat16(self):
         class Bf16Region(torch.nn.Module):
             def __init__(self):

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