From ccb8b3bd6fac47e7dfc3f7de10e2916001be260a Mon Sep 17 00:00:00 2001
From: RJ Ascani <rja@meta.com>
Date: Mon, 18 May 2026 12:25:25 -0700
Subject: [PATCH 1/8] Cortex-M backend: dispatch quantized_linear AOT layout on
 target ISA
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

CMSIS-NN's `arm_fully_connected_s8` has three runtime paths, gated by
compile-time `ARM_MATH_MVEI` / `ARM_MATH_DSP`. They split the bias and
input_offset×sum(weight) offset term between two inputs, in
incompatible conventions:

* MVE (Helium): reads `ctx.buf` as a precomputed kernel_sum that must
  already include `input_offset × sum(weight)` and the bias
  contribution. The `bias` argument is `(void)bias;` — ignored.
* DSP / scalar (Armv7E-M, Armv7-M, Armv6-M, Armv8-M Baseline): read
  the `bias` argument directly and fold the input_offset contribution
  at runtime. `ctx.buf` (kernel_sum) is `(void)kernel_sum;` — ignored.

`ConvertToCortexMPass._get_linear_replacement` previously emitted only
the MVE shape (kernel_sum populated, bias=None). On any non-MVE build
the DSP/scalar path started the int32 accumulator at 0 instead of at
`bias + input_offset × sum(weight)`, dropping both the bias and the
offset contribution. The accumulator wound up much smaller than
intended, requantization collapsed it to the output zero point, and
every classifier with a deep, narrow tail produced essentially
uniform near-zero outputs on Cortex-M0/0+/M3/M4/M7/M23/M33 builds —
exactly where `CortexMTargetConfig.backend != cmsis_nn.Backend.MVE`.

