Skip to content

Commit 91a1d5f

Browse files
pssrawatfacebook-github-bot
authored andcommitted
Quantize expand_as through the existing shared-qparams expand path (#24153)
Summary: `SharedQspecQuantizer` already gives `expand` and `expand_copy` shared qparams when a neighbor is quantized, and the lowering then runs `expand_copy` on int8. `expand_as` is not in that list, and it only decomposes to `expand` at `to_edge`, after quantization, so every `expand_as` broadcast stayed in float behind a dequantize and a requantize. `ExpandAsToExpandPass` runs in `transform_for_annotation` and rewrites `expand_as(x, other)` into `expand(x, other.shape)`, emitting `sym_size` nodes for symbolic dims. `expand_as` then gets the same neighbor-driven treatment as `expand`: int8 when its neighbors are quantized, float otherwise. The shape-only `other` operand is no longer an input of the op, so it is never pulled into the shared-qparams cluster. No new pattern check or support-dict entry is added for `expand`, so an expand between float-only ops is not forced to int8. Differential Revision: D127380012
1 parent d7157b7 commit 91a1d5f

5 files changed

Lines changed: 221 additions & 0 deletions

File tree

‎backends/cortex_m/passes/BUCK‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ fbcode_target(_kind = runtime.python_library,
3737
"decompose_hardswish_pass.py",
3838
"decompose_mean_pass.py",
3939
"decompose_sdpa_pass.py",
40+
"expand_as_to_expand_pass.py",
4041
"explicit_layout_pass.py",
4142
"fuse_conv_padding_pass.py",
4243
"initialize_scratch_buffers_pass.py",

‎backends/cortex_m/passes/cortex_m_pass_manager.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
from .decompose_hardswish_pass import DecomposeHardswishPass
5353
from .decompose_mean_pass import DecomposeMeanPass
5454
from .decompose_sdpa_pass import DecomposeSDPAPass
55+
from .expand_as_to_expand_pass import ExpandAsToExpandPass
5556
from .explicit_layout_pass import (
5657
CortexMCanonicalizeViewCopyPermutePass,
5758
CortexMReplaceOpsWithChannelsLastVariants,
@@ -195,6 +196,7 @@ class CortexMPassManager(ExportedProgramPassManager):
195196
ClampHardswishPass,
196197
DecomposeMeanPass,
197198
MatmulToBmmPass,
199+
ExpandAsToExpandPass,
198200
DeduplicateGetAttrPass,
199201
]
200202

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
from typing import cast, Dict
8+
9+
import torch
10+
from executorch.exir.pass_base import ExportPass, NodeMetadata, ProxyValue
11+
12+
from torch._ops import OpOverload
13+
from torch.fx.node import Argument
14+
15+
16+
class ExpandAsToExpandPass(ExportPass):
17+
"""
18+
Rewrites ``aten.expand_as.default`` into ``aten.expand.default`` so it gets
19+
the shared-qparams treatment ``SharedQspecQuantizer`` already gives
20+
``expand``. ``expand_as`` only decomposes at ``to_edge``, after
21+
quantization, and its second operand is used for its shape alone, so it
22+
must not be tied into the shared-qparams cluster.
23+
24+
This is what the ATen composite itself does -- ``expand_as(self, other)``
25+
is ``self.expand_symint(other.sym_sizes())`` -- read off the second
26+
operand a dimension at a time. A dynamic dimension cannot be copied
27+
across as a value, because the shape entry is a symbol rather than a
28+
number; ``call_size_operator_all`` emits the ``sym_size.int`` that reads
29+
it off the tensor at run time.
30+
"""
31+
32+
def call_operator(
33+
self,
34+
op: OpOverload,
35+
args: tuple[Argument, ...],
36+
kwargs: Dict[str, Argument],
37+
meta: NodeMetadata,
38+
) -> ProxyValue:
39+
if op != torch.ops.aten.expand_as.default:
40+
return super().call_operator(op, args, kwargs, meta)
41+
42+
source = cast(ProxyValue, args[0])
43+
other = cast(ProxyValue, args[1])
44+
sizes = self.call_size_operator_all(other, meta)
45+
return super().call_operator(
46+
torch.ops.aten.expand.default, (source, sizes), {}, meta
47+
)

‎backends/cortex_m/test/targets.bzl‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,26 @@ def define_common_targets(is_fbcode = False):
144144

145145

146146

147+
python_pytest(
148+
name = "test_expand_lowering",
149+
srcs = ["test_expand_lowering.py"],
150+
compile = "with-source",
151+
typing = False,
152+
deps = [
153+
"//caffe2:torch",
154+
"//pytorch/ao:torchao", # @manual
155+
"//executorch/backends/cortex_m:edge_compile_config",
156+
"//executorch/backends/cortex_m/ops:ops",
157+
"//executorch/backends/cortex_m/passes:cortex_passes",
158+
"//executorch/backends/cortex_m/quantizer:quantizer",
159+
"//executorch/backends/test/harness:tester",
160+
"//executorch/exir:lib",
161+
"//executorch/exir/dialects:lib",
162+
"fbsource//third-party/pypi/pytest:pytest",
163+
":tester",
164+
],
165+
)
166+
147167
python_pytest(
148168
name = "test_fuse_conv_padding",
149169
srcs = ["test_fuse_conv_padding.py"],
Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
import pytest
8+
import torch
9+
from executorch.backends.cortex_m.edge_compile_config import (
10+
cortex_m_edge_compile_config,
11+
)
12+
from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager
13+
from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer
14+
from executorch.backends.cortex_m.test.tester import CortexMTester
15+
from executorch.backends.test.harness.stages import StageType
16+
from executorch.exir import to_edge_transform_and_lower
17+
from executorch.exir.dialects._ops import ops as exir_ops
18+
from torch.export import Dim, export
19+
from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
20+
21+
22+
class ExpandBetweenQuantizedOps(torch.nn.Module):
23+
def forward(self, x, y, z):
24+
return (x * y).expand(4, 8) + z
25+
26+
27+
class ExpandAsBetweenQuantizedOps(torch.nn.Module):
28+
def forward(self, x, y, z):
29+
return (x * y).expand_as(z) + z
30+
31+
32+
class ExpandBetweenFloatOps(torch.nn.Module):
33+
def forward(self, x, z):
34+
return torch.cumsum(torch.cumsum(x, dim=1).expand(4, 8), dim=0)
35+
36+
37+
class ExpandAsBetweenFloatOps(torch.nn.Module):
38+
def forward(self, x, z):
39+
return torch.cumsum(torch.cumsum(x, dim=1).expand_as(z), dim=0)
40+
41+
42+
class ExpandAsWithQuantizedShapeOperand(torch.nn.Module):
43+
def forward(self, x, y, z):
44+
other = z + z
45+
return (x * y).expand_as(other) + other
46+
47+
48+
def quantize(model, inputs, dynamic_shapes=None):
49+
captured = export(model, inputs, dynamic_shapes=dynamic_shapes).module()
50+
prepared = prepare_pt2e(captured, CortexMQuantizer())
51+
prepared(*inputs)
52+
return convert_pt2e(prepared)
53+
54+
55+
def is_qdq(node) -> bool:
56+
return "quantize_per_tensor" in str(node.target)
57+
58+
59+
def single_expand(gm):
60+
(expand,) = gm.graph.find_nodes(
61+
op="call_function", target=torch.ops.aten.expand.default
62+
)
63+
return expand
64+
65+
66+
@pytest.mark.parametrize(
67+
"model,inputs",
68+
[
69+
(
70+
ExpandBetweenQuantizedOps(),
71+
(torch.randn(1, 8), torch.randn(1, 8), torch.randn(4, 8)),
72+
),
73+
(
74+
ExpandBetweenQuantizedOps(),
75+
(torch.randn(8), torch.randn(8), torch.randn(4, 8)),
76+
),
77+
(
78+
ExpandAsBetweenQuantizedOps(),
79+
(torch.randn(1, 8), torch.randn(1, 8), torch.randn(4, 8)),
80+
),
81+
(
82+
ExpandAsBetweenQuantizedOps(),
83+
(torch.randn(8), torch.randn(8), torch.randn(4, 8)),
84+
),
85+
],
86+
ids=["expand", "expand_rank_change", "expand_as", "expand_as_rank_change"],
87+
)
88+
def test_expand_with_quantized_neighbors_runs_on_int8(model, inputs):
89+
tester = CortexMTester(model.eval(), inputs)
90+
tester.quantize().export().to_edge().run_passes()
91+
program = tester.get_artifact(StageType.RUN_PASSES).exported_program()
92+
93+
(expand,) = program.graph.find_nodes(
94+
op="call_function", target=exir_ops.edge.aten.expand_copy.default
95+
)
96+
assert expand.args[0].meta["val"].dtype == torch.int8
97+
tester.run_method_and_compare_outputs(inputs=inputs, qtol=1)
98+
99+
100+
@pytest.mark.parametrize(
101+
"model", [ExpandBetweenFloatOps(), ExpandAsBetweenFloatOps()], ids=["expand", "expand_as"]
102+
)
103+
def test_expand_between_float_ops_stays_float(model):
104+
expand = single_expand(quantize(model, (torch.randn(1, 8), torch.randn(4, 8))))
105+
106+
assert not is_qdq(expand.args[0])
107+
assert not any(is_qdq(user) for user in expand.users)
108+
109+
110+
def test_expand_as_does_not_share_qparams_with_its_shape_operand():
111+
inputs = (torch.randn(1, 8), torch.randn(1, 8), 10 * torch.randn(4, 8))
112+
expand = single_expand(quantize(ExpandAsWithQuantizedShapeOperand(), inputs))
113+
114+
dequantize = expand.args[0]
115+
assert is_qdq(dequantize)
116+
assert dequantize.args[0].args[0].target == torch.ops.aten.mul.Tensor
117+
(quantize_after_expand,) = expand.users
118+
(add,) = [
119+
user
120+
for dq in quantize_after_expand.users
121+
for user in dq.users
122+
if user.target == torch.ops.aten.add.Tensor
123+
]
124+
other_dequantize = add.args[1]
125+
assert is_qdq(other_dequantize)
126+
assert dequantize.args[1] != other_dequantize.args[1]
127+
assert all(isinstance(size, int) for size in expand.args[1])
128+
129+
130+
def test_expand_as_with_dynamic_shape_reads_the_size_at_runtime():
131+
inputs = (torch.randn(1, 8), torch.randn(1, 8), torch.randn(4, 8))
132+
dynamic_shapes = ({}, {}, {0: Dim("B")})
133+
quantized = quantize(ExpandAsBetweenQuantizedOps(), inputs, dynamic_shapes)
134+
expand = single_expand(quantized)
135+
136+
assert any(
137+
isinstance(size, torch.fx.Node) and size.target == torch.ops.aten.sym_size.int
138+
for size in expand.args[1]
139+
)
140+
assert is_qdq(expand.args[0])
141+
assert any(is_qdq(user) for user in expand.users)
142+
143+
edge = to_edge_transform_and_lower(
144+
export(quantized, inputs, dynamic_shapes=dynamic_shapes),
145+
compile_config=cortex_m_edge_compile_config(),
146+
transform_passes=CortexMPassManager(),
147+
)
148+
(expand_copy,) = edge.exported_program().graph.find_nodes(
149+
op="call_function", target=exir_ops.edge.aten.expand_copy.default
150+
)
151+
assert expand_copy.args[0].meta["val"].dtype == torch.int8

0 commit comments

Comments
 (0)