From 824ec35f0e2ac86ddf59acd975b10fe7bc71e5b0 Mon Sep 17 00:00:00 2001 From: Yufeng Shi Date: Fri, 21 Aug 2026 13:03:33 +0100 Subject: [PATCH] Add partitioner pre-decomposition transform hook Invoke partitioner-provided transforms before EXIR's initial run_decompositions({}) call in to_edge_transform_and_lower(). This avoids backend-specific frontend wrappers for required transformations. The lowering order becomes: to_edge_transform_and_lower() -> partitioner.transform_for_pre_decomposition() -> program.run_decompositions({}) -> Edge transformation and partitioning Run multiple partitioner transforms in the supplied order and retain a no-op default for backward compatibility. Assisted by Codex. Change-Id: I9e67b896c22fd2043c3b4f5439016016acc05362 Signed-off-by: Yufeng Shi --- exir/backend/partitioner.py | 25 +++++++++ exir/program/_program.py | 14 ++++- .../test_partitioner_pre_decomposition.py | 54 +++++++++++++++++++ 3 files changed, 91 insertions(+), 2 deletions(-) create mode 100644 exir/tests/test_partitioner_pre_decomposition.py diff --git a/exir/backend/partitioner.py b/exir/backend/partitioner.py index 68d5c246906..c0edf0c71d8 100644 --- a/exir/backend/partitioner.py +++ b/exir/backend/partitioner.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -112,3 +113,27 @@ def ops_to_not_decompose( in the list returned by ops_to_not_decompose. """ return ([], None) + + def transform_for_pre_decomposition( + self, exported_program: ExportedProgram + ) -> ExportedProgram: + """Transform an ATen program before default decompositions run. + + EXIR invokes this hook automatically from + ``to_edge_transform_and_lower``. Backend partitioners may override it + to apply transformations that must run before decomposition. Callers + should use ``to_edge_transform_and_lower`` instead of invoking this + method directly. + + When multiple partitioners are used for a method, their transforms run + sequentially in the order supplied to ``to_edge_transform_and_lower``. + + Args: + exported_program (ExportedProgram): The ATen-dialect program to + transform. + + Returns: + ExportedProgram: The transformed ATen-dialect program. + + """ + return exported_program diff --git a/exir/program/_program.py b/exir/program/_program.py index 2dc9560924d..33cccf86776 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -1,6 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. -# Copyright 2025 Arm Limited and/or its affiliates. +# Copyright 2025-2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -1114,6 +1114,14 @@ def _can_skip_using_EDGE_DO_NOT_DECOMP( return check_op_support is None +def _apply_pre_decomposition_transforms( + program: ExportedProgram, partitioners: List[Partitioner] +) -> ExportedProgram: + for partitioner in partitioners: + program = partitioner.transform_for_pre_decomposition(program) + return program + + def _gen_edge_manager_for_partitioners( partitioner: Dict[str, List[Partitioner]], aten_programs: Dict[str, ExportedProgram], @@ -1137,11 +1145,13 @@ def _gen_edge_manager_for_partitioners( ops_set_to_not_decompose_by_program = defaultdict(list) edge_programs: Dict[str, ExportedProgram] = {} for name, program in aten_programs.items(): + partitioners_for_program = partitioner.get(name, []) + program = _apply_pre_decomposition_transforms(program, partitioners_for_program) + # Functionalize program before asking partitioners to preserve ops program = program.run_decompositions({}) if partitioner is not None: - partitioners_for_program = partitioner.get(name, []) final_ops_to_preserve = set() # Decompose by default if there are no partitioners for the method diff --git a/exir/tests/test_partitioner_pre_decomposition.py b/exir/tests/test_partitioner_pre_decomposition.py new file mode 100644 index 00000000000..52c0ef32cbd --- /dev/null +++ b/exir/tests/test_partitioner_pre_decomposition.py @@ -0,0 +1,54 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from executorch.exir import to_edge_transform_and_lower +from executorch.exir.backend.partitioner import Partitioner, PartitionResult +from torch.export import export, ExportedProgram + + +class _SDPA(torch.nn.Module): + def forward( + self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor + ) -> torch.Tensor: + return torch.nn.functional.scaled_dot_product_attention(query, key, value) + + +class _RecordingPartitioner(Partitioner): + def __init__(self, name: str, calls: list[str]) -> None: + super().__init__() + self.name = name + self.calls = calls + self.saw_sdpa = False + + def transform_for_pre_decomposition( + self, exported_program: ExportedProgram + ) -> ExportedProgram: + self.calls.append(self.name) + self.saw_sdpa = any( + node.target == torch.ops.aten.scaled_dot_product_attention.default + for node in exported_program.graph.nodes + ) + return exported_program + + def partition(self, exported_program: ExportedProgram) -> PartitionResult: + return PartitionResult(exported_program, {}) + + +def test_partitioner_transforms_run_before_decomposition_in_order() -> None: + inputs = tuple(torch.randn(1, 3, 4, 5) for _ in range(3)) + exported_program = export(_SDPA(), inputs, strict=True) + calls: list[str] = [] + first = _RecordingPartitioner("first", calls) + second = _RecordingPartitioner("second", calls) + + to_edge_transform_and_lower( + exported_program, + partitioner=[first, second], + ) + + assert calls == ["first", "second"] + assert first.saw_sdpa + assert second.saw_sdpa