diff --git a/backends/cortex_m/passes/BUCK b/backends/cortex_m/passes/BUCK index d301e14823c..2a66ca3e80c 100644 --- a/backends/cortex_m/passes/BUCK +++ b/backends/cortex_m/passes/BUCK @@ -35,6 +35,7 @@ fbcode_target(_kind = runtime.python_library, "cortex_m_pass_manager.py", "decompose_hardswish_pass.py", "decompose_mean_pass.py", + "fold_scale_into_quantize_pass.py", "matmul_to_bmm_pass.py", "quantized_clamp_activation_pass.py", "scratch_buffer_sizes.py", diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index 892baf136ed..b66de52ff7d 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/backends/cortex_m/passes/cortex_m_pass_manager.py @@ -27,6 +27,7 @@ from .clamp_hardswish_pass import ClampHardswishPass from .decompose_hardswish_pass import DecomposeHardswishPass from .decompose_mean_pass import DecomposeMeanPass +from .fold_scale_into_quantize_pass import FoldScaleIntoQuantizePass from .matmul_to_bmm_pass import MatmulToBmmPass from .quantized_clamp_activation_pass import QuantizedClampActivationPass from .replace_quant_nodes_pass import ReplaceQuantNodesPass @@ -39,6 +40,9 @@ class CortexMPassManager(PassManager): # Run before folding so qparams attach to max_pool2d values, not tuple + getitem. RemoveGetItemPass, FoldAndAnnotateQParamsPass, + # After qparam-folding: drop a constant scale (e.g. attention /sqrt(d)) + # the quantizer already absorbed into the adjacent quantize scale. + FoldScaleIntoQuantizePass, ReplaceScalarWithTensorArgPass, ReplaceQuantNodesPass, ActivationFusionPass, diff --git a/backends/cortex_m/passes/fold_scale_into_quantize_pass.py b/backends/cortex_m/passes/fold_scale_into_quantize_pass.py new file mode 100644 index 00000000000..3f3b742e9b7 --- /dev/null +++ b/backends/cortex_m/passes/fold_scale_into_quantize_pass.py @@ -0,0 +1,129 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import math +from typing import cast, Optional, Tuple + +from executorch.backends.arm._passes.arm_pass_utils import ( + get_param_tensor, + is_param_node, +) +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult +from torch.fx import GraphModule, Node + +_QUANTIZE = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default +_DEQUANTIZE = exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default +_DIV = exir_ops.edge.aten.div.Tensor +_MUL = exir_ops.edge.aten.mul.Tensor + + +class FoldScaleIntoQuantizePass(ExportPass): + """Drop a constant elementwise scale (``x / c`` or ``x * c``) that the + quantizer already absorbed into the surrounding quantization scale. + + Runs after ``FoldAndAnnotateQParamsPass``. An attention-score ``/sqrt(d)`` + survives that pass as an unannotated fp32 op left inside a + ``dequantize_per_tensor(S_in) -> div/mul(c) -> quantize_per_tensor(S_out)`` + sandwich: the observer that set ``S_out`` saw the post-scale range, so + ``S_out == S_in / c`` (div) or ``S_in * c`` (mul) and the two zero-points + match. Under that relation the sandwich is the identity on the int8 values + (``quantize(dequantize(q, S_in) / c, S_out) == q``), so it is removed and the + producer's int8 wired straight into the consumer. + + The pass rewrites no quantization parameters -- it only recognises that the + constant is already in ``S_out`` and deletes the redundant fp32 round-trip. + Because no scale changes, a downstream SharedQspec consumer (view/reshape/ + pool) that shares ``S_out`` is never disturbed. If the ``S_out == S_in/c`` + (or ``*c``) and matching-zero-point relation does not hold -- per-channel + qparams, a non-scalar constant, or a scale the quantizer did not absorb -- + the fold is skipped and the fp32 op is left in place. + + A constant add/sub cannot fold this way: an additive shift moves the affine + quantization parameters by a generally non-integer amount, so no + integer-preserving requantize identity exists -- a separate, lossy transform, + and no current model needs it. + """ + + def __init__(self, exported_program: Optional[ExportedProgram] = None) -> None: + super().__init__() + self.exported_program = exported_program + + def call(self, graph_module: GraphModule) -> PassResult: + ep = self.exported_program + if ep is None: + return PassResult(graph_module, False) + + modified = False + for node in list(graph_module.graph.nodes): + match = self._absorbed_scale_sandwich(node, ep) + if match is None: + continue + dequantize, quantize, producer = match + quantize.replace_all_uses_with(producer) + graph_module.graph.erase_node(quantize) + graph_module.graph.erase_node(node) + graph_module.graph.erase_node(dequantize) + modified = True + + if modified: + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, modified) + + def _absorbed_scale_sandwich( + self, node: Node, ep: ExportedProgram + ) -> Optional[Tuple[Node, Node, Node]]: + """Return ``(dequantize, quantize, producer)`` when ``node`` is a + constant div/mul the quantizer already folded into the surrounding + quantize scale, so the sandwich can be deleted bit-exactly; else None.""" + if node.op != "call_function" or node.target not in (_DIV, _MUL): + return None + scaled, const = node.args[0], node.args[1] + if not (isinstance(scaled, Node) and isinstance(const, Node)): + return None + # A dequantize -> op -> quantize sandwich, each side used only by the op. + if scaled.target is not _DEQUANTIZE or len(scaled.users) != 1: + return None + # A dequantize's input is always a Node. + producer = cast(Node, scaled.args[0]) + users = list(node.users) + if len(users) != 1 or users[0].target is not _QUANTIZE: + return None + quantize = users[0] + + c = self._scalar_constant(const, ep) + if c is None: + return None + + # (de)quantize args: (input, scale, zero_point, qmin, qmax, dtype). + s_in, zp_in = scaled.args[1], scaled.args[2] + s_out, zp_out = quantize.args[1], quantize.args[2] + if not (isinstance(s_in, float) and isinstance(s_out, float)): + return None + expected = s_in / c if node.target is _DIV else s_in * c + # Bit-exact only if the quantizer folded c into s_out, the zero-points + # match, and the quantize's clamp range/dtype (args 3..5) match the + # dequantize's -- so the removed requantize is a no-op on the producer int8. + if ( + zp_in != zp_out + or scaled.args[3:6] != quantize.args[3:6] + or not math.isclose(s_out, expected, rel_tol=1e-6) + ): + return None + return scaled, quantize, producer + + def _scalar_constant(self, const: Node, ep: ExportedProgram) -> Optional[float]: + if not is_param_node(ep, const): + return None + const_t = get_param_tensor(ep, const) + if const_t is None or const_t.numel() != 1: + return None + c = float(const_t.reshape(-1)[0]) + if c == 0.0 or not math.isfinite(c): + return None + return c diff --git a/backends/cortex_m/test/ops/test_attention_scale.py b/backends/cortex_m/test/ops/test_attention_scale.py new file mode 100644 index 00000000000..1636b128331 --- /dev/null +++ b/backends/cortex_m/test/ops/test_attention_scale.py @@ -0,0 +1,85 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import math + +import torch +from executorch.backends.arm.test.common import parametrize +from executorch.backends.cortex_m.test.tester import CortexMTester, McuTestCase + + +class CortexMScaledAttentionDiv(torch.nn.Module): + """``softmax(bmm(q, k^T) / sqrt(d))`` -- the attention-score scale is an fp32 + ``aten.div.Tensor`` by a constant that otherwise stays between the QK^T bmm + and softmax. It must fold into the softmax-input quantize scale + (``quantize(x / c, S) == quantize(x, S*c)``) so no fp32 div remains and the + chain lowers to ``cortex_m.quantized_batch_matmul`` + ``cortex_m.softmax``. + """ + + ops_before_transforms = { + "executorch_exir_dialects_edge__ops_aten_bmm_default": 1, + "executorch_exir_dialects_edge__ops_aten_div_Tensor": 1, + "executorch_exir_dialects_edge__ops_aten__softmax_default": 1, + } + + ops_after_transforms = { + "executorch_exir_dialects_edge__ops_cortex_m_quantized_batch_matmul_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_softmax_default": 1, + "executorch_exir_dialects_edge__ops_aten_div_Tensor": 0, + "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 0, + } + + def forward(self, q, k): + scores = torch.bmm(q, k.transpose(-2, -1)) / math.sqrt(q.shape[-1]) + return torch.softmax(scores, dim=-1) + + +class CortexMScaledAttentionMul(torch.nn.Module): + """Same, but the scale is applied as ``* (1/sqrt(d))`` -- an + ``aten.mul.Tensor`` by a constant, folded via ``quantize(x*c, S) == + quantize(x, S/c)``. + """ + + ops_before_transforms = { + "executorch_exir_dialects_edge__ops_aten_bmm_default": 1, + "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 1, + "executorch_exir_dialects_edge__ops_aten__softmax_default": 1, + } + + ops_after_transforms = { + "executorch_exir_dialects_edge__ops_cortex_m_quantized_batch_matmul_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_softmax_default": 1, + "executorch_exir_dialects_edge__ops_aten_div_Tensor": 0, + "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 0, + } + + def forward(self, q, k): + scores = torch.bmm(q, k.transpose(-2, -1)) * (1.0 / math.sqrt(q.shape[-1])) + return torch.softmax(scores, dim=-1) + + +test_cases = { + "scaled_attn_div": McuTestCase( + CortexMScaledAttentionDiv(), + (torch.rand(1, 8, 16), torch.rand(1, 8, 16)), + ), + "scaled_attn_mul": McuTestCase( + CortexMScaledAttentionMul(), + (torch.rand(1, 8, 16), torch.rand(1, 8, 16)), + ), +} + + +@parametrize("test_case", test_cases) +def test_dialect_attention_scale(test_case, cortex_m_target): + tester = CortexMTester( + test_case.model, test_case.example_inputs, target_config=cortex_m_target + ) + tester.test_dialect( + test_case.model.ops_before_transforms, + test_case.model.ops_after_transforms, + qtol=1, + ) diff --git a/backends/cortex_m/test/ops/test_fold_scale_into_quantize.py b/backends/cortex_m/test/ops/test_fold_scale_into_quantize.py new file mode 100644 index 00000000000..fb16bf73a23 --- /dev/null +++ b/backends/cortex_m/test/ops/test_fold_scale_into_quantize.py @@ -0,0 +1,295 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Pass-level tests for ``FoldScaleIntoQuantizePass``. + +These run the pass on the post-``FoldAndAnnotateQParamsPass`` graph (its real +input state) and check the graph rewrite and its guards directly, without the +full lowering that ``test_attention_scale.py`` exercises end to end. +""" + +import copy +import math + +import torch +from executorch.backends.arm._passes import FoldAndAnnotateQParamsPass +from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import ( + get_input_qparams, + get_output_qparams, +) +from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager +from executorch.backends.cortex_m.passes.fold_scale_into_quantize_pass import ( + FoldScaleIntoQuantizePass, +) +from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer +from executorch.backends.transforms.remove_getitem_op import RemoveGetItemPass +from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import export +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + +_QD = torch.ops.quantized_decomposed +_QUANTIZE = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default +_EDGE_CFG = EdgeCompileConfig( + _check_ir_validity=False, + _core_aten_ops_exception_list=[torch.ops.aten.max_pool2d.default], +) + + +class _AttnDiv(torch.nn.Module): + def forward(self, q, k): + return torch.softmax( + torch.bmm(q, k.transpose(-2, -1)) / math.sqrt(q.shape[-1]), dim=-1 + ) + + +class _AttnMul(torch.nn.Module): + def forward(self, q, k): + return torch.softmax( + torch.bmm(q, k.transpose(-2, -1)) * (1.0 / math.sqrt(q.shape[-1])), dim=-1 + ) + + +class _PoolScale(torch.nn.Module): + # A constant scale feeding a SharedQspec pool -- the passthrough case where + # folding must not disturb the shared scale. + def forward(self, x): + return torch.nn.functional.max_pool2d(x / 4.0, 2) + + +class _ScaleFirstMul(torch.nn.Module): + # Constant as the FIRST operand: mul(const, x). The constant lifts to a + # placeholder (empty args), so the pass must not read its args[0] before + # confirming the operand is a dequantize. + def __init__(self): + super().__init__() + self.register_buffer("scale", torch.tensor(0.125)) + + def forward(self, q, k): + return torch.softmax(self.scale * torch.bmm(q, k.transpose(-2, -1)), dim=-1) + + +class _MultiHeadMatmulAttn(torch.nn.Module): + # Rank-4 multi-head attention with q @ k^T scaled by 1/sqrt(head_dim), like + # SAM's mask-decoder attention. MatmulToBmmPass rewrites the matmul to a + # quantizable bmm at annotation, so the scale lands in the requantize sandwich + # the fold removes -- the realistic path a bare rank-3 torch.bmm skips. + def __init__(self, dim=32, heads=4): + super().__init__() + self.heads = heads + self.head_dim = dim // heads + self.q_proj = torch.nn.Linear(dim, dim) + self.k_proj = torch.nn.Linear(dim, dim) + self.v_proj = torch.nn.Linear(dim, dim) + + def _heads(self, x): + b, n, c = x.shape + return x.reshape(b, n, self.heads, self.head_dim).transpose(1, 2) + + def forward(self, x): + q = self._heads(self.q_proj(x)) + k = self._heads(self.k_proj(x)) + v = self._heads(self.v_proj(x)) + attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) + attn = torch.softmax(attn, dim=-1) + return (attn @ v).transpose(1, 2).reshape(x.shape) + + +class _RuntimeDiv(torch.nn.Module): + # Divisor is a runtime graph input, not a param/buffer, so is_param_node is + # False -- the scale is non-constant and must not fold. + def forward(self, q, k, d): + return torch.softmax(torch.bmm(q, k.transpose(-2, -1)) / d, dim=-1) + + +_ATTN_INPUTS = (torch.rand(1, 8, 16), torch.rand(1, 8, 16)) +_POOL_INPUTS = (torch.rand(1, 1, 8, 8),) +_MHA_INPUT = (torch.rand(1, 8, 32),) +_RUNTIME_INPUTS = ( + torch.rand(1, 8, 16), + torch.rand(1, 8, 16), + torch.rand(1, 8, 8) + 0.5, +) + + +def _to_edge(model, inputs): + gm = export(model.eval(), inputs, strict=True).module() + gm = prepare_pt2e(gm, CortexMQuantizer()) + gm(*inputs) + gm = convert_pt2e(gm) + return to_edge(export(gm, inputs, strict=True), compile_config=_EDGE_CFG) + + +def _run(edge, extra_passes): + passes = [RemoveGetItemPass, FoldAndAnnotateQParamsPass, *extra_passes] + return CortexMPassManager( + copy.deepcopy(edge).exported_program(), passes=passes + ).transform() + + +def _count(gm, needle): + return sum( + 1 for n in gm.graph.nodes if n.op == "call_function" and needle in str(n.target) + ) + + +def _find(gm, needle): + return next(n for n in gm.graph.nodes if needle in str(n.target)) + + +def test_div_folds_and_removes_the_sandwich(): + edge = _to_edge(_AttnDiv(), _ATTN_INPUTS) + annotated = _run(edge, []) + folded = _run(edge, [FoldScaleIntoQuantizePass]) + + assert _count(annotated.graph_module, "aten.div.Tensor") == 1 + assert _count(folded.graph_module, "aten.div.Tensor") == 0 + # The dequantize/quantize that wrapped the div go with it (one of each). + assert _count(folded.graph_module, "dequantize_per_tensor") == ( + _count(annotated.graph_module, "dequantize_per_tensor") - 1 + ) + # bmm now feeds softmax directly. + assert "bmm" in str(_find(folded.graph_module, "_softmax").args[0].target) + + +def test_mul_folds(): + edge = _to_edge(_AttnMul(), _ATTN_INPUTS) + assert _count(_run(edge, []).graph_module, "aten.mul.Tensor") == 1 + assert ( + _count(_run(edge, [FoldScaleIntoQuantizePass]).graph_module, "aten.mul.Tensor") + == 0 + ) + + +def test_matmul_multihead_attention_scale_folds(): + # SAM-faithful path: the rank-4 q @ k^T matmul becomes a quantizable bmm via + # MatmulToBmmPass at annotation, so the /sqrt(head_dim) scale lands in a + # dequantize -> div -> quantize sandwich the fold removes. A bare rank-3 + # torch.bmm (the tests above) has its output quantized directly, so it does + # not exercise this matmul->bmm dependency -- if that pass regresses, the + # scale would stay fp32 and this test fails. + edge = _to_edge(_MultiHeadMatmulAttn(), _MHA_INPUT) + assert _count(_run(edge, []).graph_module, "aten.div.Tensor") == 1 + assert ( + _count(_run(edge, [FoldScaleIntoQuantizePass]).graph_module, "aten.div.Tensor") + == 0 + ) + + +def test_fold_is_bit_exact(): + # Removing the sandwich is bit-exact iff quantize(dequantize(q) / c) == q for + # every int8 at the calibrated qparams; verify against the real kernels. + edge = _to_edge(_AttnDiv(), _ATTN_INPUTS) + div = _find(_run(edge, []).graph_module, "aten.div.Tensor") + dq, quant = div.args[0], next(iter(div.users)) + s_in, zp_in = dq.args[1], dq.args[2] + s_out, zp_out = quant.args[1], quant.args[2] + + c = math.sqrt(_ATTN_INPUTS[0].shape[-1]) + v = torch.arange(-128, 128, dtype=torch.int8) + deq = _QD.dequantize_per_tensor.default(v, s_in, zp_in, -128, 127, torch.int8) + requant = _QD.quantize_per_tensor.default( + deq / c, s_out, zp_out, -128, 127, torch.int8 + ) + assert torch.equal(requant, v) + + +def test_sharedqspec_consumer_scale_is_untouched(): + # A constant scale feeding a SharedQspec pool (RFC #19299). The fold + # rewrites no scale, so the pool's shared in/out qparams are identical before + # and after -- the shared cluster can never be corrupted -- and the divide + # still folds. + edge = _to_edge(_PoolScale(), _POOL_INPUTS) + annotated = _run(edge, []) + folded = _run(edge, [FoldScaleIntoQuantizePass]) + + assert _count(folded.graph_module, "aten.div.Tensor") == 0 + pool_before = _find(annotated.graph_module, "max_pool2d") + pool_after = _find(folded.graph_module, "max_pool2d") + # Pin the premise: max_pool2d is a SharedQspec op (input scale == output + # scale). Without this the test would silently stop guarding that case. + assert ( + get_input_qparams(pool_before)[0].scale + == get_output_qparams(pool_before)[0].scale + ) + assert ( + get_input_qparams(pool_before)[0].scale + == get_input_qparams(pool_after)[0].scale + ) + assert ( + get_output_qparams(pool_before)[0].scale + == get_output_qparams(pool_after)[0].scale + ) + + +def _perturb_softmax_quantize(program, arg_index, fn): + for node in program.graph_module.graph.nodes: + if ( + node.op == "call_function" + and node.target is _QUANTIZE + and any("softmax" in str(u.target) for u in node.users) + ): + args = list(node.args) + args[arg_index] = fn(args[arg_index]) + node.args = tuple(args) + + +def test_skips_when_scale_not_absorbed(): + program = _run(_to_edge(_AttnDiv(), _ATTN_INPUTS), []) + _perturb_softmax_quantize(program, 1, lambda s: s * 1.5) + result = FoldScaleIntoQuantizePass(exported_program=program).call( + program.graph_module + ) + assert not result.modified + assert _count(program.graph_module, "aten.div.Tensor") == 1 + + +def test_skips_on_zero_point_mismatch(): + program = _run(_to_edge(_AttnDiv(), _ATTN_INPUTS), []) + _perturb_softmax_quantize(program, 2, lambda z: z + 1) + result = FoldScaleIntoQuantizePass(exported_program=program).call( + program.graph_module + ) + assert not result.modified + assert _count(program.graph_module, "aten.div.Tensor") == 1 + + +def test_constant_first_operand_does_not_crash(): + # `const * tensor` puts the constant (a placeholder with empty args) at + # args[0]; the pass must not read its args[0] before confirming it is a + # dequantize. It should skip, not raise (regression). + program = _run(_to_edge(_ScaleFirstMul(), _ATTN_INPUTS), []) + result = FoldScaleIntoQuantizePass(exported_program=program).call( + program.graph_module + ) + assert not result.modified + assert _count(program.graph_module, "aten.mul.Tensor") == 1 + + +def test_skips_on_dtype_or_range_mismatch(): + # A quantize whose dtype/clamp range (args 3..5) differs from the dequantize + # makes the removed requantize a real conversion, not a no-op (e.g. the int16 + # activation path). The scale/zero-point relation still holds, so only the + # args[3:6] guard prevents the fold. + program = _run(_to_edge(_AttnDiv(), _ATTN_INPUTS), []) + _perturb_softmax_quantize(program, 5, lambda _: torch.int16) + result = FoldScaleIntoQuantizePass(exported_program=program).call( + program.graph_module + ) + assert not result.modified + assert _count(program.graph_module, "aten.div.Tensor") == 1 + + +def test_skips_non_param_divisor(): + # The divisor is a runtime graph input, not a param/buffer, so is_param_node + # is False: the scale is genuinely non-constant and must not fold -- and the + # pass must not crash reading a param tensor for a non-param node. + program = _run(_to_edge(_RuntimeDiv(), _RUNTIME_INPUTS), []) + result = FoldScaleIntoQuantizePass(exported_program=program).call( + program.graph_module + ) + assert not result.modified + assert _count(program.graph_module, "aten.div.Tensor") == 1