XNNPACK: Lift constant mul scalars for partitioning

backends/xnnpack/_passes/__init__.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.
@@ -28,6 +29,9 @@
 from executorch.backends.xnnpack._passes.fuse_activation_pass import FuseActivationPass
 from executorch.backends.xnnpack._passes.fuse_batch_norm import FuseBatchNormPass
 from executorch.backends.xnnpack._passes.insert_pad_qdq import InsertPadQDQPass
+from executorch.backends.xnnpack._passes.lift_constant_scalar_operands_pass import (
+    LiftConstantScalarOperandsPass,
+)
 from executorch.backends.xnnpack._passes.prelu_reshape_pass import PReLUReshapePass
 from executorch.backends.xnnpack._passes.propagate_custom_meta_pass import (
     PropagateCustomMetaPass,
@@ -76,6 +80,7 @@ def __init__(
                 ConvertToLinearPass,
                 PropagateCustomMetaPass,
                 ConvertToSDPAPass,
+                LiftConstantScalarOperandsPass,
                 ConstPropPass,
                 FuseBatchNormPass,
                 DecomposeBatchNorm,

backends/xnnpack/_passes/convert_to_sdpa.py

@@ -1,19 +1,22 @@
 # 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.
 
 import logging
+from copy import deepcopy
 from typing import Optional
 
 import torch
 from executorch.backends.transforms import get_shape
-
+from executorch.backends.xnnpack._passes.remove_noop_expand_copy_pass import (
+    RemoveNoopExpandCopyPass,
+)
 from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass
 from executorch.backends.xnnpack.partition.graphs import sdpa
 from executorch.exir.dialects._ops import ops as exir_ops
-
 from torch.fx.passes.infra.pass_base import PassResult
 from torch.fx.passes.utils.matcher_utils import InternalMatch, SubgraphMatcher
 
@@ -24,31 +27,31 @@
 class ConvertToSDPAPass(XNNPACKPass):
     def get_scale(self, match: InternalMatch) -> Optional[float]:
         """
-        Returns the scale of the SDPA op.
+        Return the SDPA scale recovered from the matched pre-QK^T multiplications.
 
-        Scale: Optional[float] doesn't change the graph pattern.
-        The default value can be calulated however we need to extract
-        it for lowering when it is the user supplied value anyway.
+        The decomposition applies the square root of the attention scale before
+        QK^T, so the extracted multiplier is squared to recover the original value.
         """
         for node in match.nodes_map.values():
             if (
-                node.op == "call_function"
-                and node.target == exir_ops.edge.aten.mul.Scalar
+                node.op != "call_function"
+                or node.target != exir_ops.edge.aten.mul.Scalar
             ):
-                scale = node.args[1]
+                continue
 
-                dtype = torch.float
-                mul_val = node.meta.get("val", None)
-                if mul_val is not None:
-                    dtype = mul_val.dtype
+            scale = node.args[1]
 
-                if isinstance(scale, float):
-                    # Convert scale value to fp16 (reducing precision)
-                    scale = torch.tensor(scale, dtype=dtype).item()
+            dtype = torch.float
+            mul_val = node.meta.get("val", None)
+            if mul_val is not None:
+                dtype = mul_val.dtype
 
-                    # since scale we extracted this before the QK^T.
-                    return scale**2
-                break
+            if isinstance(scale, float):
+                # Convert scale value to fp16 (reducing precision)
+                scale = torch.tensor(scale, dtype=dtype).item()
+
+                # since scale we extracted this before the QK^T.
+                return scale**2
         return None
 
     def assert_2d_mask(self, match: InternalMatch) -> None:
@@ -99,11 +102,16 @@ def call(self, graph_module: torch.fx.GraphModule):
         logger.debug("ConvertToSDPA Begin: ")
         logger.debug(graph_module.print_readable(print_output=False))
 
-        for pattern in sdpa.get_graphs():
-            sm = SubgraphMatcher(pattern.graph, ignore_literals=True)
-            matches = list(sm.match(graph_module.graph))
-            for partition_to_replace in matches:
-                self.create_sdpa(graph_module, partition_to_replace)
+        for scalar_pattern in sdpa.get_graphs():
+            normalized_pattern = RemoveNoopExpandCopyPass()(
+                deepcopy(scalar_pattern)
+            ).graph_module
+
+            for pattern in (scalar_pattern, normalized_pattern):
+                sm = SubgraphMatcher(pattern.graph, ignore_literals=True)
+                matches = list(sm.match(graph_module.graph))
+                for partition_to_replace in matches:
+                    self.create_sdpa(graph_module, partition_to_replace)
 
         graph_module.recompile()
         graph_module = super().call(graph_module).graph_module

backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py

@@ -0,0 +1,119 @@
+# 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 __future__ import annotations
+
+from numbers import Number
+from typing import Dict, Optional, Union
+
+import torch
+from executorch.backends.transforms.utils import create_constant_placeholder
+from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.dialects.edge._ops import EdgeOpOverload
+from executorch.exir.pass_base import PassResult
+from torch._ops import OpOverload
+from torch.export import ExportedProgram
+from torch.export.graph_signature import InputKind
+
+ScalarOp = Union[EdgeOpOverload, OpOverload]
+
+
+class LiftConstantScalarOperandsPass(XNNPACKPass):
+    """
+    Lift scalar operands into tensor constants for selected binary ops.
+
+    XNNPACK already supports the tensor overloads for these binary operations.
+    This pass converts explicitly listed scalar overloads to their tensor
+    overloads by replacing constant scalar operands with small tensor constants.
+    The constants are registered as exported-program constant tensor inputs.
+    Keep the op map narrow until each new scalar overload is covered by tests.
+    """
+
+    default_scalar_to_tensor_ops: Dict[ScalarOp, ScalarOp] = {
+        exir_ops.edge.aten.mul.Scalar: exir_ops.edge.aten.mul.Tensor,
+    }
+
+    def __init__(
+        self,
+        exported_program: ExportedProgram,
+        scalar_to_tensor_ops: Optional[Dict[ScalarOp, ScalarOp]] = None,
+    ) -> None:
+        super().__init__(exported_program)
+        self.scalar_to_tensor_ops = (
+            scalar_to_tensor_ops
+            if scalar_to_tensor_ops is not None
+            else self.default_scalar_to_tensor_ops
+        )
+
+    def _create_constant_node(
+        self,
+        graph_module: torch.fx.GraphModule,
+        node: torch.fx.Node,
+        value: Number,
+    ) -> torch.fx.Node:
+        input_node = node.args[0]
+        if not isinstance(input_node, torch.fx.Node):
+            raise RuntimeError("Expected scalar op input to be an FX node.")
+
+        input_value = input_node.meta["val"]
+        tensor = torch.tensor(value, dtype=input_value.dtype, device=input_value.device)
+        name = self._get_new_constant_name(graph_module)
+        first_placeholder = next(
+            graph_node
+            for graph_node in graph_module.graph.nodes
+            if graph_node.op == "placeholder"
+        )
+        with graph_module.graph.inserting_before(first_placeholder):
+            return create_constant_placeholder(
+                self.exported_program,
+                graph_module.graph,
+                name,
+                InputKind.CONSTANT_TENSOR,
+                tensor,
+            )
+
+    def _get_new_constant_name(self, graph_module: torch.fx.GraphModule) -> str:
+        prefix = "_tensor_constant_"
+        existing_names = {node.name for node in graph_module.graph.nodes}
+        existing_names.update(self.exported_program.constants)
+        existing_names.update(self.exported_program.state_dict)
+        index = 0
+        while f"{prefix}{index}" in existing_names:
+            index += 1
+        return f"{prefix}{index}"
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        modified = False
+
+        for node in list(graph_module.graph.nodes):
+            if (
+                node.op != "call_function"
+                or node.target not in self.scalar_to_tensor_ops
+                or len(node.args) != 2
+                or not isinstance(node.args[0], torch.fx.Node)
+                or not isinstance(node.args[1], Number)
+            ):
+                continue
+
+            input_value = node.args[0].meta.get("val")
+            output_value = node.meta.get("val")
+            if (
+                not isinstance(input_value, torch.Tensor)
+                or not isinstance(output_value, torch.Tensor)
+                or input_value.dtype != output_value.dtype
+            ):
+                continue
+
+            tensor_arg = self._create_constant_node(graph_module, node, node.args[1])
+            node.args = (node.args[0], tensor_arg)
+            node.target = self.scalar_to_tensor_ops[node.target]
+            modified = True
+
+        graph_module.graph.eliminate_dead_code()
+        graph_module.graph.lint()
+        graph_module.recompile()
+
+        return PassResult(graph_module, modified)

backends/xnnpack/partition/config/__init__.py

@@ -43,6 +43,7 @@
     MeanDimConfig,
     MinimumConfig,
     MulConfig,
+    MulScalarConfig,
     NegConfig,
     PermuteConfig,
     PowConfig,
@@ -106,6 +107,7 @@
     MinimumConfig,
     MMConfig,
     MulConfig,
+    MulScalarConfig,
     NegConfig,
     PermuteConfig,
     PowConfig,

backends/xnnpack/partition/config/generic_node_configs.py

@@ -8,6 +8,7 @@
 # pyre-unsafe
 
 import logging
+from numbers import Number
 from typing import cast, List, Optional
 
 import numpy as np
@@ -416,6 +417,35 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]:
         return [ConfigPrecisionType.FP32, ConfigPrecisionType.STATIC_QUANT]
 
 
+class MulScalarConfig(GenericNodePartitionerConfig):
+    target_name = "mul.Scalar"
+
+    def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
+        if not self.check_common_constraints(node, ep):
+            return False
+
+        if (
+            len(node.args) != 2
+            or not isinstance(node.args[0], torch.fx.Node)
+            or not isinstance(node.args[1], Number)
+        ):
+            return False
+
+        input_value = node.args[0].meta.get("val")
+        output_value = node.meta.get("val")
+        return (
+            isinstance(input_value, torch.Tensor)
+            and isinstance(output_value, torch.Tensor)
+            and input_value.dtype == output_value.dtype
+        )
+
+    def supported_precision_types(self) -> List[ConfigPrecisionType]:
+        return [ConfigPrecisionType.FP32]
+
+    def get_original_aten(self) -> Optional[torch._ops.OpOverload]:
+        return torch.ops.aten.mul.Scalar
+
+
 class MaximumConfig(GenericNodePartitionerConfig):
     target_name = "maximum.default"
 

backends/xnnpack/test/ops/test_multiply.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.
@@ -29,6 +30,10 @@ def forward(self, x, y):
             z = torch.mul(x, y) * torch.functional.torch.mul(x, y)
             return z
 
+    class MulScalar(torch.nn.Module):
+        def forward(self, x):
+            return torch.ops.aten.mul.Scalar(x, 0.5)
+
     class MulRelu(torch.nn.Module):
         def forward(self, x, y):
             z = x * y
@@ -58,6 +63,23 @@ def test_fp32_mul(self):
         inputs = (torch.randn((1, 3)), torch.randn((4, 3)))
         self._test_mul(inputs)
 
+    def test_fp32_mul_scalar(self):
+        (
+            Tester(self.MulScalar(), (torch.randn(2, 3),))
+            .export()
+            .to_edge_transform_and_lower()
+            .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+            .check_not(
+                [
+                    "executorch_exir_dialects_edge__ops_aten_mul_Tensor",
+                    "executorch_exir_dialects_edge__ops_aten_mul_Scalar",
+                ]
+            )
+            .to_executorch()
+            .serialize()
+            .run_method_and_compare_outputs()
+        )
+
     def test_qs8_mul(self):
         inputs = (torch.randn(1, 1, 4, 4), torch.randn(1, 1, 4, 1))
         (