Partitioner: run backend-contributed edge transform passes

backends/qualcomm/partition/qnn_partitioner.py

@@ -5,10 +5,13 @@
 # LICENSE file in the root directory of this source tree.
 import copy
 import logging
-from collections import defaultdict
+from collections import defaultdict, OrderedDict
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import torch
+from executorch.backends.qualcomm._passes.qnn_pass_manager import (
+    get_qnn_pass_manager_cls,
+)
 from executorch.backends.qualcomm.builders import node_visitor_manager
 from executorch.backends.qualcomm.builders.node_visitor import (
     PER_CHANNEL_GROUP_ENCODING,
@@ -41,6 +44,7 @@
     PartitionResult,
 )
 from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer
+from executorch.exir.pass_manager import PassType
 from torch.export.exported_program import ExportedProgram
 from torch.fx.passes.infra.partitioner import Partition
 from torch.fx.passes.operator_support import OperatorSupportBase
@@ -199,14 +203,25 @@ def __init__(
         skip_node_id_set: set = None,
         skip_node_op_set: set = None,
         skip_mutable_buffer: bool = False,
+        dep_table: Optional[Dict] = None,
+        passes_job: Optional[OrderedDict] = None,
+        convert_linear_to_conv2d: bool = False,
     ):
         """
         Args:
             compiler_specs (List[CompileSpec]): Backend compiler specifications.
             skip_node_id_set (set, optional): Set of node IDs to exclude from partitioning.
             skip_node_op_set (set, optional): Set of OpOverload to exclude from partitioning.
             skip_mutable_buffer (bool, optional): If True, mutable buffers are not delegated to QNN.
+            dep_table (Dict, optional): Dependency table for the transformation passes.
+            passes_job (OrderedDict, optional): Ordered dictionary of transformation passes.
+            convert_linear_to_conv2d (bool, optional): Convert linear to conv2d where that
+                improves HTP performance.
         """
