diff --git a/backends/xnnpack/_passes/__init__.py b/backends/xnnpack/_passes/__init__.py index 22147fa4215..f4aea344d26 100644 --- a/backends/xnnpack/_passes/__init__.py +++ b/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, diff --git a/backends/xnnpack/_passes/convert_to_sdpa.py b/backends/xnnpack/_passes/convert_to_sdpa.py index c7982db750f..b926d62df8f 100644 --- a/backends/xnnpack/_passes/convert_to_sdpa.py +++ b/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 diff --git a/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py b/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py new file mode 100644 index 00000000000..01524a3268d --- /dev/null +++ b/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) diff --git a/backends/xnnpack/partition/config/__init__.py b/backends/xnnpack/partition/config/__init__.py index c6c54f083d6..651eba6234a 100644 --- a/backends/xnnpack/partition/config/__init__.py +++ b/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, diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py index f2b946e412d..ecb3ba9c8f5 100644 --- a/backends/xnnpack/partition/config/generic_node_configs.py +++ b/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" diff --git a/backends/xnnpack/test/ops/test_multiply.py b/backends/xnnpack/test/ops/test_multiply.py index 3315200005d..a44b5e6b406 100644 --- a/backends/xnnpack/test/ops/test_multiply.py +++ b/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)) ( diff --git a/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py b/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py new file mode 100644 index 00000000000..dd59aab5204 --- /dev/null +++ b/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py @@ -0,0 +1,142 @@ +# 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 unittest + +import torch +from executorch.backends.xnnpack._passes import XNNPACKPassManager +from executorch.backends.xnnpack._passes.convert_to_sdpa import ConvertToSDPAPass +from executorch.backends.xnnpack._passes.lift_constant_scalar_operands_pass import ( + LiftConstantScalarOperandsPass, +) +from executorch.backends.xnnpack.utils.configs import ( + get_transform_passes, + get_xnnpack_edge_compile_config, +) +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export.graph_signature import InputKind + + +class TestLiftConstantScalarOperandsPass(unittest.TestCase): + def setUp(self): + torch._dynamo.reset() + + class MulScalar(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.mul.Scalar(x, 0.5) + + class AddScalar(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.add.Scalar(x, 0.5) + + class SDPA(torch.nn.Module): + def forward(self, q, k, v, mask): + return torch.nn.functional.scaled_dot_product_attention(q, k, v, mask) + + def _to_edge_program_manager(self, module): + return to_edge( + torch.export.export(module, (torch.randn(2, 3),), strict=True), + compile_config=get_xnnpack_edge_compile_config(skip_dim_order=True), + ) + + def _lift(self, exported_program): + return XNNPACKPassManager( + exported_program, passes=[LiftConstantScalarOperandsPass] + ).transform() + + def test_lifts_mul_scalar_operand(self): + exported_program = self._lift( + self._to_edge_program_manager(self.MulScalar()).exported_program() + ) + graph = exported_program.graph_module.graph + + self.assertFalse( + any(node.target == exir_ops.edge.aten.mul.Scalar for node in graph.nodes) + ) + self.assertTrue( + any(node.target == exir_ops.edge.aten.mul.Tensor for node in graph.nodes) + ) + self.assertFalse(any(node.op == "get_attr" for node in graph.nodes)) + + constant_specs = [ + spec + for spec in exported_program.graph_signature.input_specs + if spec.kind == InputKind.CONSTANT_TENSOR + ] + self.assertEqual(len(constant_specs), 1) + constant_spec = constant_specs[0] + self.assertIn(constant_spec.target, exported_program.constants) + + placeholders = [node for node in graph.nodes if node.op == "placeholder"] + self.assertEqual(placeholders[0].name, constant_spec.arg.name) + mul_node = next( + node for node in graph.nodes if node.target == exir_ops.edge.aten.mul.Tensor + ) + self.assertIs(mul_node.args[1], placeholders[0]) + + def test_is_idempotent(self): + exported_program = self._lift( + self._to_edge_program_manager(self.MulScalar()).exported_program() + ) + exported_program = self._lift(exported_program) + + constant_specs = [ + spec + for spec in exported_program.graph_signature.input_specs + if spec.kind == InputKind.CONSTANT_TENSOR + ] + self.assertEqual(len(constant_specs), 1) + self.assertEqual(len(exported_program.constants), 1) + + def test_keeps_unmapped_scalar_op(self): + exported_program = self._lift( + self._to_edge_program_manager(self.AddScalar()).exported_program() + ) + graph = exported_program.graph_module.graph + + self.assertTrue( + any(node.target == exir_ops.edge.aten.add.Scalar for node in graph.nodes) + ) + self.assertFalse(exported_program.constants) + + def test_converts_sdpa_after_default_transform_passes(self): + q = torch.randn(2, 4, 8, 16) + k = torch.randn(2, 4, 8, 16) + v = torch.randn(2, 4, 8, 16) + mask = torch.randn(8, 8) + for use_default_transforms in (False, True): + with self.subTest(use_default_transforms=use_default_transforms): + edge = to_edge( + torch.export.export(self.SDPA(), (q, k, v, mask), strict=True), + compile_config=get_xnnpack_edge_compile_config(), + ) + if use_default_transforms: + edge = edge.transform(get_transform_passes()) + exported_program = XNNPACKPassManager( + edge.exported_program(), + passes=[ConvertToSDPAPass, LiftConstantScalarOperandsPass], + ).transform() + + graph = exported_program.graph_module.graph + self.assertTrue( + any( + node.target + == exir_ops.edge.aten.scaled_dot_product_attention.default + for node in graph.nodes + ) + ) + self.assertFalse( + any( + node.target == exir_ops.edge.aten.bmm.default + for node in graph.nodes + ) + ) + self.assertFalse( + any( + node.target == exir_ops.edge.aten.mul.Scalar + for node in graph.nodes + ) + ) diff --git a/backends/xnnpack/test/test_xnnpack_partitioner.py b/backends/xnnpack/test/test_xnnpack_partitioner.py index 894fab4098f..278646b9cf1 100644 --- a/backends/xnnpack/test/test_xnnpack_partitioner.py +++ b/backends/xnnpack/test/test_xnnpack_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. @@ -30,6 +31,39 @@ def __init__(self): def forward(self, x): return self.linear(x) + class MulScalar(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.mul.Scalar(x, 0.5) + + def test_mul_scalar_ops_to_not_decompose(self): + partitioner = XnnpackPartitioner() + exported_program = export(self.MulScalar(), (torch.randn(2, 3),)) + ops, filter_fn = partitioner.ops_to_not_decompose(exported_program) + + self.assertIn(torch.ops.aten.mul.Scalar, ops) + self.assertIsNotNone(filter_fn) + mul_node = next( + node + for node in exported_program.graph.nodes + if node.target == torch.ops.aten.mul.Scalar + ) + self.assertTrue(filter_fn(mul_node)) + + def test_mul_scalar_ops_to_not_decompose_rejects_unsupported_dtype(self): + partitioner = XnnpackPartitioner() + exported_program = export( + self.MulScalar(), (torch.ones(2, 3, dtype=torch.int32),) + ) + _, filter_fn = partitioner.ops_to_not_decompose(exported_program) + + self.assertIsNotNone(filter_fn) + mul_node = next( + node + for node in exported_program.graph.nodes + if node.target == torch.ops.aten.mul.Scalar + ) + self.assertFalse(filter_fn(mul_node)) + def test_deprecation_warning_for_to_backend_workflow(self): """ Test that the deprecated to_edge + to_backend workflow shows a deprecation warning.