Skip to content

Commit 3413fdf

Browse files
XNNPACK: Lift constant mul scalars for partitioning (pytorch#20515)
XNNPACK supports the tensor overload for multiply, but aten.mul.Scalar is normally decomposed before the XNNPACK partitioner can select it. Add a MulScalarConfig that preserves eligible scalar multiplications during partitioning. During XNNPACK preprocessing, lift the scalar into an exported constant-tensor placeholder and rewrite the operation to aten.mul.Tensor, allowing the existing multiply serialization path to handle it. This avoids introducing an aten.full operation and works with the default to_edge_transform_and_lower flow without explicit XNNPACK transform passes. Run scalar lifting after SDPA conversion so ConvertToSDPAPass can recover attention scales from the original aten.mul.Scalar operations. Also allow SDPA conversion to match graphs processed by the default no-op-expand transform. For DeIT Tiny, this removes 24 portable aten.mul.Scalar nodes and reduces delegate count from 62 to 50. In current local timing checks the latency impact is modest: about 1% faster on both Android SME2 and the aarch64 XNNPACK/KleidiAI NEON-class host runner. These are modest uplifts but may introduce more opportunities for improvements. cc @GregoryComer @digantdesai @cbilgin @freddan80 @per @zingo @oscarandersson8218 @Sebastian-Larsson @robell @rascani --------- Signed-off-by: Måns Nilsson <mans.nilsson@arm.com> Co-authored-by: Jacob Stevens <stevens.jacob1492@gmail.com>
1 parent a33ac10 commit 3413fdf

8 files changed

Lines changed: 420 additions & 24 deletions

File tree

‎backends/xnnpack/_passes/__init__.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,9 @@
2929
from executorch.backends.xnnpack._passes.fuse_activation_pass import FuseActivationPass
3030
from executorch.backends.xnnpack._passes.fuse_batch_norm import FuseBatchNormPass
3131
from executorch.backends.xnnpack._passes.insert_pad_qdq import InsertPadQDQPass
32+
from executorch.backends.xnnpack._passes.lift_constant_scalar_operands_pass import (
33+
LiftConstantScalarOperandsPass,
34+
)
3235
from executorch.backends.xnnpack._passes.prelu_reshape_pass import PReLUReshapePass
3336
from executorch.backends.xnnpack._passes.propagate_custom_meta_pass import (
3437
PropagateCustomMetaPass,
@@ -79,6 +82,7 @@ def __init__(
7982
ConvertToLinearPass,
8083
PropagateCustomMetaPass,
8184
ConvertToSDPAPass,
85+
LiftConstantScalarOperandsPass,
8286
ConstPropPass,
8387
FuseBatchNormPass,
8488
DecomposeBatchNorm,

‎backends/xnnpack/_passes/convert_to_sdpa.py‎

Lines changed: 32 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,22 @@
11
# Copyright (c) Meta Platforms, Inc. and affiliates.
22
# All rights reserved.
3+
# Copyright 2026 Arm Limited and/or its affiliates.
34
#
45
# This source code is licensed under the BSD-style license found in the
56
# LICENSE file in the root directory of this source tree.
67

78
import logging
9+
from copy import deepcopy
810
from typing import Optional
911

1012
import torch
1113
from executorch.backends.transforms import get_shape
12-
14+
from executorch.backends.xnnpack._passes.remove_noop_expand_copy_pass import (
15+
RemoveNoopExpandCopyPass,
16+
)
1317
from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass
1418
from executorch.backends.xnnpack.partition.graphs import sdpa
1519
from executorch.exir.dialects._ops import ops as exir_ops
16-
1720
from torch.fx.passes.infra.pass_base import PassResult
1821
from torch.fx.passes.utils.matcher_utils import InternalMatch, SubgraphMatcher
1922

@@ -24,31 +27,31 @@
2427
class ConvertToSDPAPass(XNNPACKPass):
2528
def get_scale(self, match: InternalMatch) -> Optional[float]:
2629
"""
27-
Returns the scale of the SDPA op.
30+
Return the SDPA scale recovered from the matched pre-QK^T multiplications.
2831
29-
Scale: Optional[float] doesn't change the graph pattern.
30-
The default value can be calulated however we need to extract
31-
it for lowering when it is the user supplied value anyway.
32+
The decomposition applies the square root of the attention scale before
33+
QK^T, so the extracted multiplier is squared to recover the original value.
3234
"""
3335
for node in match.nodes_map.values():
3436
if (
35-
node.op == "call_function"
36-
and node.target == exir_ops.edge.aten.mul.Scalar
37+
node.op != "call_function"
38+
or node.target != exir_ops.edge.aten.mul.Scalar
3739
):
38-
scale = node.args[1]
40+
continue
3941

40-
dtype = torch.float
41-
mul_val = node.meta.get("val", None)
42-
if mul_val is not None:
43-
dtype = mul_val.dtype
42+
scale = node.args[1]
4443

45-
if isinstance(scale, float):
46-
# Convert scale value to fp16 (reducing precision)
47-
scale = torch.tensor(scale, dtype=dtype).item()
44+
dtype = torch.float
45+
mul_val = node.meta.get("val", None)
46+
if mul_val is not None:
47+
dtype = mul_val.dtype
4848

49-
# since scale we extracted this before the QK^T.
50-
return scale**2
51-
break
49+
if isinstance(scale, float):
50+
# Convert scale value to fp16 (reducing precision)
51+
scale = torch.tensor(scale, dtype=dtype).item()
52+
53+
# since scale we extracted this before the QK^T.
54+
return scale**2
5255
return None
5356

5457
def assert_2d_mask(self, match: InternalMatch) -> None:
@@ -99,11 +102,16 @@ def call(self, graph_module: torch.fx.GraphModule):
99102
logger.debug("ConvertToSDPA Begin: ")
100103
logger.debug(graph_module.print_readable(print_output=False))
101104

102-
for pattern in sdpa.get_graphs():
103-
sm = SubgraphMatcher(pattern.graph, ignore_literals=True)
104-
matches = list(sm.match(graph_module.graph))
105-
for partition_to_replace in matches:
106-
self.create_sdpa(graph_module, partition_to_replace)
105+
for scalar_pattern in sdpa.get_graphs():
106+
normalized_pattern = RemoveNoopExpandCopyPass()(
107+
deepcopy(scalar_pattern)
108+
).graph_module
109+
110+
for pattern in (scalar_pattern, normalized_pattern):
111+
sm = SubgraphMatcher(pattern.graph, ignore_literals=True)
112+
matches = list(sm.match(graph_module.graph))
113+
for partition_to_replace in matches:
114+
self.create_sdpa(graph_module, partition_to_replace)
107115

108116
graph_module.recompile()
109117
graph_module = super().call(graph_module).graph_module
Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
# Copyright 2026 Arm Limited and/or its affiliates.
2+
#
3+
# This source code is licensed under the BSD-style license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
from __future__ import annotations
7+
8+
from numbers import Number
9+
from typing import Dict, Optional, Union
10+
11+
import torch
12+
from executorch.backends.transforms.utils import create_constant_placeholder
13+
from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass
14+
from executorch.exir.dialects._ops import ops as exir_ops
15+
from executorch.exir.dialects.edge._ops import EdgeOpOverload
16+
from executorch.exir.pass_base import PassResult
17+
from torch._ops import OpOverload
18+
from torch.export import ExportedProgram
19+
from torch.export.graph_signature import InputKind
20+
21+
ScalarOp = Union[EdgeOpOverload, OpOverload]
22+
23+
24+
class LiftConstantScalarOperandsPass(XNNPACKPass):
25+
"""
26+
Lift scalar operands into tensor constants for selected binary ops.
27+
28+
XNNPACK already supports the tensor overloads for these binary operations.
29+
This pass converts explicitly listed scalar overloads to their tensor
30+
overloads by replacing constant scalar operands with small tensor constants.
31+
The constants are registered as exported-program constant tensor inputs.
32+
Keep the op map narrow until each new scalar overload is covered by tests.
33+
"""
34+
35+
default_scalar_to_tensor_ops: Dict[ScalarOp, ScalarOp] = {
36+
exir_ops.edge.aten.mul.Scalar: exir_ops.edge.aten.mul.Tensor,
37+
}
38+
39+
def __init__(
40+
self,
41+
exported_program: ExportedProgram,
42+
scalar_to_tensor_ops: Optional[Dict[ScalarOp, ScalarOp]] = None,
43+
) -> None:
44+
super().__init__(exported_program)
45+
self.scalar_to_tensor_ops = (
46+
scalar_to_tensor_ops
47+
if scalar_to_tensor_ops is not None
48+
else self.default_scalar_to_tensor_ops
49+
)
50+
51+
def _create_constant_node(
52+
self,
53+
graph_module: torch.fx.GraphModule,
54+
node: torch.fx.Node,
55+
value: Number,
56+
) -> torch.fx.Node:
57+
input_node = node.args[0]
58+
if not isinstance(input_node, torch.fx.Node):
59+
raise RuntimeError("Expected scalar op input to be an FX node.")
60+
61+
input_value = input_node.meta["val"]
62+
tensor = torch.tensor(value, dtype=input_value.dtype, device=input_value.device)
63+
name = self._get_new_constant_name(graph_module)
64+
first_placeholder = next(
65+
graph_node
66+
for graph_node in graph_module.graph.nodes
67+
if graph_node.op == "placeholder"
68+
)
69+
with graph_module.graph.inserting_before(first_placeholder):
70+
return create_constant_placeholder(
71+
self.exported_program,
72+
graph_module.graph,
73+
name,
74+
InputKind.CONSTANT_TENSOR,
75+
tensor,
76+
)
77+
78+
def _get_new_constant_name(self, graph_module: torch.fx.GraphModule) -> str:
79+
prefix = "_tensor_constant_"
80+
existing_names = {node.name for node in graph_module.graph.nodes}
81+
existing_names.update(self.exported_program.constants)
82+
existing_names.update(self.exported_program.state_dict)
83+
index = 0
84+
while f"{prefix}{index}" in existing_names:
85+
index += 1
86+
return f"{prefix}{index}"
87+
88+
def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
89+
modified = False
90+
91+
for node in list(graph_module.graph.nodes):
92+
if (
93+
node.op != "call_function"
94+
or node.target not in self.scalar_to_tensor_ops
95+
or len(node.args) != 2
96+
or not isinstance(node.args[0], torch.fx.Node)
97+
or not isinstance(node.args[1], Number)
98+
):
99+
continue
100+
101+
input_value = node.args[0].meta.get("val")
102+
output_value = node.meta.get("val")
103+
if (
104+
not isinstance(input_value, torch.Tensor)
105+
or not isinstance(output_value, torch.Tensor)
106+
or input_value.dtype != output_value.dtype
107+
):
108+
continue
109+
110+
tensor_arg = self._create_constant_node(graph_module, node, node.args[1])
111+
node.args = (node.args[0], tensor_arg)
112+
node.target = self.scalar_to_tensor_ops[node.target]
113+
modified = True
114+
115+
graph_module.graph.eliminate_dead_code()
116+
graph_module.graph.lint()
117+
graph_module.recompile()
118+
119+
return PassResult(graph_module, modified)

‎backends/xnnpack/partition/config/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
MeanDimConfig,
4444
MinimumConfig,
4545
MulConfig,
46+
MulScalarConfig,
4647
NegConfig,
4748
PermuteConfig,
4849
PowConfig,
@@ -107,6 +108,7 @@
107108
MinimumConfig,
108109
MMConfig,
109110
MulConfig,
111+
MulScalarConfig,
110112
NegConfig,
111113
PermuteConfig,
112114
PowConfig,

‎backends/xnnpack/partition/config/generic_node_configs.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
# pyre-unsafe
99

1010
import logging
11+
from numbers import Number
1112
from typing import cast, List, Optional
1213

1314
import numpy as np
@@ -439,6 +440,38 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]:
439440
return [ConfigPrecisionType.FP32, ConfigPrecisionType.STATIC_QUANT]
440441

441442

443+
class MulScalarConfig(GenericNodePartitionerConfig):
444+
target_name = "mul.Scalar"
445+
446+
def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool:
447+
if node.graph is not ep.graph:
448+
return False
449+
450+
if not self.check_common_constraints(node, ep):
451+
return False
452+
453+
if (
454+
len(node.args) != 2
455+
or not isinstance(node.args[0], torch.fx.Node)
456+
or not isinstance(node.args[1], Number)
457+
):
458+
return False
459+
460+
input_value = node.args[0].meta.get("val")
461+
output_value = node.meta.get("val")
462+
return (
463+
isinstance(input_value, torch.Tensor)
464+
and isinstance(output_value, torch.Tensor)
465+
and input_value.dtype == output_value.dtype
466+
)
467+
468+
def supported_precision_types(self) -> List[ConfigPrecisionType]:
469+
return [ConfigPrecisionType.FP32]
470+
471+
def get_original_aten(self) -> Optional[torch._ops.OpOverload]:
472+
return torch.ops.aten.mul.Scalar
473+
474+
442475
class MaximumConfig(GenericNodePartitionerConfig):
443476
target_name = "maximum.default"
444477

‎backends/xnnpack/test/ops/test_multiply.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
# Copyright (c) Meta Platforms, Inc. and affiliates.
22
# All rights reserved.
3+
# Copyright 2026 Arm Limited and/or its affiliates.
34
#
45
# This source code is licensed under the BSD-style license found in the
56
# LICENSE file in the root directory of this source tree.
@@ -29,6 +30,10 @@ def forward(self, x, y):
2930
z = torch.mul(x, y) * torch.functional.torch.mul(x, y)
3031
return z
3132

33+
class MulScalar(torch.nn.Module):
34+
def forward(self, x):
35+
return torch.ops.aten.mul.Scalar(x, 0.5)
36+
3237
class MulRelu(torch.nn.Module):
3338
def forward(self, x, y):
3439
z = x * y
@@ -58,6 +63,23 @@ def test_fp32_mul(self):
5863
inputs = (torch.randn((1, 3)), torch.randn((4, 3)))
5964
self._test_mul(inputs)
6065

66+
def test_fp32_mul_scalar(self):
67+
(
68+
Tester(self.MulScalar(), (torch.randn(2, 3),))
69+
.export()
70+
.to_edge_transform_and_lower()
71+
.check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
72+
.check_not(
73+
[
74+
"executorch_exir_dialects_edge__ops_aten_mul_Tensor",
75+
"executorch_exir_dialects_edge__ops_aten_mul_Scalar",
76+
]
77+
)
78+
.to_executorch()
79+
.serialize()
80+
.run_method_and_compare_outputs()
81+
)
82+
6183
def test_qs8_mul(self):
6284
inputs = (torch.randn(1, 1, 4, 4), torch.randn(1, 1, 4, 1))
6385
(

0 commit comments

Comments
 (0)