+        self.compiler_specs = compiler_specs
+        self.dep_table = dep_table
+        self.passes_job = passes_job
+        self.convert_linear_to_conv2d = convert_linear_to_conv2d
         self.compiler_specs_snapshot = copy.deepcopy(compiler_specs)
         self.backend = flatbuffer_to_option(
             generate_qnn_executorch_option(self.compiler_specs_snapshot)
@@ -220,6 +235,26 @@ def __init__(
         self.skip_node_op_set = set() if skip_node_op_set is None else skip_node_op_set
         self.skip_mutable_buffer = skip_mutable_buffer
 
+    def transform_for_pre_decomposition(
+        self, exported_program: ExportedProgram
+    ) -> ExportedProgram:
+        """Run the QNN ATen-level passes before default decompositions."""
+        return get_qnn_pass_manager_cls(self.backend)().transform_for_export_pipeline(
+            exported_program
+        )
+
+    def transform_passes(self, exported_program: ExportedProgram) -> List[PassType]:
+        """Run the QNN edge-level passes before partitioning."""
+        return get_qnn_pass_manager_cls(self.backend)().get_to_edge_transform_passes(
+            exported_program,
+            passes_job=self.passes_job,
+            dep_table=self.dep_table,
+            compiler_specs=self.compiler_specs,
+            skip_node_id_set=self.skip_node_id_set,
+            skip_node_op_set=self.skip_node_op_set,
+            convert_linear_to_conv2d=self.convert_linear_to_conv2d,
+        )
+
     def generate_partitions(
         self, edge_program: torch.export.ExportedProgram
     ) -> List[Any]:

backends/qualcomm/tests/test_passes.py

@@ -28,7 +28,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 (
@@ -850,6 +853,70 @@ def forward(self, x):
                 self.skipTest(f"QNN SDK not available: {e}")
             raise
 
+    def test_partitioner_carries_qnn_edge_pass_pipeline(self):
+        """QnnPartitioner.transform_passes must reproduce the edge pass list that
+        to_edge_transform_and_lower_to_qnn previously built by hand, so a plain
+        to_edge_transform_and_lower(partitioner=[QnnPartitioner(...)]) lowers
+        through the same pipeline.
+        """
+
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(4, 4)
+
+            def forward(self, x):
+                return torch.nn.functional.relu(self.linear(x))
+
+        sample_input = (torch.randn(1, 4),)
+        exported = torch.export.export(Model().eval(), sample_input, strict=True)
+        compiler_specs = generate_qnn_executorch_compiler_spec(
+            soc_model=QcomChipset.SM8650,
+            backend_options=generate_htp_compiler_spec(use_fp16=True),
+        )
+
+        partitioner = QnnPartitioner(compiler_specs)
+        expected = get_qnn_pass_manager_cls(
+            partitioner.backend
+        )().get_to_edge_transform_passes(exported, compiler_specs=compiler_specs)
+        actual = partitioner.transform_passes(exported)
+
+        self.assertTrue(expected, "expected a non-empty QNN edge pass pipeline")
+        self.assertEqual([type(p) for p in actual], [type(p) for p in expected])
+
+    def test_partitioner_runs_qnn_aten_passes_pre_decomposition(self):
+        """transform_for_pre_decomposition must run the QNN ATen pipeline.
+
+        DecomposeScaledDotProductAttention is part of that pipeline, so an SDPA
+        node must be gone after the hook runs.
+        """
+
+        class Sdpa(torch.nn.Module):
+            def forward(self, q, k, v):
+                return torch.nn.functional.scaled_dot_product_attention(q, k, v)
+
+        sample_input = tuple(torch.randn(1, 2, 4, 8) for _ in range(3))
+        exported = torch.export.export(Sdpa().eval(), sample_input, strict=True)
+
+        def has_sdpa(program):
+            return any(
+                node.op == "call_function"
+                and "scaled_dot_product_attention" in str(node.target)
+                for node in program.graph.nodes
+            )
+
+        # Guard against a vacuous test.
+        self.assertTrue(has_sdpa(exported), "expected SDPA in the traced graph")
+
+        compiler_specs = generate_qnn_executorch_compiler_spec(
+            soc_model=QcomChipset.SM8650,
+            backend_options=generate_htp_compiler_spec(use_fp16=True),
+        )
+        transformed = QnnPartitioner(compiler_specs).transform_for_pre_decomposition(
+            exported
+        )
+        self.assertFalse(has_sdpa(transformed))
+
     def test_index_put_int64_value_not_quantized(self):
         """QNN's IndexPut annotator must skip a non-float (int64) value arg.
 

backends/qualcomm/utils/utils.py

@@ -462,16 +462,21 @@ def ensure_graph_specific_dict(value, graph_names):
     dep_table = ensure_graph_specific_dict(dep_table, graph_names)
     passes_job = ensure_graph_specific_dict(passes_job, graph_names)
 
-    # Prepare programs and partitioners
+    # Prepare programs and partitioners. The partitioner carries the QNN pass
+    # pipeline (transform_for_pre_decomposition / transform_passes), so lowering
+    # runs the same passes whether it is driven from here or from a plain
+    # to_edge_transform_and_lower call.
     aten_programs = {}
-    transform_passes = {}
     qnn_partitioners = {
         graph_name: [
             QnnPartitioner(
                 compiler_specs[graph_name],
                 skip_node_id_set=skip_node_id_set,
                 skip_node_op_set=skip_node_op_set,
                 skip_mutable_buffer=skip_mutable_buffer,
+                dep_table=dep_table[graph_name],
+                passes_job=passes_job[graph_name],
+                convert_linear_to_conv2d=convert_linear_to_conv2d,
             )
         ]
         for graph_name in graph_names
@@ -503,31 +508,10 @@ def ensure_graph_specific_dict(value, graph_names):
                 dynamic_shapes=dynamic_shapes[graph_name],
                 strict=True,
             )
-        option = generate_qnn_executorch_option(compiler_specs[graph_name])
-        python_options = flatbuffer_to_option(option)
-        backend_type = python_options.backend_options.backend_type
-        pass_manager = get_qnn_pass_manager_cls(backend_type)()
-        # This transformation is primarily intended for the LiftConstantScalarOperands pass
-        # to avoid creating temporary tensors in the operation builder.
-        # However, this pass will create a get_attr node, which should be converted
-        # into a lifted tensor constant by the lift_constant_tensor_pass.
-        # If placed in the to_edge_transform_passes, it will be executed
-        # after the lift_constant_tensor_pass, causing the operation builder
-        # to fail to correctly retrieve the parameter by the get_parameter.
-        aten_programs[graph_name] = pass_manager.transform_for_export_pipeline(ep)
-        transform_passes[graph_name] = pass_manager.get_to_edge_transform_passes(
-            ep,
-            passes_job=passes_job[graph_name],
-            dep_table=dep_table[graph_name],
-            compiler_specs=compiler_specs[graph_name],
-            skip_node_id_set=skip_node_id_set,
-            skip_node_op_set=skip_node_op_set,
-            convert_linear_to_conv2d=convert_linear_to_conv2d,
-        )
+        aten_programs[graph_name] = ep
     with QnnManagerContext(compiler_specs):
         return to_edge_transform_and_lower(
             aten_programs,
-            transform_passes=transform_passes,
             partitioner=qnn_partitioners,
             constant_methods=constant_methods,
             compile_config=qnn_edge_config(),

exir/backend/partitioner.py

@@ -14,6 +14,7 @@
 
 from executorch.exir.backend.backend_details import enforcedmethod
 from executorch.exir.backend.compile_spec_schema import CompileSpec
+from executorch.exir.pass_manager import PassType
 from torch.export import ExportedProgram
 
 
@@ -114,6 +115,31 @@ def ops_to_not_decompose(
         """
         return ([], None)
 
+    def transform_passes(self, exported_program: ExportedProgram) -> List[PassType]:
+        """Return passes to run on the edge program before partitioning.
+
+        EXIR invokes this hook automatically from
+        ``to_edge_transform_and_lower`` and runs the returned passes in the
+        same slot as its ``transform_passes`` argument, after any passes the
+        caller supplied there. Backend partitioners may override it to carry
+        edge-dialect rewrites that must run before ``partition`` sees the
+        graph, so callers do not have to thread backend-specific pass lists
+        through every lowering call.
+
+        When multiple partitioners are used for a method, their passes are
+        concatenated in the order supplied to ``to_edge_transform_and_lower``.
+
+        Args:
+            exported_program (ExportedProgram): The ATen-dialect program being
+                lowered, after ``transform_for_pre_decomposition`` has run and
+                before default decompositions. Supplied because pass
+                construction may depend on the graph.
+
+        Returns:
+            List[PassType]: Passes to run, in order. Empty by default.
+        """
+        return []
+
     def transform_for_pre_decomposition(
         self, exported_program: ExportedProgram
     ) -> ExportedProgram:

exir/program/_program.py

@@ -1131,12 +1131,13 @@ def _apply_pre_decomposition_transforms(
     return program
 
 
-def _gen_edge_manager_for_partitioners(
+def _gen_edge_manager_for_partitioners(  # noqa: C901
     partitioner: Dict[str, List[Partitioner]],
     aten_programs: Dict[str, ExportedProgram],
     config: EdgeCompileConfig,
     constant_methods: Optional[Dict[str, Any]],
     generate_etrecord: Optional[bool] = False,
+    collected_transform_passes: Optional[Dict[str, List[PassType]]] = None,
 ) -> "EdgeProgramManager":
     """
     Generates EdgeProgramManager for subsequent lowering to the
@@ -1157,6 +1158,16 @@ def _gen_edge_manager_for_partitioners(
         partitioners_for_program = partitioner.get(name, [])
         program = _apply_pre_decomposition_transforms(program, partitioners_for_program)
 
+        # Ask partitioners for edge passes here, while the program still matches
+        # what they saw in transform_for_pre_decomposition: pass construction can
+        # depend on the graph.
+        if collected_transform_passes is not None:
+            passes_for_program = []
+            for curr_partitioner in partitioners_for_program:
+                passes_for_program.extend(curr_partitioner.transform_passes(program))
+            if passes_for_program:
+                collected_transform_passes[name] = passes_for_program
+
         # Functionalize program before asking partitioners to preserve ops
         program = program.run_decompositions({})
 
@@ -1285,6 +1296,46 @@ def collect_named_data_store_outputs(
     collect_named_data_store_outputs(exported_program.graph_module)
 
 
+def _merge_partitioner_transform_passes(
+    transform_passes: Optional[
+        Union[
+            Sequence[PassType],
+            Dict[str, Sequence[PassType]],
+            GraphModulePassManager,
+            ExportedProgramPassManager,
+        ]
+    ],
+    partitioner_passes: Dict[str, List[PassType]],
+    method_names: Sequence[str],
+) -> Dict[str, Sequence[PassType]]:
+    """Fold passes contributed by partitioners into ``transform_passes``.
+
+    Caller-supplied passes run first, so an explicit argument is never
+    reordered behind a backend's own passes.
+    """
+    if isinstance(
+        transform_passes, (GraphModulePassManager, ExportedProgramPassManager)
+    ):
+        raise ValueError(
+            "transform_passes was given as a PassManager, which cannot be "
+            "combined with passes contributed by a partitioner's "
+            "transform_passes(). Pass a list, or a dict of lists, instead."
+        )
+
+    merged: Dict[str, Sequence[PassType]] = {}
+    for name in method_names:
+        if isinstance(transform_passes, dict):
+            caller_passes = transform_passes.get(name, [])
+        elif transform_passes is not None:
+            caller_passes = transform_passes
+        else:
+            caller_passes = []
+        combined = [*caller_passes, *partitioner_passes.get(name, [])]
+        if combined:
+            merged[name] = combined
+    return merged
+
+
 @et_logger("to_edge_transform_and_lower")
 def to_edge_transform_and_lower(  # noqa: C901
     programs: Union[ExportedProgram, Dict[str, ExportedProgram]],
@@ -1363,10 +1414,21 @@ def to_edge_transform_and_lower(  # noqa: C901
     elif partitioner is None:
         partitioner = {name: [] for name in aten_programs.keys()}
 
+    partitioner_transform_passes: Dict[str, List[PassType]] = {}
     edge_manager = _gen_edge_manager_for_partitioners(
-        partitioner, aten_programs, config, constant_methods, generate_etrecord
+        partitioner,
+        aten_programs,
+        config,
+        constant_methods,
+        generate_etrecord,
+        collected_transform_passes=partitioner_transform_passes,
     )
 
+    if partitioner_transform_passes:
+        transform_passes = _merge_partitioner_transform_passes(
+            transform_passes, partitioner_transform_passes, tuple(aten_programs)
+        )
+
     if transform_passes is not None:
         edge_manager = edge_manager.transform(transform_passes)
         for method in edge_manager.methods:

exir/tests/targets.bzl

@@ -265,6 +265,20 @@ def define_common_targets(is_fbcode = False):
         ],
     )
 
+    python_unittest(
+        name = "test_partitioner_transform_passes",
+        srcs = [
+            "test_partitioner_transform_passes.py",
+        ],
+        deps = [
+            "//caffe2:torch",
+            "//executorch/exir:lib",
+            "//executorch/exir:pass_base",
+            "//executorch/exir:pass_manager",
+            "//executorch/exir/backend:partitioner",
+        ],
+    )
+
     python_unittest(
         name = "common",
         srcs = [