diff --git a/backends/qualcomm/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py index 23016de006b..b27b2eb94a3 100644 --- a/backends/qualcomm/partition/qnn_partitioner.py +++ b/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,6 +203,9 @@ 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: @@ -206,7 +213,15 @@ def __init__( 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]: diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 606fed01d72..235d3975006 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/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. diff --git a/backends/qualcomm/utils/utils.py b/backends/qualcomm/utils/utils.py index 2f831a15b55..eea8831d246 100644 --- a/backends/qualcomm/utils/utils.py +++ b/backends/qualcomm/utils/utils.py @@ -462,9 +462,11 @@ 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( @@ -472,6 +474,9 @@ def ensure_graph_specific_dict(value, graph_names): 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(), diff --git a/exir/backend/partitioner.py b/exir/backend/partitioner.py index c0edf0c71d8..286c39579a0 100644 --- a/exir/backend/partitioner.py +++ b/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: diff --git a/exir/program/_program.py b/exir/program/_program.py index ea9b09b847f..a3889a62572 100644 --- a/exir/program/_program.py +++ b/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: diff --git a/exir/tests/targets.bzl b/exir/tests/targets.bzl index f68ee23cc16..7fc46ed309b 100644 --- a/exir/tests/targets.bzl +++ b/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 = [ diff --git a/exir/tests/test_partitioner_transform_passes.py b/exir/tests/test_partitioner_transform_passes.py new file mode 100644 index 00000000000..4fe404c22c0 --- /dev/null +++ b/exir/tests/test_partitioner_transform_passes.py @@ -0,0 +1,104 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import List + +import pytest + +import torch +from executorch.exir import to_edge_transform_and_lower +from executorch.exir.backend.partitioner import Partitioner, PartitionResult +from executorch.exir.pass_base import ExportPass, PassResult +from executorch.exir.pass_manager import PassType +from torch.export import export, ExportedProgram + + +class _Add(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + x + + +class _RecordingPass(ExportPass): + def __init__(self, name: str, calls: List[str]) -> None: + super().__init__() + self.name = name + self.calls = calls + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + self.calls.append(self.name) + return PassResult(graph_module, False) + + +class _ContributingPartitioner(Partitioner): + def __init__(self, names: List[str], calls: List[str]) -> None: + super().__init__() + self.passes = [_RecordingPass(name, calls) for name in names] + self.saw_program = False + + def transform_passes(self, exported_program: ExportedProgram) -> List[PassType]: + self.saw_program = isinstance(exported_program, ExportedProgram) + return list(self.passes) + + def partition(self, exported_program: ExportedProgram) -> PartitionResult: + return PartitionResult(exported_program, {}) + + +class _NoopPartitioner(Partitioner): + def partition(self, exported_program: ExportedProgram) -> PartitionResult: + return PartitionResult(exported_program, {}) + + +def _lower(partitioners, transform_passes=None): + exported_program = export(_Add(), (torch.randn(2, 3),), strict=True) + return to_edge_transform_and_lower( + exported_program, + transform_passes=transform_passes, + partitioner=partitioners, + ) + + +def test_partitioner_transform_passes_run() -> None: + calls: List[str] = [] + partitioner = _ContributingPartitioner(["a"], calls) + + _lower([partitioner]) + + assert calls == ["a"] + assert partitioner.saw_program + + +def test_partitioner_transform_passes_default_is_empty() -> None: + # A partitioner that does not override the hook must not change behavior. + _lower([_NoopPartitioner()]) + + +def test_partitioner_transform_passes_concatenate_in_order() -> None: + calls: List[str] = [] + first = _ContributingPartitioner(["first"], calls) + second = _ContributingPartitioner(["second"], calls) + + _lower([first, second]) + + assert calls == ["first", "second"] + + +def test_caller_transform_passes_run_before_partitioner_passes() -> None: + calls: List[str] = [] + partitioner = _ContributingPartitioner(["from_partitioner"], calls) + + _lower([partitioner], transform_passes=[_RecordingPass("from_caller", calls)]) + + assert calls == ["from_caller", "from_partitioner"] + + +def test_pass_manager_argument_with_contributed_passes_is_rejected() -> None: + from torch.fx.passes.infra.pass_manager import PassManager + + calls: List[str] = [] + partitioner = _ContributingPartitioner(["a"], calls) + + with pytest.raises(ValueError, match="cannot be combined"): + _lower([partitioner], transform_passes=PassManager([]))