Add partitioner pre-decomposition transform hook

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

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

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