Arm backend: Support quantization of torch.einsum (pytorch#18539)

- Refactor GetDecompositionPass to expose small override hooks for trace
inputs, placeholder mapping, and output extraction.

- Implement DecomposeEinsumPass as a GetDecompositionPass subclass to
handle aten.einsum.default's (equation, [operands]) signature.

Change-Id: I0546572f44903e400dfd75ea067ed2785e2a326e

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell

Signed-off-by: Yufeng Shi <yufeng.shi@arm.com>

Author: YufengShi-dudu

backends/arm/_passes/__init__.py

@@ -43,6 +43,7 @@
 from .decompose_cumsum_pass import DecomposeCumsumPass  # noqa
 from .decompose_div_pass import DecomposeDivPass  # noqa
 from .decompose_div_tensor_mode import DecomposeDivTensorModePass  # noqa
+from .decompose_einsum_pass import DecomposeEinsumPass  # noqa
 from .decompose_elu_pass import DecomposeEluPass  # noqa
 from .decompose_embedding_pass import DecomposeEmbeddingPass  # noqa  # noqa
 from .decompose_erfinv_pass import DecomposeErfinvPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -51,6 +51,7 @@
     DecomposeCumsumPass,
     DecomposeDivPass,
     DecomposeDivTensorModePass,
+    DecomposeEinsumPass,
     DecomposeEluPass,
     DecomposeEmbeddingPass,
     DecomposeErfinvPass,
@@ -560,6 +561,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
                 DecomposeFloorDividePass(tfa_pass=True),
                 DecomposeDivTensorModePass(tfa_pass=True),
                 DecomposeWhereScalarOtherPass(tfa_pass=True),
+                DecomposeEinsumPass(tfa_pass=True),
                 RewriteInplaceArithmeticPass(tfa_pass=True),
                 DecomposeAddSubAlphaPass(tfa_pass=True),
                 DecomposeLeakyReLUPass(tfa_pass=True),

backends/arm/_passes/decompose_einsum_pass.py

@@ -0,0 +1,66 @@
+# 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.backends.arm._passes.get_decomposition_pass import GetDecompositionPass
+
+
+class DecomposeEinsumPass(GetDecompositionPass):
+    """Decomposes aten.einsum.default into more primitive ops.
+
+    This pass is intended to be called in transform_for_annotation to prepare
+    the graph for quantization. Einsum is not annotated directly by the Arm
+    quantizer, but the decomposed ops are.
+
+    """
+
+    targeted_ops = [torch.ops.aten.einsum.default]
+
+    def _get_input_tensors(self, node: torch.fx.Node) -> list:
+        """Override the base hook because aten.einsum.default takes (equation,
+        [operands]), which cannot be handled by the generic one-arg-per-input
+        logic.
+        """
+        equation, operands = node.args  # type: ignore[union-attr]
+        fake_operands = [operand.meta["val"] for operand in operands]  # type: ignore[union-attr]
+        return [equation, fake_operands]
+
+    def _get_placeholder_map(
+        self,
+        node: torch.fx.Node,
+        decomposed_module: torch.fx.GraphModule,
+    ) -> dict[str, torch.fx.Node]:
+        """Override the base hook because einsum does not trace placeholders
+        one-to-one with node.args.
+
+        The traced graph includes arg0_1 for the equation string and arg1_i for
+        each tensor inside the operand list, so we must skip the equation
+        placeholder, which is not an original FX tensor node, and map each
+        operand placeholder back to the corresponding original FX node.
+
+        """
+        _, operands = node.args
+        name_to_input_tensor_map = {}
+
+        for decomposed_node in decomposed_module.graph.nodes:
+            if decomposed_node.op != "placeholder":
+                continue
+            if decomposed_node.name == "arg0_1":
+                continue
+            if not decomposed_node.name.startswith("arg1_"):
+                raise RuntimeError(
+                    f"Unexpected einsum placeholder name {decomposed_node.name!r}."
+                )
+
+            operand_idx = int(decomposed_node.name.split("_")[1]) - 1
+            name_to_input_tensor_map[decomposed_node.name] = operands[operand_idx]  # type: ignore[index]
+
+        return name_to_input_tensor_map  # type: ignore[return-value]
+
+    def _get_output_node(self, output_node: torch.fx.Node) -> torch.fx.Node:
+        """Return the traced value node for einsum graphs that emit
+        output([node]).
+        """
+        return output_node.args[0][0]  # type: ignore[index, return-value]

backends/arm/_passes/get_decomposition_pass.py

@@ -34,6 +34,31 @@ def __init__(self, tfa_pass=False, *args, **kwargs):
     def _skip_pass(self, input_tensors: list) -> bool:
         return False
 
+    def _get_input_tensors(self, node: torch.fx.Node) -> list:
+        input_tensors = []
+        for arg in node.args:
+            if hasattr(arg, "meta"):
+                input_tensors.append(arg.meta["val"])  # type: ignore[union-attr]
+            elif isinstance(arg, int):
+                input_tensors.append(arg)
+        return input_tensors
+
+    def _get_placeholder_map(
+        self,
+        node: torch.fx.Node,
+        decomposed_module: torch.fx.GraphModule,
+    ) -> dict[str, torch.fx.Node]:
+        # Keep decomposed_module in the hook signature so subclasses can inspect
+        # traced placeholder structure when the mapping is not one-to-one.
+        name_to_input_tensor_map = {}
+        for i, arg in enumerate(node.args):
+            name_to_input_tensor_map[f"arg{i}_1"] = arg
+        return name_to_input_tensor_map  # type: ignore[return-value]
+
+    def _get_output_node(self, output_node: torch.fx.Node) -> torch.fx.Node:
+        """Return the traced value node for graphs that emit output(node)."""
+        return output_node.args[0]  # type: ignore[return-value]
+
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
         modified = False
         for node in graph_module.graph.nodes:
@@ -44,13 +69,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
             ):
                 continue
 
