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

@@ -14,7 +14,7 @@
 from executorch.backends.cortex_m.passes.passes_utils import (
     dequantize_per_tensor_cmsis,
     is_channels_last,
-    is_channel_broadcast,
+    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,25 @@ 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.
+
+    Deliberately layout-free. Channels are innermost under both activation
+    contracts, so the broadcast operand's element count is the repeat length and
+    the larger operand is a whole number of copies of it. Deriving the channel
+    axis instead would be unsound: a serialized dim order does not identify it
+    when the channel count is one.
+    """
+    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)
+    )
+    if sum(1 for size in smaller.shape if size != 1) > 1:
+        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,
+    )