diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index ede60fbcbee..a7a21118dd8 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 .quantized_clamp_activation_pass import QuantizedClampActivationPass from .replace_quant_nodes_pass import ReplaceQuantNodesPass @@ -35,6 +36,9 @@ class CortexMPassManager(PassManager): pass_list: list[PassClass] = [ + # Fold constant scales (e.g. attention /sqrt(d)) into the adjacent + # quantize before its scale is folded into op meta. + FoldScaleIntoQuantizePass, # Run before folding so qparams attach to max_pool2d values, not tuple + getitem. RemoveGetItemPass, FoldAndAnnotateQParamsPass, 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..25f3f9b5549 --- /dev/null +++ b/backends/cortex_m/passes/fold_scale_into_quantize_pass.py @@ -0,0 +1,81 @@ +# 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. + +from typing import Optional, Set, Type + +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 +_DIV = exir_ops.edge.aten.div.Tensor +_MUL = exir_ops.edge.aten.mul.Tensor + + +class FoldScaleIntoQuantizePass(ExportPass): + """Fold a constant elementwise scale (``x / c`` or ``x * c``) into the scale + of the per-tensor quantize that consumes it, then drop the div/mul. + + Because ``quantize(x / c, scale=S) == quantize(x, scale=S*c)`` and + ``quantize(x * c, scale=S) == quantize(x, scale=S/c)`` produce identical int8 + values, the constant scale can be absorbed into the adjacent quantize with no + numerical change. This erases the attention-score ``/sqrt(d)`` scale -- an + fp32 div that otherwise stays between the QK^T bmm and softmax -- so the + attention chain is int8 through softmax. + + Runs before ``FoldAndAnnotateQParamsPass`` while the softmax-input quantize is + still an explicit ``quantized_decomposed.quantize_per_tensor`` node. + """ + + _passes_required_after: Set[Type[ExportPass]] = set() + + 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): + if node.op != "call_function" or node.target not in (_DIV, _MUL): + continue + + scaled, const = node.args[0], node.args[1] + if not (isinstance(scaled, Node) and isinstance(const, Node)): + continue + if not is_param_node(ep, const): + continue + const_t = get_param_tensor(ep, const) + if const_t is None or const_t.numel() != 1: + continue + c = float(const_t.reshape(-1)[0]) + if c == 0.0: + continue + + users = list(node.users) + if len(users) != 1 or users[0].target != _QUANTIZE: + continue + + quantize = users[0] + scale = quantize.args[1] + new_scale = scale * c if node.target is _DIV else scale / c + quantize.update_arg(1, new_scale) + quantize.replace_input_with(node, scaled) + graph_module.graph.erase_node(node) + modified = True + + if modified: + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, modified) 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, + )