-            input_tensors = []
-            for arg in node.args:
-                if hasattr(arg, "meta"):
-                    input_tensors.append(arg.meta["val"])
-
-                elif isinstance(arg, int):
-                    input_tensors.append(arg)
+            input_tensors = self._get_input_tensors(node)
 
             if self._skip_pass(input_tensors):
                 continue
@@ -70,22 +89,26 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
             )(*input_tensors)
 
             with graph_module.graph.inserting_before(node):
-                name_to_input_tensor_map = {}
-                for i, arg in enumerate(node.args):
-                    name_to_input_tensor_map[f"arg{i}_1"] = arg
+                name_to_input_tensor_map = self._get_placeholder_map(
+                    node, decomposed_module
+                )
 
                 decomposed_node_to_subgraph_node = {}
                 last_decomposed_node = None
                 # Create a mapping from input nodes in decomposed module to original nodes.
                 # In decomposed module, there are only input tensors for placeholder op.
                 for decomposed_node in decomposed_module.graph.nodes:
                     if decomposed_node.op == "placeholder":
+                        # Some ops, such as einsum, trace extra placeholders that do
+                        # not map back to original graph tensor inputs.
+                        if decomposed_node.name not in name_to_input_tensor_map:
+                            continue
                         decomposed_node_to_subgraph_node[decomposed_node] = (
                             name_to_input_tensor_map[decomposed_node.name]
                         )
 
                     if decomposed_node.op == "output":
-                        last_decomposed_node = decomposed_node.args[0]
+                        last_decomposed_node = self._get_output_node(decomposed_node)
 
                 # Copy node from decompose graph module
                 for decomposed_node in decomposed_module.graph.nodes:

backends/arm/test/models/stable_diffusion/test_SD3Transformer2DModel.py

@@ -42,7 +42,6 @@ class TestSD3Transformer2DModel:
     ops_after_partitioner_INT = {
         "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 2,
         "torch.ops.higher_order.executorch_call_delegate": 2,
-        "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1,
     }
 
     ops_after_partitioner_vgf_quantize = {

backends/arm/test/passes/test_decompose_einsum_pass.py

@@ -0,0 +1,73 @@
+# 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.
+
+from typing import Tuple
+
+import torch
+from executorch.backends.arm._passes import DecomposeEinsumPass
+from executorch.backends.arm.quantizer import (
+    get_symmetric_quantization_config,
+    TOSAQuantizer,
+)
+from executorch.backends.arm.test.tester.test_pipeline import QuantizationPipeline
+from executorch.backends.arm.tosa import TosaSpecification
+from torch.export import export
+
+
+input_t = Tuple[torch.Tensor]
+
+
+class EinsumPermuteModule(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.einsum("nhwpqc->nchpwq", x)
+
+    @staticmethod
+    def get_inputs() -> input_t:
+        return (torch.randn(2, 4, 16, 1, 16, 1),)
+
+
+def _get_int8_quantizer() -> TOSAQuantizer:
+    quantizer = TOSAQuantizer(TosaSpecification.create_from_string("TOSA-1.0+INT"))
+    quantizer.set_global(get_symmetric_quantization_config())
+    return quantizer
+
+
+def test_decompose_einsum_no_target_rewrites_export_graph() -> None:
+    module = EinsumPermuteModule().eval()
+    exported_program = export(module, module.get_inputs())
+
+    before_targets = [
+        str(node.target)
+        for node in exported_program.graph_module.graph.nodes
+        if node.op == "call_function"
+    ]
+    assert before_targets == ["aten.einsum.default"]
+
+    pass_result = DecomposeEinsumPass().call(exported_program.graph_module)
+
+    after_targets = [
+        str(node.target)
+        for node in pass_result.graph_module.graph.nodes
+        if node.op == "call_function"
+    ]
+    assert "aten.einsum.default" not in after_targets
+    assert after_targets == ["aten.permute.default"]
+
+
+def test_decompose_einsum_tosa_INT_quantizes_after_transform_for_annotation() -> None:
+    module = EinsumPermuteModule().eval()
+    quantization_annotations = {
+        "aten.permute.default": {
+            get_symmetric_quantization_config().output_activation: 1
+        }
+    }
+
+    pipeline = QuantizationPipeline[input_t](
+        module,
+        module.get_inputs(),
+        quantizer=_get_int8_quantizer(),
+        qspecs=quantization_annotations,  # type: ignore[arg-type]
+    )
+    pipeline.run()