diff --git a/backends/cortex_m/ops/cortex_m_ops_common.h b/backends/cortex_m/ops/cortex_m_ops_common.h index bcaed0a1bc7..9752019633e 100644 --- a/backends/cortex_m/ops/cortex_m_ops_common.h +++ b/backends/cortex_m/ops/cortex_m_ops_common.h @@ -17,6 +17,7 @@ #include #include +#include #include #include #include @@ -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( diff --git a/backends/cortex_m/ops/op_quantized_add.cpp b/backends/cortex_m/ops/op_quantized_add.cpp index f93bb6c1be9..fb9b81bddd7 100644 --- a/backends/cortex_m/ops/op_quantized_add.cpp +++ b/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(input1_shift_val, input2_shift_val); std::swap(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( + std::min(input1_int8.numel(), input2_int8.numel())); } else { adds_per_loop = out.numel(); } diff --git a/backends/cortex_m/ops/op_quantized_mul.cpp b/backends/cortex_m/ops/op_quantized_mul.cpp index 93ce2303d64..17977079ef0 100644 --- a/backends/cortex_m/ops/op_quantized_mul.cpp +++ b/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(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( + std::min(input1_int8.numel(), input2_int8.numel())); } else { muls_per_loop = out.numel(); } diff --git a/backends/cortex_m/ops/operators.py b/backends/cortex_m/ops/operators.py index 96a7bdc0165..37d03ccbe04 100644 --- a/backends/cortex_m/ops/operators.py +++ b/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}" ) diff --git a/backends/cortex_m/passes/passes_utils.py b/backends/cortex_m/passes/passes_utils.py index bcb828c5928..467893f5bdc 100644 --- a/backends/cortex_m/passes/passes_utils.py +++ b/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 diff --git a/backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py b/backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py index a64e583be49..09ec4ca2656 100644 --- a/backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py +++ b/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, + ) diff --git a/backends/cortex_m/test/test_quantized_conv2d_layout.py b/backends/cortex_m/test/test_quantized_conv2d_layout.py index 896bd950b20..585212ae8d0 100644 --- a/backends/cortex_m/test/test_quantized_conv2d_layout.py +++ b/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( @@ -213,21 +271,43 @@ def test_nhwc_conv2d_fake_shape_is_logical_nhwc(): assert output.dim_order() == (0, 1, 2, 3) -def test_pad_preserves_singleton_height_layout(): +def test_pad_contiguous_preserves_singleton_height_layout(): x = torch.arange(1 * 1 * 5 * 3, dtype=torch.int8).reshape(1, 1, 5, 3) pre_pad = [0, 0, 1, 0] post_pad = [0, 0, 2, 0] - actual = torch.ops.cortex_m.pad(x, pre_pad, post_pad, -7) + actual = torch.ops.cortex_m.pad_contiguous(x, pre_pad, post_pad, -7) expected = torch.nn.functional.pad(x, (0, 0, 1, 2, 0, 0, 0, 0), value=-7) assert actual.shape == torch.Size([1, 1, 8, 3]) torch.testing.assert_close(actual, expected) -def test_pad_rejects_invalid_padding_length(): +def test_pad_contiguous_handles_every_supported_rank(): + """Ranks below four are contiguous by construction, so this entry point + covers them too. That is what lets it eventually replace ``cortex_m::pad`` + outright rather than sitting beside it forever.""" + for shape, pad in ( + ((2, 3, 4, 5), [0, 0, 1, 2]), + ((2, 3, 4), [0, 0, 1, 2]), + ((3, 5), [0, 0, 1, 2]), + ((6,), [0, 0, 0, 2]), + ): + x = torch.randint(-8, 8, shape, dtype=torch.int8) + rank = len(shape) + offset = 4 - rank + actual = torch.ops.cortex_m.pad_contiguous(x, pad, pad, -7) + flat = [] + for dim in reversed(range(rank)): + flat.extend([pad[offset + dim], pad[offset + dim]]) + expected = torch.nn.functional.pad(x, flat, value=-7) + assert actual.shape == expected.shape, shape + torch.testing.assert_close(actual, expected) + + +def test_pad_contiguous_rejects_invalid_padding_length(): with pytest.raises(RuntimeError, match="expects four padding values per side"): - torch.ops.cortex_m.pad( + torch.ops.cortex_m.pad_contiguous( torch.zeros((1, 1, 5, 3), dtype=torch.int8), [0, 0, 0], [0, 0, 0, 0], @@ -236,7 +316,7 @@ def test_pad_rejects_invalid_padding_length(): with FakeTensorMode(): with pytest.raises(RuntimeError, match="expects four padding values per side"): - torch.ops.cortex_m.pad( + torch.ops.cortex_m.pad_contiguous( torch.zeros((1, 1, 5, 3), dtype=torch.int8), [0, 0, 0], [0, 0, 0, 0], @@ -360,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))