Cortex-M: make quantized add and multiply layout-agnostic

backends/cortex_m/ops/cortex_m_ops_common.h

@@ -17,6 +17,7 @@
 #include <executorch/runtime/core/exec_aten/util/dim_order_util.h>
 #include <executorch/runtime/platform/assert.h>
 
+#include <algorithm>
 #include <cinttypes>
 #include <limits>
 #include <optional>
@@ -149,23 +150,30 @@ inline bool is_channels_last_tensor(const Tensor& tensor) {
   return tensor.dim_order() == channels_last_order;
 }
 
+// A channel broadcast the elementwise kernels can serve as a flat repeat: one
+// operand holds a single value per channel, and channels are contiguous under
+// both activation contracts, so the broadcast operand's element count is the
+// repeat length. Deliberately layout-free -- deriving the channel axis from the
+// dim order would be unsound, because a serialized dim order does not identify
+// it when the channel count is one.
 inline bool is_channel_broadcast(const Tensor& tensor1, const Tensor& tensor2) {
-  if (tensor1.dim() != tensor2.dim()) {
+  if (tensor1.dim() != tensor2.dim() || tensor1.dim() != 4) {
     return false;
   }
-
-  if (tensor1.dim() != 4) {
+  if (tensor1.numel() == tensor2.numel()) {
     return false;
   }
 
-  if (tensor1.size(1) != tensor2.size(1)) {
-    return false;
+  const bool first_is_larger = tensor1.numel() > tensor2.numel();
+  const Tensor& larger = first_is_larger ? tensor1 : tensor2;
+  const Tensor& smaller = first_is_larger ? tensor2 : tensor1;
+  int64_t non_unit = 0;
+  for (int64_t i = 0; i < smaller.dim(); ++i) {
+    if (smaller.size(i) != 1) {
+      ++non_unit;
+    }
   }
-
-  const bool tensor1_channels_only = tensor1.numel() == tensor1.size(1);
-  const bool tensor2_channels_only = tensor2.numel() == tensor2.size(1);
-
-  return tensor1_channels_only || tensor2_channels_only;
+  return non_unit <= 1 && larger.numel() % smaller.numel() == 0;
 }
 
 inline bool check_int32_within_range(

backends/cortex_m/ops/op_quantized_add.cpp

@@ -30,15 +30,21 @@ Tensor& quantized_add_out(
     const int64_t activation_min,
     const int64_t activation_max,
     Tensor& out) {
-  // Validate tensor types and dim order
   bool channel_broadcast = is_channel_broadcast(input1_int8, input2_int8);
   validate_cmsis_nn_tensor_requirements(
       input1_int8,
       input2_int8,
       out,
       ScalarType::Char,
-      /*require_channels_last=*/channel_broadcast,
+      /*require_channels_last=*/false,
       /*require_same_sizes=*/!channel_broadcast);
+  if (channel_broadcast) {
+    const Tensor& full_input =
+        input1_int8.numel() > input2_int8.numel() ? input1_int8 : input2_int8;
+    ET_CHECK_MSG(
+        out.sizes() == full_input.sizes(),
+        "quantized_add_out: output must have the broadcast result shape");
+  }
 
   // Validate quantization parameters
   validate_quantization_params(
@@ -101,7 +107,10 @@ Tensor& quantized_add_out(
       std::swap<int>(input1_shift_val, input2_shift_val);
       std::swap<int8_t*>(input1_ptr, input2_ptr);
     }
-    adds_per_loop = input1_int8.size(1);
+    // The broadcast operand holds one value per channel and channels are
+    // contiguous, so its element count is the repeat length.
+    adds_per_loop = static_cast<int32_t>(
+        std::min(input1_int8.numel(), input2_int8.numel()));
   } else {
     adds_per_loop = out.numel();
   }

backends/cortex_m/ops/op_quantized_mul.cpp

@@ -29,16 +29,21 @@ Tensor& quantized_mul_out(
     const int64_t output_multiplier,
     const int64_t output_shift,
     Tensor& out) {
-  // Validate tensor types and quantization parameters
-
   bool channel_broadcast = is_channel_broadcast(input1_int8, input2_int8);
   validate_cmsis_nn_tensor_requirements(
       input1_int8,
       input2_int8,
       out,
       ScalarType::Char,
-      /*require_channels_last=*/channel_broadcast,
+      /*require_channels_last=*/false,
       /*require_same_sizes=*/!channel_broadcast);
+  if (channel_broadcast) {
+    const Tensor& full_input =
+        input1_int8.numel() > input2_int8.numel() ? input1_int8 : input2_int8;
+    ET_CHECK_MSG(
+        out.sizes() == full_input.sizes(),
+        "quantized_mul_out: output must have the broadcast result shape");
+  }
 
   const int32_t kIdentityMultiplier(/*value=*/1);
   const int32_t kZeroShift(/*value=*/0);
@@ -70,7 +75,10 @@ Tensor& quantized_mul_out(
       std::swap<int8_t*>(input1_ptr, input2_ptr);
     }
 
-    muls_per_loop = input1_int8.size(1);
+    // The broadcast operand holds one value per channel and channels are
+    // contiguous, so its element count is the repeat length.
+    muls_per_loop = static_cast<int32_t>(
+        std::min(input1_int8.numel(), input2_int8.numel()));
   } else {
     muls_per_loop = out.numel();
   }

backends/cortex_m/ops/operators.py

@@ -13,8 +13,8 @@
 import torch.nn.functional as F
 from executorch.backends.cortex_m.passes.passes_utils import (
     dequantize_per_tensor_cmsis,
-    is_channel_broadcast,
     is_channels_last,
+    is_flat_channel_broadcast,
     quantize_per_tensor_cmsis,
     requantize_cmsis,
     SHIFT_INT8,
@@ -154,7 +154,7 @@ def quantized_add_meta(
     activation_min: int,
     activation_max: int,
 ) -> torch.Tensor:
-    assert self.shape == other.shape or is_channel_broadcast(self, other), (
+    assert self.shape == other.shape or is_flat_channel_broadcast(self, other), (
         "Cortex-M quantized_add: broadcasting is not yet supported except for channel dim — "
         f"got self.shape={self.shape}, other.shape={other.shape}"
     )
@@ -181,7 +181,7 @@ def quantized_add_impl(
     activation_min: int,
     activation_max: int,
 ) -> torch.Tensor:
-    assert self.shape == other.shape or is_channel_broadcast(self, other), (
+    assert self.shape == other.shape or is_flat_channel_broadcast(self, other), (
         "Cortex-M quantized_add: broadcasting is not yet supported except for channel dim — "
         f"got self.shape={self.shape}, other.shape={other.shape}"
     )
@@ -225,7 +225,7 @@ def quantized_mul_meta(
     output_shift: int,
 ) -> torch.Tensor:
     # Broadcast to output shape
-    assert self.shape == other.shape or is_channel_broadcast(self, other), (
+    assert self.shape == other.shape or is_flat_channel_broadcast(self, other), (
         "Cortex-M quantized_mul: broadcasting is not yet supported except for channel dim — "
         f"got self.shape={self.shape}, other.shape={other.shape}"
     )
@@ -249,7 +249,7 @@ def quantized_mul_impl(
     # CMSIS-NN kernel multiplies raw int8 tensors (after zero-point offset) and
     # only uses the output multiplier/shift for rescaling. Mirror that here to
     # keep the composite implementation numerically aligned with the backend.
-    assert self.shape == other.shape or is_channel_broadcast(self, other), (
+    assert self.shape == other.shape or is_flat_channel_broadcast(self, other), (
         "Cortex-M quantized_mul: broadcasting is not yet supported except for channel dim — "
         f"got self.shape={self.shape}, other.shape={other.shape}"
     )

backends/cortex_m/passes/passes_utils.py

@@ -337,9 +337,11 @@ def to_physical_order(logical_pad: list[int], tensor: torch.Tensor) -> list[int]
 
 
 def is_channel_broadcast(tensor1: torch.Tensor, tensor2: torch.Tensor) -> bool:
-    """
-    Check if tensor1 is broadcasted to tensor2 along channel dimension.
-    Assumes tensor2 has shape [N, C, ...] and tensor1 has shape [N, 1, ...] or [1, C, ...].
+    """Check for a broadcast of one value per channel, on logical NCHW shapes.
+
+    This is the question the quantizer asks, before any layout transform, so
+    the channel is dimension one regardless of memory format. Callers that also
+    require a particular memory format check it separately.
     """
     if tensor1.dim() != tensor2.dim():
         return False
@@ -353,3 +355,43 @@ def is_channel_broadcast(tensor1: torch.Tensor, tensor2: torch.Tensor) -> bool:
     tensor2_channels_only = tensor2.numel() == tensor2.size(1)
 
     return channel_match and (tensor1_channels_only or tensor2_channels_only)
+
+
+def is_flat_channel_broadcast(tensor1: torch.Tensor, tensor2: torch.Tensor) -> bool:
+    """Check that a broadcast is a flat repeat in memory, as the kernels need.
+
+    The kernels repeat the broadcast operand along the innermost axis: element
+    ``i`` of the result reads ``small[i % small.numel()]``. The operand itself is
+    a flat run whatever its dim order, since it has at most one non-unit extent.
+    What has to hold is that the *larger* operand carries that same axis
+    innermost in memory, and its dim order is what says so.
+
+    Both activation contracts satisfy this and disagree on which logical axis it
+    is -- legacy is a channels-last ``[N, C, H, W]``, explicit a contiguous
+    ``[N, H, W, C]`` -- so the axis is taken from the broadcast operand's own
+    non-unit extent rather than assumed. A genuinely contiguous ``[N, C, H, W]``
+    is refused: its innermost axis is W, and the flat repeat would stride across
+    the wrong elements.
+
+    An operand with no non-unit extent repeats a single value and is accepted
+    without consulting a dim order, which is the one case where the dim order
+    cannot identify the channel axis.
+    """
+    if tensor1.dim() != tensor2.dim() or tensor1.dim() != 4:
+        return False
+    if tensor1.numel() == tensor2.numel():
+        return False
+
+    larger, smaller = (
+        (tensor1, tensor2) if tensor1.numel() > tensor2.numel() else (tensor2, tensor1)
+    )
+    non_unit_dims = [dim for dim, size in enumerate(smaller.shape) if size != 1]
+    if len(non_unit_dims) > 1:
+        return False
+    if non_unit_dims:
+        channel_axis = non_unit_dims[0]
+        if larger.shape[channel_axis] != smaller.shape[channel_axis]:
+            return False
+        if larger.dim_order()[-1] != channel_axis:
+            return False
+    return larger.numel() % smaller.numel() == 0

backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py

@@ -227,6 +227,46 @@ def forward(self, x):
         )
 
 
+class AddNhwc(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("bias", _int8_values((1, 1, 1, 3)))
+
+    def forward(self, x):
+        return torch.ops.cortex_m.quantized_add.default(
+            x,
+            0,
+            1 << 30,
+            -1,
+            self.bias,
+            0,
+            1 << 30,
+            -1,
+            0,
+            1 << 30,
+            -1,
+            -128,
+            127,
+        )
+
+
+class MulNhwc(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("bias", _int8_values((1, 1, 1, 3)))
+
+    def forward(self, x):
+        return torch.ops.cortex_m.quantized_mul.default(
+            x,
+            0,
+            self.bias,
+            0,
+            0,
+            1 << 30,
+            -1,
+        )
+
+
 def test_conv2d_nhwc_runs_on_fvp(cortex_m_target):
     _run_on_fvp(
         Conv2dNhwc(),
@@ -294,3 +334,23 @@ def test_pad_contiguous_runs_on_fvp_with_singleton_height(cortex_m_target):
         exir_ops.edge.cortex_m.pad_contiguous.default,
         cortex_m_target,
     )
+
+
+def test_channel_broadcast_add_nhwc_runs_on_fvp(cortex_m_target):
+    _run_on_fvp(
+        AddNhwc(),
+        _int8_values((1, 5, 7, 3)),
+        exir_ops.edge.cortex_m.quantized_add.default,
+        cortex_m_target,
+        atol=1,
+    )
+
+
+def test_channel_broadcast_mul_nhwc_runs_on_fvp(cortex_m_target):
+    _run_on_fvp(
+        MulNhwc(),
+        _int8_values((1, 5, 7, 3)),
+        exir_ops.edge.cortex_m.quantized_mul.default,
+        cortex_m_target,
+        atol=1,
+    )

backends/cortex_m/test/test_quantized_conv2d_layout.py

@@ -104,6 +104,28 @@ def _run_max_pool2d(op, x):
     )
 
 
+def _run_add(op, x, bias):
+    return op(
+        x,
+        0,
+        1 << 30,
+        -1,
+        bias,
+        0,
+        1 << 30,
+        -1,
+        0,
+        1 << 30,
+        -1,
+        -128,
+        127,
+    )
+
+
+def _run_mul(op, x, bias):
+    return op(x, 0, bias, 0, 0, 1 << 30, -1)
+
+
 def test_nhwc_conv2d_matches_legacy_layout():
     torch.manual_seed(0)
     x = torch.randint(-8, 8, (1, 3, 8, 8), dtype=torch.int8)
@@ -200,6 +222,42 @@ def test_nhwc_max_pool2d_matches_legacy_layout():
     torch.testing.assert_close(explicit, legacy.permute(0, 2, 3, 1))
 
 
+def test_nhwc_channel_broadcast_add_matches_legacy_layout():
+    x = torch.randint(-8, 8, (1, 4, 5, 7), dtype=torch.int8)
+    bias = torch.randint(-4, 4, (1, 4, 1, 1), dtype=torch.int8)
+
+    legacy = _run_add(
+        torch.ops.cortex_m.quantized_add,
+        x.to(memory_format=torch.channels_last),
+        bias.to(memory_format=torch.channels_last),
+    )
+    explicit = _run_add(
+        torch.ops.cortex_m.quantized_add,
+        x.permute(0, 2, 3, 1).contiguous(),
+        bias.permute(0, 2, 3, 1).contiguous(),
+    )
+
+    torch.testing.assert_close(explicit, legacy.permute(0, 2, 3, 1))
+
+
+def test_nhwc_channel_broadcast_mul_matches_legacy_layout():
+    x = torch.randint(-8, 8, (1, 4, 5, 7), dtype=torch.int8)
+    bias = torch.randint(-4, 4, (1, 4, 1, 1), dtype=torch.int8)
+
+    legacy = _run_mul(
+        torch.ops.cortex_m.quantized_mul,
+        x.to(memory_format=torch.channels_last),
+        bias.to(memory_format=torch.channels_last),
+    )
+    explicit = _run_mul(
+        torch.ops.cortex_m.quantized_mul,
+        x.permute(0, 2, 3, 1).contiguous(),
+        bias.permute(0, 2, 3, 1).contiguous(),
+    )
+
+    torch.testing.assert_close(explicit, legacy.permute(0, 2, 3, 1))
+
+
 def test_nhwc_conv2d_fake_shape_is_logical_nhwc():
     with FakeTensorMode():
         output = _run_conv2d(
@@ -382,3 +440,27 @@ def make_node(
             assert required_cmsis_nn_buffer_sizes(
                 legacy, backend
             ) == required_cmsis_nn_buffer_sizes(explicit, backend)
+
+
+def test_single_channel_broadcast_add_matches_legacy_layout():
+    """A single-channel broadcast is the case a dim-order derivation gets wrong.
+
+    A channels-last tensor with one channel serializes its dim order as
+    (0, 2, 1, 3), which names no channel axis at all, so the repeat length has
+    to come from the broadcast operand's element count instead.
+    """
+    x = torch.randint(-8, 8, (1, 1, 5, 7), dtype=torch.int8)
+    bias = torch.randint(-4, 4, (1, 1, 1, 1), dtype=torch.int8)
+
+    legacy = _run_add(
+        torch.ops.cortex_m.quantized_add,
+        x.to(memory_format=torch.channels_last),
+        bias.to(memory_format=torch.channels_last),
+    )
+    explicit = _run_add(
+        torch.ops.cortex_m.quantized_add,
+        x.permute(0, 2, 3, 1).contiguous(),
+        bias.permute(0, 2, 3, 1).contiguous(),
+    )
+
+    torch.testing.assert_close(explicit, legacy.permute(0, 2, 3, 1))