Use the target-ISA plumbing added by the CortexMTargetConfig PR
(#19470) to dispatch the right input shape at AOT time: on MVE
targets emit kernel_sum with bias folded in (bias=None); on DSP and
scalar targets emit the raw int32 bias directly (kernel_sum=None).
The CMSIS-NN runtime then matches exactly what it expects — no
redundant copy of the bias in the .pte, no silent miscompute on
target mismatch (the runtime errors loudly if ctx.buf is None on an
MVE build, instead of producing garbage).

Update `quantized_linear_impl` in `operators.py` to mirror the same
contract: dispatch off whichever of kernel_sum / bias is non-None.
Threading happens automatically via `CortexMPassManager`'s signature
injection of `target_config` into the pass's `__init__`.

Add `backends/cortex_m/test/misc/test_quantized_linear_small_magnitude.py`
as a regression. A tiny `nn.Linear(512, 10)` on uniform[0, 0.002]
input is the minimal reproducer for the small-magnitude regime where
the missing offset terms dominate. The dialect test parametrizes over
MVE/DSP/scalar target configs; the implementation test runs against
whatever path the runner build matches.

Identified by bisecting STResNet Pico's int8-output collapse on
Corstone-300. The same collapse explains the historical MV2 / MV3
"deep classifier PTQ flakiness" xfails — both classifiers have
small-magnitude inputs to their final quantized_linear and likely
hit the same bug on any non-MVE deployment target.

Authored with Claude.
---
 backends/cortex_m/ops/operators.py            |  13 +-
 .../passes/convert_to_cortex_m_pass.py        |  92 +++++++---
 .../test_quantized_linear_small_magnitude.py  | 173 ++++++++++++++++++
 3 files changed, 245 insertions(+), 33 deletions(-)
 create mode 100644 backends/cortex_m/test/misc/test_quantized_linear_small_magnitude.py

diff --git a/backends/cortex_m/ops/operators.py b/backends/cortex_m/ops/operators.py
index 2c35ed8730b..f617a960189 100644
--- a/backends/cortex_m/ops/operators.py
+++ b/backends/cortex_m/ops/operators.py
@@ -434,8 +434,8 @@ def quantized_linear_meta(
 def quantized_linear_impl(
     input: torch.Tensor,
     weights: torch.Tensor,
-    bias: torch.Tensor,
-    kernel_sum: torch.Tensor,
+    bias: torch.Tensor | None,
+    kernel_sum: torch.Tensor | None,
     input_offset: int,
     filter_offset: int,
     output_offset: int,
@@ -448,10 +448,11 @@ def quantized_linear_impl(
     Functional variant - creates output tensor and calls out variant
     """
 
-    # Leaving both implementations for debugging purposes.
-    compute_using_kernel_sum = True
-
-    if compute_using_kernel_sum:
+    # Mirror CMSIS-NN's arm_fully_connected_s8 contract: the MVE path reads
+    # kernel_sum (ctx.buf) and ignores bias; the DSP and scalar paths read
+    # bias and ignore kernel_sum. The AOT pass populates exactly one of them
+    # based on the target ISA, so dispatch off which one is present.
+    if kernel_sum is not None:
         weights_int32 = weights.to(torch.int32)
 
         input_int32 = input.to(torch.int32)
diff --git a/backends/cortex_m/passes/convert_to_cortex_m_pass.py b/backends/cortex_m/passes/convert_to_cortex_m_pass.py
index 418f6cd63ff..e8b9583889f 100644
--- a/backends/cortex_m/passes/convert_to_cortex_m_pass.py
+++ b/backends/cortex_m/passes/convert_to_cortex_m_pass.py
@@ -7,10 +7,12 @@
 
 import executorch.backends.cortex_m.ops.operators  # noqa
 
+import cmsis_nn  # type: ignore[import-not-found, import-untyped]
 import torch
 import torch.fx
 from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
 from executorch.backends.cortex_m.passes.passes_utils import quantize_multiplier_aot
+from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig
 
 from executorch.backends.transforms.utils import (
     create_constant_placeholder,
@@ -20,6 +22,7 @@
 
 from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass
 from executorch.exir.dialects._ops import ops as exir_ops
+from torch.export import ExportedProgram
 from torch.export.graph_signature import InputKind
 from torch.fx.passes.infra.pass_manager import PassResult
 
@@ -33,21 +36,35 @@ class ConvertToCortexMPass(XNNPACKPass):
     by call_operator.
     """
 
-    def _compute_kernel_sum(self, weights, bias, input_offset, weight_offset):
-        """
-        Computes the precomputed kernel sum term (bias optional)
-            a * sum_j(wij + b) + ci
-
-        for i = (1, ..., n), where j indexes the input activations.
+    def __init__(
+        self,
+        exported_program: ExportedProgram,
+        target_config: CortexMTargetConfig | None = None,
+    ) -> None:
+        super().__init__(exported_program)
+        # Default mirrors CortexMPassManager: MVE-capable M55 (the previous
+        # behavior for any caller that constructs the pass without a config).
+        self._target_config = target_config or CortexMTargetConfig(cpu=CortexM.M55)
+
+    @property
+    def target_config(self) -> CortexMTargetConfig:
+        return self._target_config
+
+    def _compute_kernel_sum(self, weights, bias_int32, neg_input_zp, neg_weight_zp):
+        """Precompute the MVE kernel_sum term:  a * sum_j(wij + b) + ci
+
+        Where `a = -input_zp` and `b = -weight_zp` per CMSIS-NN convention.
+        Parameter names use the `neg_*_zp` form to keep that sign explicit at
+        every call site. Bias is optional; pass None for an unbiased Linear.
         """
         weights_transposed = weights.T
         weights_int32 = weights_transposed.to(torch.int32)
-        offset_weights = weights_int32 + weight_offset
+        offset_weights = weights_int32 + neg_weight_zp
         kernel_sum = torch.sum(offset_weights, dim=0, keepdim=True, dtype=torch.int32)
-        kernel_sum_offset = kernel_sum * input_offset
+        kernel_sum_offset = kernel_sum * neg_input_zp
 
-        if bias is not None:
-            kernel_sum_offset += bias
+        if bias_int32 is not None:
+            kernel_sum_offset += bias_int32
 
         return kernel_sum_offset
 
@@ -96,37 +113,58 @@ def _get_linear_replacement(self, node):
         output_min = node.meta["output_qparams"][0].qmin
         output_max = node.meta["output_qparams"][0].qmax
 
+        # CMSIS-NN's FC path treats weights as per-tensor symmetric (single
+        # `filter_offset`, single multiplier/shift). The non-zero-weight-zp
+        # paths in `arm_nn_vec_mat_mult_t_s8.c` exist but are untested in this
+        # backend — fail loudly if the quantizer ever produces asymmetric
+        # weights so we don't silently land on that codepath.
+        if weight_zp != 0:
+            raise NotImplementedError(
+                f"cortex_m::quantized_linear assumes symmetric weight "
+                f"quantization (weight_zp == 0); got weight_zp={weight_zp}"
+            )
+
         quantized_multiplier, quantized_shift = quantize_multiplier_aot(
             (input_scale * weight_scale) / output_scale
         )
 
-        # TODO: Add support for configuring the backend to support other extensions.
-        # Kernel sum is only used in the CMSIS-NN implementation for the MVE extension,
-        # so this should be optional.
+        # CMSIS-NN's MVE `arm_fully_connected_s8` path reads a precomputed
+        # kernel_sum (input_offset×sum(weight) + bias) from ctx.buf and
+        # ignores the bias argument. The DSP and scalar paths do the opposite
+        # — they read the bias argument at runtime and ignore ctx.buf
+        # (see arm_nn_vec_mat_mult_t_s8.c). Pick the right input format here
+        # based on the target ISA so the runtime gets exactly what it expects.
         weights = node.args[1]
         weights_tensor = get_param_tensor(self.exported_program, weights)
+        bias_node = node.args[2] if len(node.args) > 2 else None
         bias_tensor = (
-            get_param_tensor(self.exported_program, node.args[2])
-            if len(node.args) > 2
+            get_param_tensor(self.exported_program, bias_node)
+            if bias_node is not None
             else None
         )
-        kernel_sum_tensor = self._compute_kernel_sum(
-            weights_tensor, bias_tensor, -input_zp, -weight_zp
-        )
-        with node.graph.inserting_after(weights):
-            kernel_sum = create_constant_placeholder(
-                self.exported_program,
-                node.graph,
-                node.name + "_kernel_sum",
-                InputKind.PARAMETER,
-                kernel_sum_tensor,
+
+        if self.target_config.backend == cmsis_nn.Backend.MVE:
+            kernel_sum_tensor = self._compute_kernel_sum(
+                weights_tensor, bias_tensor, -input_zp, -weight_zp
             )
+            with node.graph.inserting_after(weights):
+                kernel_sum_arg = create_constant_placeholder(
+                    self.exported_program,
+                    node.graph,
+                    node.name + "_kernel_sum",
+                    InputKind.PARAMETER,
+                    kernel_sum_tensor,
+                )
+            bias_arg = None
+        else:
+            kernel_sum_arg = None
+            bias_arg = bias_node
 
         args = (
             node.args[0],
             weights,
-            None,
-            kernel_sum,
+            bias_arg,
+            kernel_sum_arg,
             -input_zp,
             -weight_zp,
             output_zp,
diff --git a/backends/cortex_m/test/misc/test_quantized_linear_small_magnitude.py b/backends/cortex_m/test/misc/test_quantized_linear_small_magnitude.py
new file mode 100644
index 00000000000..b07f9963e66
--- /dev/null
+++ b/backends/cortex_m/test/misc/test_quantized_linear_small_magnitude.py
@@ -0,0 +1,173 @@
+# 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.
+"""Regression test for the cortex_m::quantized_linear bias/kernel_sum bug.
+
+CMSIS-NN's `arm_fully_connected_s8` has three runtime paths, gated by
+compile-time ARM_MATH_MVEI / ARM_MATH_DSP:
+
+* MVE: reads ctx.buf (precomputed kernel_sum that includes bias plus
+  input_offset x sum(weight)), ignores the bias argument.
+* DSP / scalar: read the bias argument directly, ignore ctx.buf.
+
+`ConvertToCortexMPass._get_linear_replacement` selects which input format
+to emit based on `CortexMTargetConfig.backend`. Before the fix, the pass
+unconditionally emitted kernel_sum + None-bias, which on a non-MVE build
+silently dropped both the bias and the input-offset term. The bug only
+showed up when those terms dominated the int32 accumulator, i.e. on
+small-magnitude inputs.
+
+Coverage:
+
+* `test_dialect_small_magnitude_linear` runs each ISA through the Python
+  op impl and checks that bias=True and bias=False variants both round
+  to the same int8 outputs as the float reference.
+* `test_aot_graph_shape_small_magnitude_linear` inspects the post-pass
+  graph and asserts the bias/kernel_sum arg positions match the ISA
+  convention -- this is the direct regression check.
+* `test_implementation_small_magnitude_linear` runs the bias=True case
+  through the default (M55, MVE) build path so the impl test exercises
+  the kernel_sum codepath in simulation.
+"""
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+from executorch.backends.arm.test.common import parametrize
+from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig
+from executorch.backends.cortex_m.test.tester import CortexMTester, McuTestCase
+from executorch.backends.test.harness.stages import StageType
+from executorch.exir.dialects._ops import ops as exir_ops
+
+torch.manual_seed(0)
+
+
+class _SmallMagnitudeLinear(nn.Module):
+    ops_before_transforms = {
+        "executorch_exir_dialects_edge__ops_aten_linear_default": 1,
+        "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 2,
+        "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 4,
+    }
+    ops_after_transforms = {
+        "executorch_exir_dialects_edge__ops_cortex_m_quantized_linear_default": 1,
+        "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1,
+        "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1,
+    }
+
+    def __init__(self, bias: bool = True):
+        super().__init__()
+        self.fc = nn.Linear(512, 10, bias=bias)
+
+    def forward(self, x):
+        return self.fc(x)
+
+
+class _SmallMagnitudeLinearNoBias(_SmallMagnitudeLinear):
+    ops_before_transforms = {
+        "executorch_exir_dialects_edge__ops_aten_linear_default": 1,
+        "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 2,
+        "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 3,
+    }
+
+    def __init__(self):
+        super().__init__(bias=False)
+
+
+def _make_input():
+    return torch.rand(1, 512) * 0.002
+
+
+_calibration_samples = [(_make_input(),) for _ in range(8)]
+
+
+@dataclass(frozen=True)
+class _IsaVariant:
+    case: McuTestCase
+    target_config: CortexMTargetConfig
+    uses_kernel_sum: bool
+
+
+def _variant(model_cls, cpu: CortexM, uses_kernel_sum: bool) -> _IsaVariant:
+    return _IsaVariant(
+        case=McuTestCase(
+            model=model_cls().eval(),
+            example_inputs=lambda: (_make_input(),),
+        ),
+        target_config=CortexMTargetConfig(cpu=cpu),
+        uses_kernel_sum=uses_kernel_sum,
+    )
+
+
+# bias=True covers the regression directly (the bug dropped the bias
+# term); bias=False covers the symmetric case where only the
+# input-offset term is missing on the non-MVE paths.
+test_variants = {
+    "mve_bias": _variant(_SmallMagnitudeLinear, CortexM.M55, uses_kernel_sum=True),
+    "dsp_bias": _variant(_SmallMagnitudeLinear, CortexM.M4, uses_kernel_sum=False),
+    "scalar_bias": _variant(_SmallMagnitudeLinear, CortexM.M0PLUS, uses_kernel_sum=False),
+    "mve_nobias": _variant(_SmallMagnitudeLinearNoBias, CortexM.M55, uses_kernel_sum=True),
+    "dsp_nobias": _variant(_SmallMagnitudeLinearNoBias, CortexM.M4, uses_kernel_sum=False),
+    "scalar_nobias": _variant(
+        _SmallMagnitudeLinearNoBias, CortexM.M0PLUS, uses_kernel_sum=False
+    ),
+}
+
+
+@parametrize("variant", test_variants)
+def test_dialect_small_magnitude_linear(variant: _IsaVariant):
+    tester = CortexMTester(
+        variant.case.model,
+        variant.case.get_example_inputs(),
+        target_config=variant.target_config,
+    )
+    tester.test_dialect(
+        ops_before_transforms=variant.case.model.ops_before_transforms,
+        ops_after_transforms=variant.case.model.ops_after_transforms,
+        qtol=1,
+        calibration_samples=_calibration_samples,
+    )
+
+
+@parametrize("variant", test_variants)
+def test_aot_graph_shape_small_magnitude_linear(variant: _IsaVariant):
+    """Assert the post-pass node args match the ISA's CMSIS-NN convention."""
+    tester = CortexMTester(
+        variant.case.model,
+        variant.case.get_example_inputs(),
+        target_config=variant.target_config,
+    )
+    tester.quantize(None)
+    tester.export()
+    tester.to_edge()
+    tester.run_passes()
+
+    module = tester.get_artifact(StageType.RUN_PASSES).exported_program().module()
+    linear_target = exir_ops.edge.cortex_m.quantized_linear.default
+    linear_nodes = [
+        n for n in module.graph.nodes if n.op == "call_function" and n.target == linear_target
+    ]
+    assert len(linear_nodes) == 1, f"expected one quantized_linear node, got {len(linear_nodes)}"
+    bias_arg, kernel_sum_arg = linear_nodes[0].args[2], linear_nodes[0].args[3]
+
+    if variant.uses_kernel_sum:
+        assert bias_arg is None, "MVE path must not pass bias (CMSIS-NN ignores it)"
+        assert kernel_sum_arg is not None, "MVE path requires precomputed kernel_sum"
+    else:
+        assert kernel_sum_arg is None, "non-MVE path must not pass kernel_sum"
+        # bias is allowed to be None only if the source nn.Linear had bias=False.
+        expects_bias = variant.case.model.fc.bias is not None
+        if expects_bias:
+            assert bias_arg is not None, "non-MVE path with bias must forward bias to CMSIS-NN"
+
+
+def test_implementation_small_magnitude_linear():
+    """Exercise the MVE kernel_sum codepath via the default M55 simulator build."""
+    case = McuTestCase(
+        model=_SmallMagnitudeLinear().eval(),
+        example_inputs=lambda: (_make_input(),),
+    )
+    tester = CortexMTester(case.model, case.get_example_inputs())
+    tester.test_implementation(qtol=1, calibration_samples=_calibration_samples)

From ea5d36c5c4ca2de5eaa2c5e1d4c9c6024beffb78 Mon Sep 17 00:00:00 2001
From: RJ Ascani <rja@meta.com>
Date: Fri, 22 May 2026 13:10:48 -0700
Subject: [PATCH 2/8] Fix lints

---
 .../passes/convert_to_cortex_m_pass.py        |  3 +--
 .../test_quantized_linear_small_magnitude.py  | 24 ++++++++++++++-----
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/backends/cortex_m/passes/convert_to_cortex_m_pass.py b/backends/cortex_m/passes/convert_to_cortex_m_pass.py
index e8b9583889f..79654e59f4e 100644
--- a/backends/cortex_m/passes/convert_to_cortex_m_pass.py
+++ b/backends/cortex_m/passes/convert_to_cortex_m_pass.py
@@ -5,9 +5,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import executorch.backends.cortex_m.ops.operators  # noqa
-
 import cmsis_nn  # type: ignore[import-not-found, import-untyped]
+import executorch.backends.cortex_m.ops.operators  # noqa
 import torch
 import torch.fx
 from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
diff --git a/backends/cortex_m/test/misc/test_quantized_linear_small_magnitude.py b/backends/cortex_m/test/misc/test_quantized_linear_small_magnitude.py
index b07f9963e66..5cfcea74e70 100644
--- a/backends/cortex_m/test/misc/test_quantized_linear_small_magnitude.py
+++ b/backends/cortex_m/test/misc/test_quantized_linear_small_magnitude.py
@@ -107,9 +107,15 @@ def _variant(model_cls, cpu: CortexM, uses_kernel_sum: bool) -> _IsaVariant:
 test_variants = {
     "mve_bias": _variant(_SmallMagnitudeLinear, CortexM.M55, uses_kernel_sum=True),
     "dsp_bias": _variant(_SmallMagnitudeLinear, CortexM.M4, uses_kernel_sum=False),
-    "scalar_bias": _variant(_SmallMagnitudeLinear, CortexM.M0PLUS, uses_kernel_sum=False),
-    "mve_nobias": _variant(_SmallMagnitudeLinearNoBias, CortexM.M55, uses_kernel_sum=True),
-    "dsp_nobias": _variant(_SmallMagnitudeLinearNoBias, CortexM.M4, uses_kernel_sum=False),
+    "scalar_bias": _variant(
+        _SmallMagnitudeLinear, CortexM.M0PLUS, uses_kernel_sum=False
+    ),
+    "mve_nobias": _variant(
+        _SmallMagnitudeLinearNoBias, CortexM.M55, uses_kernel_sum=True
+    ),
+    "dsp_nobias": _variant(
+        _SmallMagnitudeLinearNoBias, CortexM.M4, uses_kernel_sum=False
+    ),
     "scalar_nobias": _variant(
         _SmallMagnitudeLinearNoBias, CortexM.M0PLUS, uses_kernel_sum=False
     ),
@@ -147,9 +153,13 @@ def test_aot_graph_shape_small_magnitude_linear(variant: _IsaVariant):
     module = tester.get_artifact(StageType.RUN_PASSES).exported_program().module()
     linear_target = exir_ops.edge.cortex_m.quantized_linear.default
     linear_nodes = [
-        n for n in module.graph.nodes if n.op == "call_function" and n.target == linear_target
+        n
+        for n in module.graph.nodes
+        if n.op == "call_function" and n.target == linear_target
     ]
-    assert len(linear_nodes) == 1, f"expected one quantized_linear node, got {len(linear_nodes)}"
+    assert (
+        len(linear_nodes) == 1
+    ), f"expected one quantized_linear node, got {len(linear_nodes)}"
     bias_arg, kernel_sum_arg = linear_nodes[0].args[2], linear_nodes[0].args[3]
 
     if variant.uses_kernel_sum:
@@ -160,7 +170,9 @@ def test_aot_graph_shape_small_magnitude_linear(variant: _IsaVariant):
         # bias is allowed to be None only if the source nn.Linear had bias=False.
         expects_bias = variant.case.model.fc.bias is not None
         if expects_bias:
-            assert bias_arg is not None, "non-MVE path with bias must forward bias to CMSIS-NN"
+            assert (
+                bias_arg is not None
+            ), "non-MVE path with bias must forward bias to CMSIS-NN"
 
 
 def test_implementation_small_magnitude_linear():

From 7e2f23562e1feeecb5fcc2433546b0985cbe2010 Mon Sep 17 00:00:00 2001
From: RJ Ascani <rja@meta.com>
Date: Tue, 26 May 2026 10:41:40 -0700
Subject: [PATCH 3/8] Cortex-M backend: address PR #19676 review feedback

Two small follow-ups on the linear ISA-dispatch fix:

* Make `ConvertToCortexMPass.target_config` a required constructor arg
  rather than silently defaulting to M55/MVE on None. The dispatch now
  picks a CMSIS-NN convention from the target ISA, so an implicit
  default would emit a wrong PTE for non-MVE targets. The pass manager
  threads its own config through, so production callers are unaffected;
  direct instantiators now get a clear TypeError at construction.

* Fold the small-magnitude regression test into the regular linear test
  file. The six ISA-by-bias variants are now parametrized cases on top
  of the existing `McuTestCase` + `ops_after_transforms` harness. The
  bespoke post-pass graph walk collapses to one inline assertion at the
  end of the dialect test that the slot consumed by the configured ISA
  is populated and the other is None -- enough to catch a future
  regression that emits a zero-valued kernel_sum on a no-bias DSP path
  while staying numerically inert.
---
 .../passes/convert_to_cortex_m_pass.py        |   8 +-
 .../test_quantized_linear_small_magnitude.py  | 185 ------------------
 backends/cortex_m/test/ops/test_linear.py     | 153 ++++++++++++++-
 3 files changed, 155 insertions(+), 191 deletions(-)
 delete mode 100644 backends/cortex_m/test/misc/test_quantized_linear_small_magnitude.py

diff --git a/backends/cortex_m/passes/convert_to_cortex_m_pass.py b/backends/cortex_m/passes/convert_to_cortex_m_pass.py
index 79654e59f4e..6a9425e1282 100644
--- a/backends/cortex_m/passes/convert_to_cortex_m_pass.py
+++ b/backends/cortex_m/passes/convert_to_cortex_m_pass.py
@@ -11,7 +11,7 @@
 import torch.fx
 from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
 from executorch.backends.cortex_m.passes.passes_utils import quantize_multiplier_aot
-from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig
+from executorch.backends.cortex_m.target_config import CortexMTargetConfig
 
 from executorch.backends.transforms.utils import (
     create_constant_placeholder,
@@ -38,12 +38,10 @@ class ConvertToCortexMPass(XNNPACKPass):
     def __init__(
         self,
         exported_program: ExportedProgram,
-        target_config: CortexMTargetConfig | None = None,
+        target_config: CortexMTargetConfig,
     ) -> None:
         super().__init__(exported_program)
-        # Default mirrors CortexMPassManager: MVE-capable M55 (the previous
-        # behavior for any caller that constructs the pass without a config).
-        self._target_config = target_config or CortexMTargetConfig(cpu=CortexM.M55)
+        self._target_config = target_config
 
     @property
     def target_config(self) -> CortexMTargetConfig:
diff --git a/backends/cortex_m/test/misc/test_quantized_linear_small_magnitude.py b/backends/cortex_m/test/misc/test_quantized_linear_small_magnitude.py
deleted file mode 100644
index 5cfcea74e70..00000000000
--- a/backends/cortex_m/test/misc/test_quantized_linear_small_magnitude.py
+++ /dev/null
@@ -1,185 +0,0 @@
-# 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.
-"""Regression test for the cortex_m::quantized_linear bias/kernel_sum bug.
-
-CMSIS-NN's `arm_fully_connected_s8` has three runtime paths, gated by
-compile-time ARM_MATH_MVEI / ARM_MATH_DSP:
-
-* MVE: reads ctx.buf (precomputed kernel_sum that includes bias plus
-  input_offset x sum(weight)), ignores the bias argument.
-* DSP / scalar: read the bias argument directly, ignore ctx.buf.
-
-`ConvertToCortexMPass._get_linear_replacement` selects which input format
-to emit based on `CortexMTargetConfig.backend`. Before the fix, the pass
-unconditionally emitted kernel_sum + None-bias, which on a non-MVE build
-silently dropped both the bias and the input-offset term. The bug only
-showed up when those terms dominated the int32 accumulator, i.e. on
-small-magnitude inputs.
-
-Coverage:
-
-* `test_dialect_small_magnitude_linear` runs each ISA through the Python
-  op impl and checks that bias=True and bias=False variants both round
-  to the same int8 outputs as the float reference.
-* `test_aot_graph_shape_small_magnitude_linear` inspects the post-pass
-  graph and asserts the bias/kernel_sum arg positions match the ISA
-  convention -- this is the direct regression check.
-* `test_implementation_small_magnitude_linear` runs the bias=True case
-  through the default (M55, MVE) build path so the impl test exercises
-  the kernel_sum codepath in simulation.
-"""
-
-from dataclasses import dataclass
-
-import torch
-import torch.nn as nn
-from executorch.backends.arm.test.common import parametrize
-from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig
-from executorch.backends.cortex_m.test.tester import CortexMTester, McuTestCase
-from executorch.backends.test.harness.stages import StageType
-from executorch.exir.dialects._ops import ops as exir_ops
-
-torch.manual_seed(0)
-
-
-class _SmallMagnitudeLinear(nn.Module):
-    ops_before_transforms = {
-        "executorch_exir_dialects_edge__ops_aten_linear_default": 1,
-        "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 2,
-        "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 4,
-    }
-    ops_after_transforms = {
-        "executorch_exir_dialects_edge__ops_cortex_m_quantized_linear_default": 1,
-        "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1,
-        "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1,
-    }
-
-    def __init__(self, bias: bool = True):
-        super().__init__()
-        self.fc = nn.Linear(512, 10, bias=bias)
-
-    def forward(self, x):
-        return self.fc(x)
-
-
-class _SmallMagnitudeLinearNoBias(_SmallMagnitudeLinear):
-    ops_before_transforms = {
-        "executorch_exir_dialects_edge__ops_aten_linear_default": 1,
-        "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 2,
-        "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 3,
-    }
-
-    def __init__(self):
-        super().__init__(bias=False)
-
-
-def _make_input():
-    return torch.rand(1, 512) * 0.002
-
-
-_calibration_samples = [(_make_input(),) for _ in range(8)]
-
-
-@dataclass(frozen=True)
-class _IsaVariant:
-    case: McuTestCase
-    target_config: CortexMTargetConfig
-    uses_kernel_sum: bool
-
-
-def _variant(model_cls, cpu: CortexM, uses_kernel_sum: bool) -> _IsaVariant:
-    return _IsaVariant(
-        case=McuTestCase(
-            model=model_cls().eval(),
-            example_inputs=lambda: (_make_input(),),
-        ),
-        target_config=CortexMTargetConfig(cpu=cpu),
-        uses_kernel_sum=uses_kernel_sum,
-    )
-
-
-# bias=True covers the regression directly (the bug dropped the bias
-# term); bias=False covers the symmetric case where only the
-# input-offset term is missing on the non-MVE paths.
-test_variants = {
-    "mve_bias": _variant(_SmallMagnitudeLinear, CortexM.M55, uses_kernel_sum=True),
-    "dsp_bias": _variant(_SmallMagnitudeLinear, CortexM.M4, uses_kernel_sum=False),
-    "scalar_bias": _variant(
-        _SmallMagnitudeLinear, CortexM.M0PLUS, uses_kernel_sum=False
-    ),
-    "mve_nobias": _variant(
-        _SmallMagnitudeLinearNoBias, CortexM.M55, uses_kernel_sum=True
-    ),
-    "dsp_nobias": _variant(
-        _SmallMagnitudeLinearNoBias, CortexM.M4, uses_kernel_sum=False
-    ),
-    "scalar_nobias": _variant(
-        _SmallMagnitudeLinearNoBias, CortexM.M0PLUS, uses_kernel_sum=False
-    ),
-}
-
-
-@parametrize("variant", test_variants)
-def test_dialect_small_magnitude_linear(variant: _IsaVariant):
-    tester = CortexMTester(
-        variant.case.model,
-        variant.case.get_example_inputs(),
-        target_config=variant.target_config,
-    )
-    tester.test_dialect(
-        ops_before_transforms=variant.case.model.ops_before_transforms,
-        ops_after_transforms=variant.case.model.ops_after_transforms,
-        qtol=1,
-        calibration_samples=_calibration_samples,
-    )
-
-
-@parametrize("variant", test_variants)
-def test_aot_graph_shape_small_magnitude_linear(variant: _IsaVariant):
-    """Assert the post-pass node args match the ISA's CMSIS-NN convention."""
-    tester = CortexMTester(
-        variant.case.model,
-        variant.case.get_example_inputs(),
-        target_config=variant.target_config,
-    )
-    tester.quantize(None)
-    tester.export()
-    tester.to_edge()
-    tester.run_passes()
-
-    module = tester.get_artifact(StageType.RUN_PASSES).exported_program().module()
-    linear_target = exir_ops.edge.cortex_m.quantized_linear.default
-    linear_nodes = [
-        n
-        for n in module.graph.nodes
-        if n.op == "call_function" and n.target == linear_target
-    ]
-    assert (
-        len(linear_nodes) == 1
-    ), f"expected one quantized_linear node, got {len(linear_nodes)}"
-    bias_arg, kernel_sum_arg = linear_nodes[0].args[2], linear_nodes[0].args[3]
-
-    if variant.uses_kernel_sum:
-        assert bias_arg is None, "MVE path must not pass bias (CMSIS-NN ignores it)"
-        assert kernel_sum_arg is not None, "MVE path requires precomputed kernel_sum"
-    else:
-        assert kernel_sum_arg is None, "non-MVE path must not pass kernel_sum"
-        # bias is allowed to be None only if the source nn.Linear had bias=False.
-        expects_bias = variant.case.model.fc.bias is not None
-        if expects_bias:
-            assert (
-                bias_arg is not None
-            ), "non-MVE path with bias must forward bias to CMSIS-NN"
-
-
-def test_implementation_small_magnitude_linear():
-    """Exercise the MVE kernel_sum codepath via the default M55 simulator build."""
-    case = McuTestCase(
-        model=_SmallMagnitudeLinear().eval(),
-        example_inputs=lambda: (_make_input(),),
-    )
-    tester = CortexMTester(case.model, case.get_example_inputs())
-    tester.test_implementation(qtol=1, calibration_samples=_calibration_samples)
diff --git a/backends/cortex_m/test/ops/test_linear.py b/backends/cortex_m/test/ops/test_linear.py
index e81daa7e83e..fc23bbbd5c6 100644
--- a/backends/cortex_m/test/ops/test_linear.py
+++ b/backends/cortex_m/test/ops/test_linear.py
@@ -1,16 +1,21 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
 
+from dataclasses import dataclass
+
 import torch
 from executorch.backends.arm.test.common import parametrize
+from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig
 from executorch.backends.cortex_m.test.tester import (
     CortexMTester,
     McuTestCase,
     ramp_tensor,
 )
+from executorch.backends.test.harness.stages import StageType
+from executorch.exir.dialects._ops import ops as exir_ops
 
 
 class CortexMLinear(torch.nn.Module):
@@ -128,3 +133,149 @@ def test_dialect_linear(test_case):
 def test_implementation_linear(test_case):
     tester = CortexMTester(test_case.model, test_case.example_inputs)
     tester.test_implementation(qtol=1)
+
+
+# ---------------------------------------------------------------------------
+# Regression: cortex_m::quantized_linear must pick the right CMSIS-NN input
+# convention based on the target ISA. `arm_fully_connected_s8` reads
+# kernel_sum (ctx.buf) on MVE/Helium and reads the bias argument on DSP/scalar
+# paths; the two are mutually exclusive. Previously the pass unconditionally
+# emitted the MVE shape, which silently dropped the bias and input-offset
+# terms on every non-MVE build. The regression only showed up when those
+# terms dominated the int32 accumulator -- i.e., on small-magnitude inputs.
+#
+# Coverage strategy: a single ISA-parametrized dialect test verifies the
+# numeric output against the float reference (catches the dropped-bias bug
+# directly), checks ops_after_transforms to confirm the linear lowered, and
+# asserts the post-pass node has the value in the slot the configured ISA
+# expects -- the structural guard against a regression that emits zero-valued
+# kernel_sum on a no-bias DSP path (numerically inert, but wrong shape).
+# An additional implementation test drives the default M55 MVE build path
+# through the simulator.
+# ---------------------------------------------------------------------------
+
+
+class _SmallMagnitudeLinear(torch.nn.Module):
+    ops_before_transforms = {
+        "executorch_exir_dialects_edge__ops_aten_linear_default": 1,
+        "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 2,
+        "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 4,
+    }
+    ops_after_transforms = {
+        "executorch_exir_dialects_edge__ops_cortex_m_quantized_linear_default": 1,
+        "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1,
+        "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1,
+    }
+
+    def __init__(self, bias: bool = True):
+        super().__init__()
+        self.fc = torch.nn.Linear(512, 10, bias=bias)
+
+    def forward(self, x):
+        return self.fc(x)
+
+
+class _SmallMagnitudeLinearNoBias(_SmallMagnitudeLinear):
+    ops_before_transforms = {
+        "executorch_exir_dialects_edge__ops_aten_linear_default": 1,
+        "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 2,
+        "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 3,
+    }
+
+    def __init__(self):
+        super().__init__(bias=False)
+
+
+def _small_magnitude_input():
+    return torch.rand(1, 512) * 0.002
+
+
+_small_magnitude_calibration = [(_small_magnitude_input(),) for _ in range(8)]
+
+
+@dataclass(frozen=True)
+class _SmallMagnitudeVariant:
+    case: McuTestCase
+    target_config: CortexMTargetConfig
+    uses_kernel_sum: bool
+
+
+def _small_magnitude_variant(
+    model_cls, cpu: CortexM, uses_kernel_sum: bool
+) -> _SmallMagnitudeVariant:
+    return _SmallMagnitudeVariant(
+        case=McuTestCase(
+            model=model_cls().eval(),
+            example_inputs=lambda: (_small_magnitude_input(),),
+        ),
+        target_config=CortexMTargetConfig(cpu=cpu),
+        uses_kernel_sum=uses_kernel_sum,
+    )
+
+
+# bias=True covers the regression directly (the bug dropped the bias term);
+# bias=False covers the symmetric case where only the input-offset term is
+# missing on the non-MVE paths.
+small_magnitude_variants = {
+    "mve_bias": _small_magnitude_variant(
+        _SmallMagnitudeLinear, CortexM.M55, uses_kernel_sum=True
+    ),
+    "dsp_bias": _small_magnitude_variant(
+        _SmallMagnitudeLinear, CortexM.M4, uses_kernel_sum=False
+    ),
+    "scalar_bias": _small_magnitude_variant(
+        _SmallMagnitudeLinear, CortexM.M0PLUS, uses_kernel_sum=False
+    ),
+    "mve_nobias": _small_magnitude_variant(
+        _SmallMagnitudeLinearNoBias, CortexM.M55, uses_kernel_sum=True
+    ),
+    "dsp_nobias": _small_magnitude_variant(
+        _SmallMagnitudeLinearNoBias, CortexM.M4, uses_kernel_sum=False
+    ),
+    "scalar_nobias": _small_magnitude_variant(
+        _SmallMagnitudeLinearNoBias, CortexM.M0PLUS, uses_kernel_sum=False
+    ),
+}
+
+
+@parametrize("variant", small_magnitude_variants)
+def test_dialect_linear_small_magnitude(variant: _SmallMagnitudeVariant):
+    tester = CortexMTester(
+        variant.case.model,
+        variant.case.get_example_inputs(),
+        target_config=variant.target_config,
+    )
+    tester.test_dialect(
+        ops_before_transforms=variant.case.model.ops_before_transforms,
+        ops_after_transforms=variant.case.model.ops_after_transforms,
+        qtol=1,
+        calibration_samples=_small_magnitude_calibration,
+    )
+
+    # Structural guard: numeric divergence catches the original dropped-bias
+    # bug, but a future regression that emits zero-valued kernel_sum on a
+    # no-bias DSP/scalar path would be numerically inert. Assert the slot the
+    # configured ISA actually consumes is populated and the unused one is None.
+    module = tester.get_artifact(StageType.RUN_PASSES).exported_program().module()
+    linear_target = exir_ops.edge.cortex_m.quantized_linear.default
+    [linear_node] = [
+        n
+        for n in module.graph.nodes
+        if n.op == "call_function" and n.target == linear_target
+    ]
+    bias_arg, kernel_sum_arg = linear_node.args[2], linear_node.args[3]
+    if variant.uses_kernel_sum:
+        assert kernel_sum_arg is not None
+        assert bias_arg is None
+    else:
+        assert kernel_sum_arg is None
+
+
+def test_implementation_linear_small_magnitude():
+    """Exercise the MVE kernel_sum codepath via the default M55 simulator build."""
+    case = McuTestCase(
+        model=_SmallMagnitudeLinear().eval(),
+        example_inputs=lambda: (_small_magnitude_input(),),
+    )
+    tester = CortexMTester(case.model, case.get_example_inputs())
+    tester.test_implementation(qtol=1, calibration_samples=_small_magnitude_calibration)

From 527ca222c59c7efcd7cf5f860326a25fc592428a Mon Sep 17 00:00:00 2001
From: RJ Ascani <rja@meta.com>
Date: Tue, 9 Jun 2026 08:46:49 -0700
Subject: [PATCH 4/8] Assert bias_arg is None for no-bias models on DSP/scalar
 paths

Tighten the structural guard in the small-magnitude regression test so
the non-MVE else branch also verifies that bias_arg is None when the
model has no bias. Previously only kernel_sum_arg was checked.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 backends/cortex_m/test/ops/test_linear.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/backends/cortex_m/test/ops/test_linear.py b/backends/cortex_m/test/ops/test_linear.py
index fc23bbbd5c6..adc622479c0 100644
--- a/backends/cortex_m/test/ops/test_linear.py
+++ b/backends/cortex_m/test/ops/test_linear.py
@@ -269,6 +269,8 @@ def test_dialect_linear_small_magnitude(variant: _SmallMagnitudeVariant):
         assert bias_arg is None
     else:
         assert kernel_sum_arg is None
+        if variant.case.model.fc.bias is None:
+            assert bias_arg is None
 
 
 def test_implementation_linear_small_magnitude():

From cb351e994ad45b400fdddd09ad5e9e0aff1ad6a1 Mon Sep 17 00:00:00 2001
From: RJ Ascani <rja@meta.com>
Date: Tue, 9 Jun 2026 09:09:51 -0700
Subject: [PATCH 5/8] Fix mypy lints

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 backends/cortex_m/passes/aten_to_cortex_m_pass.py | 2 +-
 backends/cortex_m/test/ops/test_linear.py         | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/backends/cortex_m/passes/aten_to_cortex_m_pass.py b/backends/cortex_m/passes/aten_to_cortex_m_pass.py
index 136fc285e51..c33fd499981 100644
--- a/backends/cortex_m/passes/aten_to_cortex_m_pass.py
+++ b/backends/cortex_m/passes/aten_to_cortex_m_pass.py
@@ -53,7 +53,7 @@ def __init__(
     ) -> None:
         super().__init__(exported_program=exported_program)
         self.target_config = target_config
-        self._DIALECT_SUBSTITUTIONS = {
+        self._DIALECT_SUBSTITUTIONS = {  # type: ignore[misc]
             **type(self)._DIALECT_SUBSTITUTIONS,
             exir_ops.edge.aten.linear.default: lambda node, ep: _get_linear_replacement(
                 node, ep, target_config
diff --git a/backends/cortex_m/test/ops/test_linear.py b/backends/cortex_m/test/ops/test_linear.py
index adc622479c0..a997d11e4ea 100644
--- a/backends/cortex_m/test/ops/test_linear.py
+++ b/backends/cortex_m/test/ops/test_linear.py
@@ -269,6 +269,7 @@ def test_dialect_linear_small_magnitude(variant: _SmallMagnitudeVariant):
         assert bias_arg is None
     else:
         assert kernel_sum_arg is None
+        assert isinstance(variant.case.model, _SmallMagnitudeLinear)
         if variant.case.model.fc.bias is None:
             assert bias_arg is None
 

From bf3c8b4cc764e11483a18970420b4ec46621312f Mon Sep 17 00:00:00 2001
From: RJ Ascani <rja@meta.com>
Date: Tue, 9 Jun 2026 10:11:03 -0700
Subject: [PATCH 6/8] Pass dialect_pass instance to substitution functions

Extend AtenToDialectPass to pass `self` instead of `self.exported_program`
to registered substitution functions. This lets substitution functions
access subclass-specific state (e.g. target_config on AtenToCortexMPass)
without workarounds.

Replaces the instance-level _DIALECT_SUBSTITUTIONS shadow with a standard
@register_dialect_substitution decorator on _get_linear_replacement.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 .../cortex_m/passes/aten_to_cortex_m_pass.py  | 29 ++++++++++---------
 backends/cortex_m/test/ops/test_linear.py     | 21 ++++++++------
 backends/transforms/aten_to_dialect_pass.py   |  4 +--
 .../test/test_aten_to_dialect_pass.py         | 22 +++++---------
 4 files changed, 37 insertions(+), 39 deletions(-)

diff --git a/backends/cortex_m/passes/aten_to_cortex_m_pass.py b/backends/cortex_m/passes/aten_to_cortex_m_pass.py
index c33fd499981..32d06578b02 100644
--- a/backends/cortex_m/passes/aten_to_cortex_m_pass.py
+++ b/backends/cortex_m/passes/aten_to_cortex_m_pass.py
@@ -53,12 +53,6 @@ def __init__(
     ) -> None:
         super().__init__(exported_program=exported_program)
         self.target_config = target_config
-        self._DIALECT_SUBSTITUTIONS = {  # type: ignore[misc]
-            **type(self)._DIALECT_SUBSTITUTIONS,
-            exir_ops.edge.aten.linear.default: lambda node, ep: _get_linear_replacement(
-                node, ep, target_config
-            ),
-        }
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         result = super().call(graph_module)
@@ -153,7 +147,7 @@ def _has_qparams(node: Node) -> bool:
 @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.tanh.default)
 @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.silu.default)
 def _get_activation_replacement(
-    node: Node, exported_program: ExportedProgram
+    node: Node, dialect_pass: AtenToDialectPass
 ) -> DialectNodeSpec | None:
     """Lower a standalone quantized sigmoid / tanh / silu to a single
     cortex_m.quantized_activation call backed by an AoT-built 256-entry
@@ -163,6 +157,7 @@ def _get_activation_replacement(
     if not _has_qparams(node):
         return None
 
+    exported_program = dialect_pass.exported_program
     input_qparams = node.meta["input_qparams"][0]
     output_qparams = node.meta["output_qparams"][0]
     lut_tensor = build_activation_lut(
@@ -192,10 +187,9 @@ def _get_activation_replacement(
     )
 
 
+@AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.linear.default)
 def _get_linear_replacement(
-    node: Node,
-    exported_program: ExportedProgram,
-    target_config: CortexMTargetConfig,
+    node: Node, dialect_pass: AtenToDialectPass
 ) -> DialectNodeSpec | None:
     """
     Let
@@ -217,6 +211,10 @@ def _get_linear_replacement(
     if not _has_qparams(node):
         return None
 
+    assert isinstance(dialect_pass, AtenToCortexMPass)
+    exported_program = dialect_pass.exported_program
+    target_config = dialect_pass.target_config
+
     input_scale = node.meta["input_qparams"][0].scale
     input_zp = node.meta["input_qparams"][0].zp
     weight_scale = node.meta["input_qparams"][1].scale
@@ -286,11 +284,12 @@ def _get_linear_replacement(
 
 @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.convolution.default)
 def _get_convolution_replacement(
-    node: Node, exported_program: ExportedProgram
+    node: Node, dialect_pass: AtenToDialectPass
 ) -> DialectNodeSpec | None:
     if not _has_qparams(node):
         return None
 
+    exported_program = dialect_pass.exported_program
     conv_args = node.args
     (
         x,
@@ -315,7 +314,7 @@ def _get_convolution_replacement(
     )
 
     if transposed:
-        return _get_transpose_conv2d_replacement(node, exported_program)
+        return _get_transpose_conv2d_replacement(node, dialect_pass)
 
     input_scale = node.meta["input_qparams"][0].scale
     input_zero_point = node.meta["input_qparams"][0].zp
@@ -460,7 +459,7 @@ def _get_convolution_replacement(
 
 
 def _get_transpose_conv2d_replacement(
-    node: Node, exported_program: ExportedProgram
+    node: Node, dialect_pass: AtenToDialectPass
 ) -> DialectNodeSpec | None:
     """
     Transform aten.convolution with transposed=True to cortex_m.quantized_transpose_conv2d.
@@ -468,6 +467,7 @@ def _get_transpose_conv2d_replacement(
     if not _has_qparams(node):
         return None
 
+    exported_program = dialect_pass.exported_program
     conv_t_args = node.args
     (
         x,
@@ -585,11 +585,12 @@ def _get_transpose_conv2d_replacement(
 
 @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.bmm.default)
 def _get_bmm_replacement(
-    node: Node, exported_program: ExportedProgram
+    node: Node, dialect_pass: AtenToDialectPass
 ) -> DialectNodeSpec | None:
     if not _has_qparams(node):
         return None
 
+    exported_program = dialect_pass.exported_program
     lhs_scale = node.meta["input_qparams"][0].scale
     lhs_zp = node.meta["input_qparams"][0].zp
     rhs_scale = node.meta["input_qparams"][1].scale
diff --git a/backends/cortex_m/test/ops/test_linear.py b/backends/cortex_m/test/ops/test_linear.py
index a997d11e4ea..ed03b535917 100644
--- a/backends/cortex_m/test/ops/test_linear.py
+++ b/backends/cortex_m/test/ops/test_linear.py
@@ -198,10 +198,11 @@ class _SmallMagnitudeVariant:
     case: McuTestCase
     target_config: CortexMTargetConfig
     uses_kernel_sum: bool
+    has_bias: bool
 
 
 def _small_magnitude_variant(
-    model_cls, cpu: CortexM, uses_kernel_sum: bool
+    model_cls, cpu: CortexM, *, uses_kernel_sum: bool, has_bias: bool
 ) -> _SmallMagnitudeVariant:
     return _SmallMagnitudeVariant(
         case=McuTestCase(
@@ -210,6 +211,7 @@ def _small_magnitude_variant(
         ),
         target_config=CortexMTargetConfig(cpu=cpu),
         uses_kernel_sum=uses_kernel_sum,
+        has_bias=has_bias,
     )
 
 
@@ -218,22 +220,22 @@ def _small_magnitude_variant(
 # missing on the non-MVE paths.
 small_magnitude_variants = {
     "mve_bias": _small_magnitude_variant(
-        _SmallMagnitudeLinear, CortexM.M55, uses_kernel_sum=True
+        _SmallMagnitudeLinear, CortexM.M55, uses_kernel_sum=True, has_bias=True
     ),
     "dsp_bias": _small_magnitude_variant(
-        _SmallMagnitudeLinear, CortexM.M4, uses_kernel_sum=False
+        _SmallMagnitudeLinear, CortexM.M4, uses_kernel_sum=False, has_bias=True
     ),
     "scalar_bias": _small_magnitude_variant(
-        _SmallMagnitudeLinear, CortexM.M0PLUS, uses_kernel_sum=False
+        _SmallMagnitudeLinear, CortexM.M0PLUS, uses_kernel_sum=False, has_bias=True
     ),
     "mve_nobias": _small_magnitude_variant(
-        _SmallMagnitudeLinearNoBias, CortexM.M55, uses_kernel_sum=True
+        _SmallMagnitudeLinearNoBias, CortexM.M55, uses_kernel_sum=True, has_bias=False
     ),
     "dsp_nobias": _small_magnitude_variant(
-        _SmallMagnitudeLinearNoBias, CortexM.M4, uses_kernel_sum=False
+        _SmallMagnitudeLinearNoBias, CortexM.M4, uses_kernel_sum=False, has_bias=False
     ),
     "scalar_nobias": _small_magnitude_variant(
-        _SmallMagnitudeLinearNoBias, CortexM.M0PLUS, uses_kernel_sum=False
+        _SmallMagnitudeLinearNoBias, CortexM.M0PLUS, uses_kernel_sum=False, has_bias=False
     ),
 }
 
@@ -269,8 +271,9 @@ def test_dialect_linear_small_magnitude(variant: _SmallMagnitudeVariant):
         assert bias_arg is None
     else:
         assert kernel_sum_arg is None
-        assert isinstance(variant.case.model, _SmallMagnitudeLinear)
-        if variant.case.model.fc.bias is None:
+        if variant.has_bias:
+            assert bias_arg is not None
+        else:
             assert bias_arg is None
 
 
diff --git a/backends/transforms/aten_to_dialect_pass.py b/backends/transforms/aten_to_dialect_pass.py
index e44b71c96dc..c9380db8e75 100644
--- a/backends/transforms/aten_to_dialect_pass.py
+++ b/backends/transforms/aten_to_dialect_pass.py
@@ -28,7 +28,7 @@ class DialectNodeSpec:
 
 # Expected type to be used for substitution functions
 SubstitutionFn: TypeAlias = Callable[
-    [torch.fx.Node, torch.export.ExportedProgram], DialectNodeSpec | None
+    [torch.fx.Node, "AtenToDialectPass"], DialectNodeSpec | None
 ]
 
 
@@ -86,7 +86,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
             if substitution_func is None:
                 continue
 
-            dialect_node_spec = substitution_func(node, self.exported_program)
+            dialect_node_spec = substitution_func(node, self)
             if dialect_node_spec is None:
                 continue
 
diff --git a/backends/transforms/test/test_aten_to_dialect_pass.py b/backends/transforms/test/test_aten_to_dialect_pass.py
index 885d1c70392..f328169ab2e 100644
--- a/backends/transforms/test/test_aten_to_dialect_pass.py
+++ b/backends/transforms/test/test_aten_to_dialect_pass.py
@@ -61,9 +61,8 @@ class _TestAtenToDialectPass(AtenToDialectPass):
 
     @_TestAtenToDialectPass.register_dialect_substitution(torch.ops.aten.add.Tensor)
     def replace_add_with_sub(
-        node: Node, exported_program: ExportedProgram
+        node: Node, dialect_pass: AtenToDialectPass
     ) -> DialectNodeSpec | None:
-        del exported_program
         return DialectNodeSpec(torch.ops.aten.sub.Tensor, node.args)
 
     exported_program = _export_add_model()
@@ -82,7 +81,7 @@ class _TestAtenToDialectPass(AtenToDialectPass):
 
     @_TestAtenToDialectPass.register_dialect_substitution(torch.ops.aten.add.Tensor)
     def replace_add_rhs_with_constant(
-        node: Node, exported_program: ExportedProgram
+        node: Node, dialect_pass: AtenToDialectPass
     ) -> DialectNodeSpec | None:
         first_placeholder = next(
             graph_node
@@ -91,7 +90,7 @@ def replace_add_rhs_with_constant(
         )
         with node.graph.inserting_before(first_placeholder):
             const_node = create_constant_placeholder(
-                exp_program=exported_program,
+                exp_program=dialect_pass.exported_program,
                 graph=node.graph,
                 name="test_constant",
                 kind=InputKind.PARAMETER,
@@ -125,9 +124,8 @@ class _TestAtenToDialectPass(AtenToDialectPass):
 
     @_TestAtenToDialectPass.register_dialect_substitution(torch.ops.aten.add.Tensor)
     def replace_add_alpha(
-        node: Node, exported_program: ExportedProgram
+        node: Node, dialect_pass: AtenToDialectPass
     ) -> DialectNodeSpec | None:
-        del exported_program
         return DialectNodeSpec(torch.ops.aten.add.Tensor, node.args, {"alpha": 3})
 
     exported_program = _export_add_alpha_model()
@@ -150,9 +148,8 @@ class _TestAtenToDialectPass(AtenToDialectPass):
 
     @_TestAtenToDialectPass.register_dialect_substitution(torch.ops.aten.add.Tensor)
     def replace_add_with_sub(
-        node: Node, exported_program: ExportedProgram
+        node: Node, dialect_pass: AtenToDialectPass
     ) -> DialectNodeSpec | None:
-        del exported_program
         return DialectNodeSpec(torch.ops.aten.sub.Tensor, node.args)
 
     exported_program = _export_add_model()
@@ -178,9 +175,8 @@ class _TestAtenToDialectPass(AtenToDialectPass):
 
     @_TestAtenToDialectPass.register_dialect_substitution(torch.ops.aten.add.Tensor)
     def do_not_replace(
-        node: Node, exported_program: ExportedProgram
+        node: Node, dialect_pass: AtenToDialectPass
     ) -> DialectNodeSpec | None:
-        del node, exported_program
         return None
 
     exported_program = _export_add_model()
@@ -199,16 +195,14 @@ class _TestAtenToDialectPass(AtenToDialectPass):
 
     @_TestAtenToDialectPass.register_dialect_substitution(torch.ops.aten.add.Tensor)
     def first_replace(
-        node: Node, exported_program: ExportedProgram
+        node: Node, dialect_pass: AtenToDialectPass
     ) -> DialectNodeSpec | None:
-        del exported_program
         return DialectNodeSpec(torch.ops.aten.sub.Tensor, node.args)
 
     with pytest.raises(RuntimeError, match="Multiple substitutions registered"):
 
         @_TestAtenToDialectPass.register_dialect_substitution(torch.ops.aten.add.Tensor)
         def second_replace(
-            node: Node, exported_program: ExportedProgram
+            node: Node, dialect_pass: AtenToDialectPass
         ) -> DialectNodeSpec | None:
-            del exported_program
             return DialectNodeSpec(torch.ops.aten.mul.Tensor, node.args)

From d07a763b80e6b24560aaeafcdb2a565b48a1a63b Mon Sep 17 00:00:00 2001
From: RJ Ascani <rja@meta.com>
Date: Tue, 9 Jun 2026 12:05:12 -0700
Subject: [PATCH 7/8] Use __future__ annotations for SubstitutionFn forward
 reference

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 backends/cortex_m/test/ops/test_linear.py   | 5 ++++-
 backends/transforms/aten_to_dialect_pass.py | 3 ++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/backends/cortex_m/test/ops/test_linear.py b/backends/cortex_m/test/ops/test_linear.py
index ed03b535917..37a02edc35f 100644
--- a/backends/cortex_m/test/ops/test_linear.py
+++ b/backends/cortex_m/test/ops/test_linear.py
@@ -235,7 +235,10 @@ def _small_magnitude_variant(
         _SmallMagnitudeLinearNoBias, CortexM.M4, uses_kernel_sum=False, has_bias=False
     ),
     "scalar_nobias": _small_magnitude_variant(
-        _SmallMagnitudeLinearNoBias, CortexM.M0PLUS, uses_kernel_sum=False, has_bias=False
+        _SmallMagnitudeLinearNoBias,
+        CortexM.M0PLUS,
+        uses_kernel_sum=False,
+        has_bias=False,
     ),
 }
 
diff --git a/backends/transforms/aten_to_dialect_pass.py b/backends/transforms/aten_to_dialect_pass.py
index c9380db8e75..29bc84b0758 100644
--- a/backends/transforms/aten_to_dialect_pass.py
+++ b/backends/transforms/aten_to_dialect_pass.py
@@ -3,6 +3,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from __future__ import annotations
 
 import traceback
 from collections.abc import Callable
@@ -28,7 +29,7 @@ class DialectNodeSpec:
 
 # Expected type to be used for substitution functions
 SubstitutionFn: TypeAlias = Callable[
-    [torch.fx.Node, "AtenToDialectPass"], DialectNodeSpec | None
+    [torch.fx.Node, AtenToDialectPass], DialectNodeSpec | None
 ]
 
 

From 1f702270cfed0b51fd78ebf4bcaf32706a1c7497 Mon Sep 17 00:00:00 2001
From: RJ Ascani <rja@meta.com>
Date: Tue, 9 Jun 2026 13:50:15 -0700
Subject: [PATCH 8/8] Move SubstitutionFn after AtenToDialectPass class
 definition

The TypeAlias RHS is a runtime expression, not an annotation, so
__future__.annotations does not defer it. Move it after the class
so AtenToDialectPass is defined when the Callable type evaluates.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 backends/transforms/aten_to_dialect_pass.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/backends/transforms/aten_to_dialect_pass.py b/backends/transforms/aten_to_dialect_pass.py
index 29bc84b0758..f26541a4b0f 100644
--- a/backends/transforms/aten_to_dialect_pass.py
+++ b/backends/transforms/aten_to_dialect_pass.py
@@ -27,12 +27,6 @@ class DialectNodeSpec:
     kwargs: dict = None
 
 
-# Expected type to be used for substitution functions
-SubstitutionFn: TypeAlias = Callable[
-    [torch.fx.Node, AtenToDialectPass], DialectNodeSpec | None
-]
-
-
 class AtenToDialectPass(ExportPass):
     """
     General pass to convert ops from ATen to a specific dialect.
@@ -117,3 +111,11 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
             graph_module = super().call(graph_module).graph_module
 
         return PassResult(graph_module, modified)
+
+
+# Defined after the class so AtenToDialectPass is available at runtime.
+# Class-body references to SubstitutionFn are annotation-only and resolve
+# via __future__.annotations.
+SubstitutionFn: TypeAlias = Callable[
+    [torch.fx.Node, AtenToDialectPass], DialectNodeSpec | None
+]