Cortex-M: fold constant attention-score scale into the softmax-input quantize

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,

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)

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,
+    )