executorch 1.6.0__cp315-cp315-win_amd64.whl

executorch/backends/cortex_m/ops/operators.py

@@ -0,0 +1,1985 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+# 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.
+
+import math
+from math import prod
+from typing import Sequence
+
+import torch
+import torch.nn.functional as F
+from executorch.backends.cortex_m.quantizer.quantization_configs import (
+    CMSIS_SOFTMAX_SCALE,
+    CMSIS_SOFTMAX_ZERO_POINT,
+)
+from executorch.backends.cortex_m.utils import (
+    dequantize_per_tensor_cmsis,
+    is_channel_broadcast,
+    is_channels_last,
+    quantize_per_tensor_cmsis,
+    requantize_cmsis,
+    SHIFT_INT8,
+)
+from executorch.exir._warnings import experimental
+from executorch.exir.dialects._ops import ops as exir_ops
+
+# To provide the implementation of the operators
+from torch.library import impl, Library, register_fake
+
+# New operator library with a custom namespace to allow fusion etc.
+lib = Library("cortex_m", "DEF")
+
+_EXPLICIT_LAYOUT_EXPERIMENTAL = (
+    "This explicit-layout Cortex-M operator may change while the legacy "
+    "dim-order operators remain supported."
+)
+
+SOFTMAX_INPUT_INTEGER_BITS = 5
+
+
+###
+# dequantize_per_tensor
+###
+
+lib.define(
+    "quantize_per_tensor(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype) -> (Tensor Z)"
+)
+
+lib.define(
+    "quantize_per_tensor.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+@register_fake("cortex_m::quantize_per_tensor")  # type: ignore[misc]
+def quantize_per_tensor_meta(
+    input: torch.Tensor,
+    scale: float,
+    zero_point: int,
+    quant_min: int,
+    quant_max: int,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    return torch.empty_like(input, dtype=dtype)
+
+
+@impl(lib, "quantize_per_tensor", "CompositeExplicitAutograd")  # type: ignore[misc]
+def quantize_per_tensor_impl(
+    input: torch.Tensor,
+    scale: float,
+    zero_point: int,
+    quant_min: int,
+    quant_max: int,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """
+    The implementation of the quantize_per_tensor operator is the same as the
+    quantize_per_tensor operator in the edge dialect.
+    """
+    return exir_ops.edge.quantized_decomposed.quantize_per_tensor.default(
+        input, scale, zero_point, quant_min, quant_max, dtype
+    )
+
+
+###
+# dequantize_per_tensor
+###
+
+lib.define(
+    "dequantize_per_tensor(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype) -> (Tensor Z)"
+)
+lib.define(
+    "dequantize_per_tensor.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+@register_fake("cortex_m::dequantize_per_tensor")  # type: ignore[misc]
+def dequantize_per_tensor_meta(
+    input: torch.Tensor,
+    scale: float,
+    zero_point: int,
+    quant_min: int,
+    quant_max: int,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    return torch.empty_like(input, dtype=torch.float)
+
+
+@impl(lib, "dequantize_per_tensor", "CompositeExplicitAutograd")  # type: ignore[misc]
+def dequantize_per_tensor_impl(
+    input: torch.Tensor,
+    scale: float,
+    zero_point: int,
+    quant_min: int,
+    quant_max: int,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """
+    The implementation of the dequantize_per_tensor operator is the same as the
+    dequantize_per_tensor operator in the edge dialect.
+    """
+    return exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default(
+        input, scale, zero_point, quant_min, quant_max, dtype
+    )
+
+
+# Define the operator schema with multipliers and shifts (11 args)
+lib.define(
+    "quantized_add("
+    "Tensor self, int self_zero_point, int self_multiplier, int self_shift, "
+    "Tensor other, int other_zero_point, int other_multiplier, int other_shift, "
+    "int output_zero_point, int output_multiplier, int output_shift, "
+    "int activation_min, int activation_max) -> Tensor"
+)
+
+lib.define(
+    "quantized_add.out("
+    "Tensor self, int self_zero_point, int self_multiplier, int self_shift, "
+    "Tensor other, int other_zero_point, int other_multiplier, int other_shift, "
+    "int output_zero_point, int output_multiplier, int output_shift, "
+    "int activation_min, int activation_max, "
+    "*, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+@register_fake("cortex_m::quantized_add")  # type: ignore[misc]
+def quantized_add_meta(
+    self: torch.Tensor,
+    self_zero_point: int,
+    self_multiplier: int,
+    self_shift: int,
+    other: torch.Tensor,
+    other_zero_point: int,
+    other_multiplier: int,
+    other_shift: int,
+    output_zero_point: int,
+    output_multiplier: int,
+    output_shift: int,
+    activation_min: int,
+    activation_max: int,
+) -> torch.Tensor:
+    assert self.shape == other.shape or is_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}"
+    )
+    if self.numel() > other.numel():
+        output_tensor = self
+    else:
+        output_tensor = other
+    return torch.empty_like(output_tensor)
+
+
+@impl(lib, "quantized_add", "CompositeExplicitAutograd")  # type: ignore[misc]
+def quantized_add_impl(
+    self: torch.Tensor,
+    self_zero_point: int,
+    self_multiplier: int,
+    self_shift: int,
+    other: torch.Tensor,
+    other_zero_point: int,
+    other_multiplier: int,
+    other_shift: int,
+    output_zero_point: int,
+    output_multiplier: int,
+    output_shift: int,
+    activation_min: int,
+    activation_max: int,
+) -> torch.Tensor:
+    assert self.shape == other.shape or is_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}"
+    )
+    self_shifted = (self.to(torch.int32) - self_zero_point) << SHIFT_INT8
+    self_fp = requantize_cmsis(self_shifted, self_multiplier, self_shift)
+
+    other_shifted = (other.to(torch.int32) - other_zero_point) << SHIFT_INT8
+    other_fp = requantize_cmsis(other_shifted, other_multiplier, other_shift)
+
+    result_fp = self_fp + other_fp
+    result_quantized = requantize_cmsis(result_fp, output_multiplier, output_shift)
+    result = torch.clamp(
+        result_quantized + output_zero_point, activation_min, activation_max
+    ).to(torch.int8)
+    return result
+
+
+# ===================================================================
+# QUANTIZED MUL OPERATION DEFINITION
+# ===================================================================
+lib.define(
+    "quantized_mul("
+    "Tensor self, int self_zero_point, "
+    "Tensor other, int other_zero_point, "
+    "int output_zero_point, int output_multiplier, int output_shift) -> Tensor"
+)
+lib.define(
+    "quantized_mul.out("
+    "Tensor self, int self_zero_point, "
+    "Tensor other, int other_zero_point, "
+    "int output_zero_point, int output_multiplier, int output_shift, "
+    "*, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+@register_fake("cortex_m::quantized_mul")  # type: ignore[misc]
+def quantized_mul_meta(
+    self: torch.Tensor,
+    self_zero_point: int,
+    other: torch.Tensor,
+    other_zero_point: int,
+    output_zero_point: int,
+    output_multiplier: int,
+    output_shift: int,
+) -> torch.Tensor:
+    # Broadcast to output shape
+    assert self.shape == other.shape or is_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}"
+    )
+    if self.numel() > other.numel():
+        output_tensor = self
+    else:
+        output_tensor = other
+    return torch.empty_like(output_tensor)
+
+
+@impl(lib, "quantized_mul", "CompositeExplicitAutograd")  # type: ignore[misc]
+def quantized_mul_impl(
+    self: torch.Tensor,
+    self_zero_point: int,
+    other: torch.Tensor,
+    other_zero_point: int,
+    output_zero_point: int,
+    output_multiplier: int,
+    output_shift: int,
+) -> torch.Tensor:
+    # 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), (
+        "Cortex-M quantized_mul: broadcasting is not yet supported except for channel dim — "
+        f"got self.shape={self.shape}, other.shape={other.shape}"
+    )
+    self_int = self.to(torch.int32) - self_zero_point
+    other_int = other.to(torch.int32) - other_zero_point
+    result_fp = self_int * other_int
+    result_quantized = requantize_cmsis(result_fp, output_multiplier, output_shift)
+    result = torch.clamp(result_quantized + output_zero_point, -128, 127).to(torch.int8)
+    return result
+
+
+# ===================================================================
+# QUANTIZED DIV OPERATION DEFINITION
+# ===================================================================
+lib.define(
+    "quantized_div("
+    "Tensor self, int self_zero_point, "
+    "Tensor other, int other_zero_point, "
+    "int output_zero_point, float output_scale) -> Tensor"
+)
+lib.define(
+    "quantized_div.out("
+    "Tensor self, int self_zero_point, "
+    "Tensor other, int other_zero_point, "
+    "int output_zero_point, float output_scale, "
+    "*, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+@register_fake("cortex_m::quantized_div")  # type: ignore[misc]
+def quantized_div_meta(
+    self: torch.Tensor,
+    self_zero_point: int,
+    other: torch.Tensor,
+    other_zero_point: int,
+    output_zero_point: int,
+    output_scale: float,
+) -> torch.Tensor:
+    # Division is not commutative, so broadcasting (handled via operand swaps in
+    # quantized_mul) is not supported: require identical shapes.
+    assert self.shape == other.shape, (
+        "Cortex-M quantized_div: broadcasting is not supported — "
+        f"got self.shape={self.shape}, other.shape={other.shape}"
+    )
+    return torch.empty_like(self)
+
+
+@impl(lib, "quantized_div", "CompositeExplicitAutograd")  # type: ignore[misc]
+def quantized_div_impl(
+    self: torch.Tensor,
+    self_zero_point: int,
+    other: torch.Tensor,
+    other_zero_point: int,
+    output_zero_point: int,
+    output_scale: float,
+) -> torch.Tensor:
+    # Mirror the kernel: the quotient of the zero-point-corrected int8/int16
+    # operands is evaluated in float and rescaled by the effective scale
+    # (scale_in1 / (scale_in2 * scale_out)) that the AoT pass carries directly.
+    assert self.shape == other.shape, (
+        "Cortex-M quantized_div: broadcasting is not supported — "
+        f"got self.shape={self.shape}, other.shape={other.shape}"
+    )
+    if self.dtype not in (torch.int8, torch.int16):
+        raise TypeError(
+            f"cortex_m.quantized_div: expected int8 or int16 inputs, got {self.dtype}"
+        )
+    self_fp = (self.to(torch.int32) - self_zero_point).to(torch.float32)
+    other_fp = (other.to(torch.int32) - other_zero_point).to(torch.float32)
+
+    quotient = torch.where(other_fp != 0, self_fp / other_fp, torch.zeros_like(self_fp))
+    result = torch.round(quotient * output_scale) + output_zero_point
+    dtype_info = torch.iinfo(self.dtype)
+    return torch.clamp(result, dtype_info.min, dtype_info.max).to(self.dtype)
+
+
+# ===================================================================
+# QUANTIZED ACTIVATION (LUT) OPERATION DEFINITION
+# ===================================================================
+# Generic table-lookup activation. The 256-entry int8 LUT is precomputed AoT
+# from the input/output qparams and the activation function (sigmoid, tanh,
+# silu, ...), so the kernel is identical regardless of which activation it
+# evaluates: out[i] = lut[input[i] + 128].
+lib.define("quantized_activation(Tensor input, Tensor lut) -> Tensor")
+lib.define(
+    "quantized_activation.out(Tensor input, Tensor lut, *, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+@register_fake("cortex_m::quantized_activation")  # type: ignore[misc]
+def quantized_activation_meta(input: torch.Tensor, lut: torch.Tensor) -> torch.Tensor:
+    assert input.dtype == torch.int8, "quantized_activation input must be int8"
+    assert lut.dtype == torch.int8 and lut.numel() == 256, (
+        "quantized_activation lut must be int8 with 256 entries; "
+        f"got dtype={lut.dtype}, numel={lut.numel()}"
+    )
+    return torch.empty_like(input)
+
+
+@impl(lib, "quantized_activation", "CompositeExplicitAutograd")  # type: ignore[misc]
+def quantized_activation_impl(input: torch.Tensor, lut: torch.Tensor) -> torch.Tensor:
+    indices = input.to(torch.int32) + 128
+    return lut[indices].to(torch.int8)
+
+
+# ===================================================================
+# QUANTIZED BATCH MATMUL OPERATION DEFINITION
+# ===================================================================
+lib.define(
+    "quantized_batch_matmul("
+    "Tensor lhs, int lhs_zero_point, "
+    "Tensor rhs_transposed, int rhs_zero_point, "
+    "int output_zero_point, int output_multiplier, int output_shift, "
+    "Tensor scratch) -> Tensor"
+)
+lib.define(
+    "quantized_batch_matmul.out("
+    "Tensor lhs, int lhs_zero_point, "
+    "Tensor rhs_transposed, int rhs_zero_point, "
+    "int output_zero_point, int output_multiplier, int output_shift, "
+    "Tensor scratch, "
+    "*, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+@register_fake("cortex_m::quantized_batch_matmul")  # type: ignore[misc]
+def quantized_batch_matmul_meta(
+    lhs: torch.Tensor,
+    lhs_zero_point: int,
+    rhs_transposed: torch.Tensor,
+    rhs_zero_point: int,
+    output_zero_point: int,
+    output_multiplier: int,
+    output_shift: int,
+    scratch: torch.Tensor,
+) -> torch.Tensor:
+    batch, lhs_rows, inner = lhs.shape
+    batch_rhs, rhs_cols, inner_rhs = rhs_transposed.shape
+    assert batch == batch_rhs and inner == inner_rhs
+    return torch.empty((batch, lhs_rows, rhs_cols), dtype=torch.int8, device=lhs.device)
+
+
+@impl(lib, "quantized_batch_matmul", "CompositeExplicitAutograd")  # type: ignore[misc]
+def quantized_batch_matmul_impl(
+    lhs: torch.Tensor,
+    lhs_zero_point: int,
+    rhs_transposed: torch.Tensor,
+    rhs_zero_point: int,
+    output_zero_point: int,
+    output_multiplier: int,
+    output_shift: int,
+    scratch: torch.Tensor,
+) -> torch.Tensor:
+    # Offsets are negated zero points (CMSIS-NN convention)
+    lhs_fp = lhs.to(torch.float32) + float(lhs_zero_point)
+    rhs_t_fp = rhs_transposed.to(torch.float32) + float(rhs_zero_point)
+    rhs_fp = rhs_t_fp.permute(0, 2, 1)
+    acc = torch.bmm(lhs_fp, rhs_fp).to(torch.int32)
+    result = requantize_cmsis(acc, output_multiplier, output_shift)
+    return torch.clamp(result + output_zero_point, -128, 127).to(torch.int8)
+
+
+# ===================================================================
+# MINIMUM/MAXIMUM OPERATION DEFINITIONS
+# ===================================================================
+lib.define("minimum(Tensor self, Tensor other) -> Tensor")
+lib.define("minimum.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)")
+
+
+@register_fake("cortex_m::minimum")  # type: ignore[misc]
+def minimum_meta(self: torch.Tensor, other: torch.Tensor) -> torch.Tensor:
+    # other is a scalar, so use initial shape.
+    if other.numel() == 1:
+        return torch.empty_like(self)
+    else:
+        # otherwise broadcast the shape.
+        return torch.empty(
+            torch.broadcast_shapes(self.shape, other.shape),
+            dtype=self.dtype,
+            device=self.device,
+        )
+
+
+@impl(lib, "minimum", "CompositeExplicitAutograd")  # type: ignore[misc]
+def minimum_impl(self: torch.Tensor, other: torch.Tensor) -> torch.Tensor:
+    return torch.minimum(self, other)
+
+
+lib.define("maximum(Tensor self, Tensor other) -> Tensor")
+lib.define("maximum.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)")
+
+
+@register_fake("cortex_m::maximum")  # type: ignore[misc]
+def maximum_meta(self: torch.Tensor, other: torch.Tensor) -> torch.Tensor:
+    assert self.dtype == other.dtype, (
+        "Cortex-M maximum: dtype mismatch — "
+        f"got self.dtype={self.dtype}, other.dtype={other.dtype}"
+    )
+    # other is a scalar, so use initial shape.
+    if other.numel() == 1:
+        return torch.empty_like(self)
+    else:
+        # otherwise broadcast the shape.
+        return torch.empty(
+            torch.broadcast_shapes(self.shape, other.shape),
+            dtype=self.dtype,
+            device=self.device,
+        )
+
+
+@impl(lib, "maximum", "CompositeExplicitAutograd")  # type: ignore[misc]
+def maximum_impl(self: torch.Tensor, other: torch.Tensor) -> torch.Tensor:
+    return torch.maximum(self, other)
+
+
+# ===================================================================
+# QUANTIZED LINEAR OPERATION DEFINITION
+# ===================================================================
+
+lib.define(
+    "quantized_linear.out("
+    "Tensor input,  "
+    "Tensor weights, "
+    "Tensor? bias, "
+    "Tensor? kernel_sum, "
+    "int input_offset, "
+    "int filter_offset, "
+    "int output_offset, "
+    "int[] requantize_multipliers, "
+    "int[] requantize_shifts, "
+    "int activation_max, "
+    "int activation_min, "
+    "*, Tensor(a!) out"
+    ") -> Tensor(a!)"
+)
+
+# Define functional variant (non-out version)
+lib.define(
+    "quantized_linear("
+    "Tensor input,  "
+    "Tensor weights, "
+    "Tensor? bias, "
+    "Tensor? kernel_sum, "
+    "int input_offset, "
+    "int filter_offset, "
+    "int output_offset, "
+    "int[] requantize_multipliers, "
+    "int[] requantize_shifts, "
+    "int activation_max, "
+    "int activation_min"
+    ") -> Tensor"
+)
+
+
+# Fake meta function for shape inference (functional variant)
+@register_fake("cortex_m::quantized_linear")  # type: ignore[misc]
+def quantized_linear_meta(
+    input: torch.Tensor,
+    weights: torch.Tensor,
+    bias: torch.Tensor | None,
+    kernel_sum: torch.Tensor | None,
+    input_offset: int,
+    filter_offset: int,
+    output_offset: int,
+    requantize_multipliers: torch.Tensor,
+    requantize_shifts: torch.Tensor,
+    activation_max: int,
+    activation_min: int,
+) -> torch.Tensor:
+
+    shape = (*input.shape[:-1], weights.shape[0])
+    return torch.empty(shape, dtype=input.dtype, device=input.device)
+
+
+# Functional variant implementation
+@impl(lib, "quantized_linear", "CompositeExplicitAutograd")  # type: ignore[misc]
+def quantized_linear_impl(
+    input: torch.Tensor,
+    weights: torch.Tensor,
+    bias: torch.Tensor | None,
+    kernel_sum: torch.Tensor | None,
+    input_offset: int,
+    filter_offset: int,
+    output_offset: int,
+    requantize_multipliers: torch.Tensor,
+    requantize_shifts: torch.Tensor,
+    activation_max: int,
+    activation_min: int,
+) -> torch.Tensor:
+    """
+    Functional variant - creates output tensor and calls out variant
+    """
+
+    # 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)
+        new_shape = (prod(input.shape[:-1]), input.shape[-1])
+        input_reshaped = input_int32.reshape(new_shape)
+
+        lhs_sum = torch.sum(input_reshaped, dim=-1, keepdim=True) * filter_offset
+        output = torch.mm(input_reshaped, weights_int32.T) + lhs_sum + kernel_sum
+        output_shape = (*input.shape[:-1], output.shape[-1])
+        output_reshaped = output.reshape(output_shape)
+    else:
+        weights_int32 = weights.to(torch.int32) + filter_offset
+
+        input_int32 = input.to(torch.int32) + input_offset
+        new_shape = (prod(input.shape[:-1]), input.shape[-1])
+        input_reshaped = input_int32.reshape(new_shape)
+
+        output = torch.mm(input_reshaped, weights_int32.T)
+        if bias is not None:
+            output = output + bias
+        output_shape = (*input.shape[:-1], output.shape[-1])
+        output_reshaped = output.reshape(output_shape)
+
+    output = requantize_cmsis(
+        output_reshaped,
+        int(requantize_multipliers[0]),
+        int(requantize_shifts[0]),
+    )
+    output += output_offset
+    output = torch.clamp(output, activation_min, activation_max).to(torch.int8)
+    return output
+
+
+# ===================================================================
+# SOFTMAX OPERATION DEFINITION
+# ===================================================================
+
+lib.define(
+    "softmax(Tensor input, int dim, int input_zero_point, int output_zero_point, int input_multiplier, int input_shift, int diff_min) -> Tensor"
+)
+lib.define(
+    "softmax.out(Tensor input, int dim, int input_zero_point, int output_zero_point, int input_multiplier, int input_shift, int diff_min, *, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+@register_fake("cortex_m::softmax")  # type: ignore[misc]
+def softmax_meta(
+    input: torch.Tensor,
+    dim: int,
+    input_zero_point: int,
+    output_zero_point: int,
+    input_multiplier: int,
+    input_shift: int,
+    diff_min: int,
+) -> torch.Tensor:
+    return torch.empty_like(input, dtype=torch.int8)
+
+
+@impl(lib, "softmax", "CompositeExplicitAutograd")  # type: ignore[misc]
+def softmax_impl(
+    input: torch.Tensor,
+    dim: int,
+    input_zero_point: int,
+    output_zero_point: int,
+    input_multiplier: int,
+    input_shift: int,
+    diff_min: int,
+) -> torch.Tensor:
+    del diff_min  # not used in reference path
+    if input.dtype != torch.int8:
+        raise TypeError(
+            f"cortex_m.softmax: expected int8 input tensor, got {input.dtype}"
+        )
+    if output_zero_point != CMSIS_SOFTMAX_ZERO_POINT:
+        raise ValueError(
+            f"cortex_m.softmax: expected output_zero_point {CMSIS_SOFTMAX_ZERO_POINT}, got {output_zero_point}"
+        )
+
+    real_multiplier = float(input_multiplier) / float(1 << 31)
+    real_multiplier = math.ldexp(real_multiplier, input_shift)
+    input_scale = real_multiplier / float(1 << (31 - SOFTMAX_INPUT_INTEGER_BITS))
+    if input_scale <= 0:
+        raise ValueError(
+            f"cortex_m.softmax: derived non-positive input scale {input_scale}"
+        )
+
+    input_fp = (input.to(torch.int32) - int(input_zero_point)).float() * input_scale
+    probs = torch.softmax(input_fp, dim=dim)
+    quantized = torch.round(probs / CMSIS_SOFTMAX_SCALE) + int(output_zero_point)
+    return quantized.clamp(-128, 127).to(torch.int8)
+
+
+# ===================================================================
+# TRANSPOSE OPERATION DEFINITION
+# ===================================================================
+lib.define("transpose(Tensor input, int[] perm) -> Tensor")
+lib.define("transpose.out(Tensor input, int[] perm, *, Tensor(a!) out) -> Tensor(a!)")
+
+
+@register_fake("cortex_m::transpose")  # type: ignore[misc]
+def transpose_meta(input: torch.Tensor, perm: Sequence[int]) -> torch.Tensor:
+    output_shape = [input.shape[idx] for idx in perm]
+    return torch.empty(output_shape, dtype=input.dtype, device=input.device)
+
+
+@impl(lib, "transpose", "CompositeExplicitAutograd")  # type: ignore[misc]
+def transpose_impl(input: torch.Tensor, perm: Sequence[int]) -> torch.Tensor:
+    return input.permute(tuple(perm)).contiguous()
+
+
+# ===================================================================
+# PAD OPERATION DEFINITION
+# ===================================================================
+lib.define("pad(Tensor input, int[] pre_pad, int[] post_pad, int pad_value) -> Tensor")
+lib.define(
+    "pad.out(Tensor input, int[] pre_pad, int[] post_pad, int pad_value, "
+    "*, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+_NHWC_INV_ORDER = [0, 3, 1, 2]
+
+
+def _pad_to_logical_order(physical_pad: list[int], input: torch.Tensor) -> list[int]:
+    """Inverse of _to_physical_order: map physical-order padding back to logical."""
+    if not is_channels_last(input):
+        return list(physical_pad)
+    return [physical_pad[_NHWC_INV_ORDER[i]] for i in range(4)]
+
+
+@register_fake("cortex_m::pad")  # type: ignore[misc]
+def pad_meta(
+    input: torch.Tensor,
+    pre_pad: list[int],
+    post_pad: list[int],
+    pad_value: int,
+) -> torch.Tensor:
+    rank = input.dim()
+    offset = 4 - rank
+    logical_pre = _pad_to_logical_order(pre_pad, input)
+    logical_post = _pad_to_logical_order(post_pad, input)
+
+    output_shape = list(input.shape)
+    for i in range(rank):
+        output_shape[i] += logical_pre[offset + i] + logical_post[offset + i]
+    result = torch.empty(output_shape, dtype=input.dtype, device=input.device)
+    if is_channels_last(input):
+        result = result.to(memory_format=torch.channels_last)
+    return result
+
+
+@impl(lib, "pad", "CompositeExplicitAutograd")  # type: ignore[misc]
+def pad_impl(
+    input: torch.Tensor,
+    pre_pad: list[int],
+    post_pad: list[int],
+    pad_value: int,
+) -> torch.Tensor:
+    rank = input.dim()
+    offset = 4 - rank
+    logical_pre = _pad_to_logical_order(pre_pad, input)
+    logical_post = _pad_to_logical_order(post_pad, input)
+
+    padding = []
+    for i in reversed(range(rank)):
+        padding.extend([logical_pre[offset + i], logical_post[offset + i]])
+    return F.pad(input, padding, mode="constant", value=pad_value)
+
+
+# ===================================================================
+# QUANTIZED CONV2D OPERATION DEFINITION
+# ===================================================================
+
+lib.define(
+    "quantized_conv2d("
+    "Tensor input, "
+    "Tensor weight, "
+    "Tensor? bias, "
+    "int[] stride, "
+    "int[] padding, "
+    "int[] dilation, "
+    "int input_offset, "
+    "int output_offset, "
+    "Tensor requantize_multipliers, "
+    "Tensor requantize_shifts, "
+    "int activation_min, "
+    "int activation_max, "
+    "Tensor scratch"
+    ") -> Tensor"
+)
+
+
+lib.define(
+    "quantized_conv2d.out("
+    "Tensor input, "
+    "Tensor weight, "
+    "Tensor? bias, "
+    "int[] stride, "
+    "int[] padding, "
+    "int[] dilation, "
+    "int input_offset, "
+    "int output_offset, "
+    "Tensor requantize_multipliers, "
+    "Tensor requantize_shifts, "
+    "int activation_min, "
+    "int activation_max, "
+    "Tensor scratch, "
+    "*, Tensor(a!) out"
+    ") -> Tensor(a!)"
+)
+
+
+def _conv2d_padding(
+    padding: Sequence[int],
+    input_shape: torch.Size,
+    weight_shape: torch.Size,
+    stride: Sequence[int],
+    dilation: Sequence[int],
+) -> tuple[int, int, int, int]:
+    # Four values mean top/left/bottom/right for supported SAME convolutions.
+    if len(padding) not in (2, 4) or any(p < 0 for p in padding):
+        raise ValueError(f"Expected 2 or 4 nonnegative padding values, got {padding}")
+    top, left = padding[:2]
+    bottom, right = padding[2:] if len(padding) == 4 else padding
+    if len(padding) == 4:
+        height, width = input_shape[2:4]
+        kernel = weight_shape[1:3]
+        if height == 1 and kernel[0] == 1:
+            raise ValueError("Four-value padding is unsupported for 1xN convolution")
+        total = [
+            max(((size + step - 1) // step - 1) * step + dil * (k - 1) + 1 - size, 0)
+            for size, step, dil, k in zip((height, width), stride, dilation, kernel)
+        ]
+        if (top, left, bottom, right) != (
+            total[0] // 2,
+            total[1] // 2,
+            total[0] - total[0] // 2,
+            total[1] - total[1] // 2,
+        ):
+            raise ValueError("Four-value padding must match SAME convolution geometry")
+    return top, left, bottom, right
+
+
+def _compute_conv2d_output_shape(
+    input_shape: torch.Size,
+    weight_shape: torch.Size,
+    stride: Sequence[int],
+    padding: Sequence[int],
+    dilation: Sequence[int],
+) -> torch.Size:
+    batch = input_shape[0]
+    in_height = input_shape[2]
+    in_width = input_shape[3]
+    # We store the weights in OHWI layout (out, kernel_h, kernel_w, in)
+    kernel_height = weight_shape[1]
+    kernel_width = weight_shape[2]
+
+    stride_h, stride_w = stride
+    pad_h, pad_w, pad_bottom, pad_right = _conv2d_padding(
+        padding, input_shape, weight_shape, stride, dilation
+    )
+    dilation_h, dilation_w = dilation
+
+    out_channels = weight_shape[0]
+    out_height = (
+        in_height + pad_h + pad_bottom - dilation_h * (kernel_height - 1) - 1
+    ) // stride_h + 1
+    out_width = (
+        in_width + pad_w + pad_right - dilation_w * (kernel_width - 1) - 1
+    ) // stride_w + 1
+    return torch.Size([batch, out_channels, out_height, out_width])
+
+
+def _compute_depthwise_conv2d_output_shape(
+    input_shape: torch.Size,
+    weight_shape: torch.Size,
+    stride: Sequence[int],
+    padding: Sequence[int],
+    dilation: Sequence[int],
+) -> torch.Size:
+    batch = input_shape[0]
+    in_height = input_shape[2]
+    in_width = input_shape[3]
+    # For depthwise conv, we store the weights in IHWO layout (1, kernel_h, kernel_w, out)
+    # where dimension 3 contains the output channels
+    kernel_height = weight_shape[1]
+    kernel_width = weight_shape[2]
+
+    stride_h, stride_w = stride
+    pad_h, pad_w, pad_bottom, pad_right = _conv2d_padding(
+        padding, input_shape, weight_shape, stride, dilation
+    )
+    dilation_h, dilation_w = dilation
+
+    out_channels = weight_shape[3]  # IHWO format: output channels at dimension 3
+    out_height = (
+        in_height + pad_h + pad_bottom - dilation_h * (kernel_height - 1) - 1
+    ) // stride_h + 1
+    out_width = (
+        in_width + pad_w + pad_right - dilation_w * (kernel_width - 1) - 1
+    ) // stride_w + 1
+    return torch.Size([batch, out_channels, out_height, out_width])
+
+
+@register_fake("cortex_m::quantized_conv2d")  # type: ignore[misc]
+def quantized_conv2d_meta(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    stride: Sequence[int],
+    padding: Sequence[int],
+    dilation: Sequence[int],
+    input_offset: int,
+    output_offset: int,
+    requantize_multipliers: torch.Tensor,
+    requantize_shifts: torch.Tensor,
+    activation_min: int,
+    activation_max: int,
+    scratch: torch.Tensor,
+) -> torch.Tensor:
+    stride_vals = list(stride)
+    padding_vals = list(padding)
+    dilation_vals = list(dilation)
+    output_shape = _compute_conv2d_output_shape(
+        input.shape, weight.shape, stride_vals, padding_vals, dilation_vals
+    )
+    return torch.empty(
+        output_shape,
+        dtype=torch.int8,
+        device=input.device,
+        memory_format=torch.channels_last,
+    )
+
+
+@impl(lib, "quantized_conv2d", "CompositeExplicitAutograd")  # type: ignore[misc]
+def quantized_conv2d_impl(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    stride: Sequence[int],
+    padding: Sequence[int],
+    dilation: Sequence[int],
+    input_offset: int,
+    output_offset: int,
+    requantize_multipliers: torch.Tensor,
+    requantize_shifts: torch.Tensor,
+    activation_min: int,
+    activation_max: int,
+    scratch: torch.Tensor,
+) -> torch.Tensor:
+    if input.dim() != 4 or weight.dim() != 4:
+        raise RuntimeError("quantized_conv2d expects 4D input and weight tensors")
+    # Convert to int32 for accumulation and apply offsets
+    input_int32 = input.to(torch.int32) + int(input_offset)
+    if len(padding) == 4:
+        top, left, bottom, right = _conv2d_padding(
+            padding, input.shape, weight.shape, stride, dilation
+        )
+        input_int32 = F.pad(input_int32, (left, right, top, bottom))
+        padding = (0, 0)
+    weight_int32 = weight.to(torch.int32)
+
+    if bias is None:
+        bias_int32 = torch.zeros(
+            weight.shape[0], dtype=torch.int32, device=input.device
+        )
+    else:
+        bias_int32 = bias.to(torch.int32)
+
+    input_channels = input.shape[1]
+    kernel_input_channels = weight.shape[3]
+    groups = input_channels // kernel_input_channels
+
+    # Convert weights back to OIHW layout expected by torch.nn.functional.conv2d
+    weight_oi_hw = weight_int32.permute(0, 3, 1, 2).contiguous()
+
+    conv_acc = F.conv2d(
+        input_int32,
+        weight_oi_hw,
+        bias_int32,
+        stride=tuple(stride),
+        padding=tuple(padding),
+        dilation=tuple(dilation),
+        groups=groups,
+    )
+
+    result_channels = []
+    for output_channel_i in range(conv_acc.shape[1]):
+        result_channel = requantize_cmsis(
+            conv_acc[:, output_channel_i, :, :],
+            int(requantize_multipliers[output_channel_i]),
+            int(requantize_shifts[output_channel_i]),
+        )
+        result_channels.append(result_channel)
+
+    result = torch.stack(result_channels, dim=1)
+
+    result += output_offset
+    result = torch.clamp(result, activation_min, activation_max)
+
+    # TODO - this enforces all convolution layers to result in channels last.
+    # This issue does comes min/mul layers from decomposition of hard swish
+    return result.to(torch.int8, memory_format=torch.channels_last)
+
+
+lib.define(
+    "quantized_conv2d_nhwc("
+    "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, "
+    "int[] dilation, int input_offset, int output_offset, "
+    "Tensor requantize_multipliers, Tensor requantize_shifts, "
+    "int activation_min, int activation_max, Tensor scratch) -> Tensor"
+)
+lib.define(
+    "quantized_conv2d_nhwc.out("
+    "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, "
+    "int[] dilation, int input_offset, int output_offset, "
+    "Tensor requantize_multipliers, Tensor requantize_shifts, "
+    "int activation_min, int activation_max, Tensor scratch, "
+    "*, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+@register_fake("cortex_m::quantized_conv2d_nhwc")  # type: ignore[misc]
+@experimental(_EXPLICIT_LAYOUT_EXPERIMENTAL)  # type: ignore[misc]
+def quantized_conv2d_nhwc_meta(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    stride: Sequence[int],
+    padding: Sequence[int],
+    dilation: Sequence[int],
+    input_offset: int,
+    output_offset: int,
+    requantize_multipliers: torch.Tensor,
+    requantize_shifts: torch.Tensor,
+    activation_min: int,
+    activation_max: int,
+    scratch: torch.Tensor,
+) -> torch.Tensor:
+    nchw = quantized_conv2d_meta(
+        input.permute(0, 3, 1, 2),
+        weight,
+        bias,
+        stride,
+        padding,
+        dilation,
+        input_offset,
+        output_offset,
+        requantize_multipliers,
+        requantize_shifts,
+        activation_min,
+        activation_max,
+        scratch,
+    )
+    return nchw.permute(0, 2, 3, 1).contiguous()
+
+
+@impl(lib, "quantized_conv2d_nhwc", "CompositeExplicitAutograd")  # type: ignore[misc]
+def quantized_conv2d_nhwc_impl(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    stride: Sequence[int],
+    padding: Sequence[int],
+    dilation: Sequence[int],
+    input_offset: int,
+    output_offset: int,
+    requantize_multipliers: torch.Tensor,
+    requantize_shifts: torch.Tensor,
+    activation_min: int,
+    activation_max: int,
+    scratch: torch.Tensor,
+) -> torch.Tensor:
+    nchw = quantized_conv2d_impl(
+        input.permute(0, 3, 1, 2).contiguous(),
+        weight,
+        bias,
+        stride,
+        padding,
+        dilation,
+        input_offset,
+        output_offset,
+        requantize_multipliers,
+        requantize_shifts,
+        activation_min,
+        activation_max,
+        scratch,
+    )
+    return nchw.permute(0, 2, 3, 1).contiguous()
+
+
+# ===================================================================
+# QUANTIZED DEPTHWISE CONV2D OPERATION DEFINITION
+# ===================================================================
+
+lib.define(
+    "quantized_depthwise_conv2d("
+    "Tensor input, "
+    "Tensor weight, "
+    "Tensor? bias, "
+    "int[] stride, "
+    "int[] padding, "
+    "int[] dilation, "
+    "int depth_multiplier, "
+    "int input_offset, "
+    "int output_offset, "
+    "Tensor requantize_multipliers, "
+    "Tensor requantize_shifts, "
+    "int activation_min, "
+    "int activation_max, "
+    "Tensor scratch"
+    ") -> Tensor"
+)
+
+
+lib.define(
+    "quantized_depthwise_conv2d.out("
+    "Tensor input, "
+    "Tensor weight, "
+    "Tensor? bias, "
+    "int[] stride, "
+    "int[] padding, "
+    "int[] dilation, "
+    "int depth_multiplier, "
+    "int input_offset, "
+    "int output_offset, "
+    "Tensor requantize_multipliers, "
+    "Tensor requantize_shifts, "
+    "int activation_min, "
+    "int activation_max, "
+    "Tensor scratch, "
+    "*, Tensor(a!) out"
+    ") -> Tensor(a!)"
+)
+
+
+@register_fake("cortex_m::quantized_depthwise_conv2d")  # type: ignore[misc]
+def quantized_depthwise_conv2d_meta(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    stride: Sequence[int],
+    padding: Sequence[int],
+    dilation: Sequence[int],
+    depth_multiplier: int,
+    input_offset: int,
+    output_offset: int,
+    requantize_multipliers: torch.Tensor,
+    requantize_shifts: torch.Tensor,
+    activation_min: int,
+    activation_max: int,
+    scratch: torch.Tensor,
+) -> torch.Tensor:
+    stride_vals = list(stride)
+    padding_vals = list(padding)
+    dilation_vals = list(dilation)
+    output_shape = _compute_depthwise_conv2d_output_shape(
+        input.shape, weight.shape, stride_vals, padding_vals, dilation_vals
+    )
+    return torch.empty(
+        output_shape,
+        dtype=torch.int8,
+        device=input.device,
+        memory_format=torch.channels_last,
+    )
+
+
+@impl(lib, "quantized_depthwise_conv2d", "CompositeExplicitAutograd")  # type: ignore[misc]
+def quantized_depthwise_conv2d_impl(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    stride: Sequence[int],
+    padding: Sequence[int],
+    dilation: Sequence[int],
+    depth_multiplier: int,
+    input_offset: int,
+    output_offset: int,
+    requantize_multipliers: torch.Tensor,
+    requantize_shifts: torch.Tensor,
+    activation_min: int,
+    activation_max: int,
+    scratch: torch.Tensor,
+) -> torch.Tensor:
+    if input.dim() != 4 or weight.dim() != 4:
+        raise RuntimeError(
+            "quantized_depthwise_conv2d expects 4D input and weight tensors"
+        )
+
+    input_channels = input.shape[1]
+    groups = input_channels
+
+    # Convert to int32 for accumulation and apply offsets
+    input_int32 = input.to(torch.int32) + int(input_offset)
+    if len(padding) == 4:
+        top, left, bottom, right = _conv2d_padding(
+            padding, input.shape, weight.shape, stride, dilation
+        )
+        input_int32 = F.pad(input_int32, (left, right, top, bottom))
+        padding = (0, 0)
+    weight_int32 = weight.to(torch.int32)
+
+    if bias is None:
+        bias_int32 = torch.zeros(
+            weight.shape[3],
+            dtype=torch.int32,
+            device=input.device,  # C_OUT is at dim 3 in IHWO
+        )
+    else:
+        bias_int32 = bias.to(torch.int32)
+
+    # Weight is in IHWO layout: [1, H, W, C_OUT]
+    # Convert to OIHW layout expected by torch.nn.functional.conv2d
+    # IHWO [1, H, W, C_OUT] -> OIHW [C_OUT, 1, H, W]
+    weight_oi_hw = weight_int32.permute(3, 0, 1, 2).contiguous()
+
+    # Depthwise convolution has groups == input_channels
+    conv_acc = F.conv2d(
+        input_int32,
+        weight_oi_hw,
+        bias_int32,
+        stride=tuple(stride),
+        padding=tuple(padding),
+        dilation=tuple(dilation),
+        groups=groups,
+    )
+
+    result_channels = []
+    for output_channel_i in range(conv_acc.shape[1]):
+        result_channel = requantize_cmsis(
+            conv_acc[:, output_channel_i, :, :],
+            int(requantize_multipliers[output_channel_i]),
+            int(requantize_shifts[output_channel_i]),
+        )
+        result_channels.append(result_channel)
+
+    result = torch.stack(result_channels, dim=1)
+
+    result += output_offset
+    result = torch.clamp(result, activation_min, activation_max)
+
+    return result.to(torch.int8, memory_format=torch.channels_last)
+
+
+lib.define(
+    "quantized_depthwise_conv2d_nhwc("
+    "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, "
+    "int[] dilation, int depth_multiplier, int input_offset, int output_offset, "
+    "Tensor requantize_multipliers, Tensor requantize_shifts, "
+    "int activation_min, int activation_max, Tensor scratch) -> Tensor"
+)
+lib.define(
+    "quantized_depthwise_conv2d_nhwc.out("
+    "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, "
+    "int[] dilation, int depth_multiplier, int input_offset, int output_offset, "
+    "Tensor requantize_multipliers, Tensor requantize_shifts, "
+    "int activation_min, int activation_max, Tensor scratch, "
+    "*, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+@register_fake("cortex_m::quantized_depthwise_conv2d_nhwc")  # type: ignore[misc]
+@experimental(_EXPLICIT_LAYOUT_EXPERIMENTAL)  # type: ignore[misc]
+def quantized_depthwise_conv2d_nhwc_meta(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    stride: Sequence[int],
+    padding: Sequence[int],
+    dilation: Sequence[int],
+    depth_multiplier: int,
+    input_offset: int,
+    output_offset: int,
+    requantize_multipliers: torch.Tensor,
+    requantize_shifts: torch.Tensor,
+    activation_min: int,
+    activation_max: int,
+    scratch: torch.Tensor,
+) -> torch.Tensor:
+    nchw = quantized_depthwise_conv2d_meta(
+        input.permute(0, 3, 1, 2),
+        weight,
+        bias,
+        stride,
+        padding,
+        dilation,
+        depth_multiplier,
+        input_offset,
+        output_offset,
+        requantize_multipliers,
+        requantize_shifts,
+        activation_min,
+        activation_max,
+        scratch,
+    )
+    return nchw.permute(0, 2, 3, 1).contiguous()
+
+
+@impl(lib, "quantized_depthwise_conv2d_nhwc", "CompositeExplicitAutograd")  # type: ignore[misc]
+def quantized_depthwise_conv2d_nhwc_impl(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    stride: Sequence[int],
+    padding: Sequence[int],
+    dilation: Sequence[int],
+    depth_multiplier: int,
+    input_offset: int,
+    output_offset: int,
+    requantize_multipliers: torch.Tensor,
+    requantize_shifts: torch.Tensor,
+    activation_min: int,
+    activation_max: int,
+    scratch: torch.Tensor,
+) -> torch.Tensor:
+    nchw = quantized_depthwise_conv2d_impl(
+        input.permute(0, 3, 1, 2).contiguous(),
+        weight,
+        bias,
+        stride,
+        padding,
+        dilation,
+        depth_multiplier,
+        input_offset,
+        output_offset,
+        requantize_multipliers,
+        requantize_shifts,
+        activation_min,
+        activation_max,
+        scratch,
+    )
+    return nchw.permute(0, 2, 3, 1).contiguous()
+
+
+# ===================================================================
+# QUANTIZED TRANSPOSE_CONV2D OPERATION DEFINITION
+# ===================================================================
+
+lib.define(
+    "quantized_transpose_conv2d("
+    "Tensor input, "
+    "Tensor weight, "
+    "Tensor? bias, "
+    "int[] stride, "
+    "int[] padding, "
+    "int[] output_padding, "
+    "int[] dilation, "
+    "int input_offset, "
+    "int output_offset, "
+    "Tensor requantize_multipliers, "
+    "Tensor requantize_shifts, "
+    "int activation_min, "
+    "int activation_max, "
+    "Tensor scratch, "
+    "Tensor output_scratch"
+    ") -> Tensor"
+)
+
+lib.define(
+    "quantized_transpose_conv2d.out("
+    "Tensor input, "
+    "Tensor weight, "
+    "Tensor? bias, "
+    "int[] stride, "
+    "int[] padding, "
+    "int[] output_padding, "
+    "int[] dilation, "
+    "int input_offset, "
+    "int output_offset, "
+    "Tensor requantize_multipliers, "
+    "Tensor requantize_shifts, "
+    "int activation_min, "
+    "int activation_max, "
+    "Tensor scratch, "
+    "Tensor output_scratch, "
+    "*, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+def _compute_conv_transpose2d_output_shape(
+    input_shape: torch.Size,
+    weight_shape: torch.Size,
+    stride: list[int],
+    padding: list[int],
+    output_padding: list[int],
+    dilation: list[int],
+) -> torch.Size:
+    """
+    Compute output shape for transposed 2D convolution.
+
+    Formula for each dimension:
+    out_size = (in_size - 1) * stride - 2 * padding + dilation * (kernel_size - 1) + output_padding + 1
+    """
+    batch = input_shape[0]
+    in_height = input_shape[2]
+    in_width = input_shape[3]
+
+    # Weight is in OHWI format after permutation
+    out_channels = weight_shape[0]
+    kernel_height = weight_shape[1]
+    kernel_width = weight_shape[2]
+
+    stride_h, stride_w = stride
+    pad_h, pad_w = padding
+    out_pad_h, out_pad_w = output_padding
+    dilation_h, dilation_w = dilation
+
+    out_height = (
+        (in_height - 1) * stride_h
+        - 2 * pad_h
+        + dilation_h * (kernel_height - 1)
+        + out_pad_h
+        + 1
+    )
+    out_width = (
+        (in_width - 1) * stride_w
+        - 2 * pad_w
+        + dilation_w * (kernel_width - 1)
+        + out_pad_w
+        + 1
+    )
+
+    return torch.Size([batch, out_channels, out_height, out_width])
+
+
+@register_fake("cortex_m::quantized_transpose_conv2d")  # type: ignore[misc]
+def quantized_transpose_conv2d_meta(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    stride: Sequence[int],
+    padding: Sequence[int],
+    output_padding: Sequence[int],
+    dilation: Sequence[int],
+    input_offset: int,
+    output_offset: int,
+    requantize_multipliers: torch.Tensor,
+    requantize_shifts: torch.Tensor,
+    activation_min: int,
+    activation_max: int,
+    scratch: torch.Tensor,
+    output_scratch: torch.Tensor,
+) -> torch.Tensor:
+    stride_vals = list(stride)
+    padding_vals = list(padding)
+    output_padding_vals = list(output_padding)
+    dilation_vals = list(dilation)
+
+    output_shape = _compute_conv_transpose2d_output_shape(
+        input.shape,
+        weight.shape,
+        stride_vals,
+        padding_vals,
+        output_padding_vals,
+        dilation_vals,
+    )
+
+    return torch.empty(
+        output_shape,
+        dtype=torch.int8,
+        device=input.device,
+        memory_format=torch.channels_last,
+    )
+
+
+@impl(lib, "quantized_transpose_conv2d", "CompositeExplicitAutograd")  # type: ignore[misc]
+def quantized_transpose_conv2d_impl(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    stride: Sequence[int],
+    padding: Sequence[int],
+    output_padding: Sequence[int],
+    dilation: Sequence[int],
+    input_offset: int,
+    output_offset: int,
+    requantize_multipliers: torch.Tensor,
+    requantize_shifts: torch.Tensor,
+    activation_min: int,
+    activation_max: int,
+    scratch: torch.Tensor,
+    output_scratch: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Reference implementation of quantized transposed convolution.
+    Simulates CMSIS-NN behavior for testing.
+    """
+    if input.dim() != 4 or weight.dim() != 4:
+        raise RuntimeError(
+            "quantized_conv_transpose2d expects 4D input and weight tensors"
+        )
+
+    # Convert to int32 for accumulation and apply offsets
+    input_int32 = input.to(torch.int32) + int(input_offset)
+    weight_int32 = weight.to(torch.int32)
+
+    if bias is None:
+        bias_int32 = torch.zeros(
+            weight.shape[0], dtype=torch.int32, device=input.device
+        )
+    else:
+        bias_int32 = bias.to(torch.int32)
+
+    input_channels = input.shape[1]
+    kernel_input_channels = weight.shape[3]
+    groups = input_channels // kernel_input_channels
+
+    # Convert weights from OHWI to IOHW layout for torch.nn.functional.conv_transpose2d
+    # Weight is in OHWI (out_channels, H, W, in_channels)
+    # F.conv_transpose2d expects IOHW (in_channels, out_channels, H, W)
+    weight_iohw = weight_int32.permute(3, 0, 1, 2).contiguous()
+
+    # PyTorch doesn't support int32 for conv_transpose2d, so convert to float
+    # Perform transposed convolution with float tensors
+    conv_transpose_acc = F.conv_transpose2d(
+        input_int32.to(torch.float32),
+        weight_iohw.to(torch.float32),
+        bias_int32.to(torch.float32),
+        stride=tuple(stride),
+        padding=tuple(padding),
+        output_padding=tuple(output_padding),
+        dilation=tuple(dilation),
+        groups=groups,
+    ).to(torch.int32)
+
+    # Apply per-channel requantization
+    result_channels = []
+    for output_channel_i in range(conv_transpose_acc.shape[1]):
+        result_channel = requantize_cmsis(
+            conv_transpose_acc[:, output_channel_i, :, :],
+            int(requantize_multipliers[output_channel_i]),
+            int(requantize_shifts[output_channel_i]),
+        )
+        result_channels.append(result_channel)
+
+    result = torch.stack(result_channels, dim=1)
+    result += output_offset
+    result = result.clamp(activation_min, activation_max)
+
+    return result.to(torch.int8).to(memory_format=torch.channels_last)
+
+
+lib.define(
+    "quantized_transpose_conv2d_nhwc("
+    "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, "
+    "int[] output_padding, int[] dilation, int input_offset, int output_offset, "
+    "Tensor requantize_multipliers, Tensor requantize_shifts, "
+    "int activation_min, int activation_max, Tensor scratch, "
+    "Tensor output_scratch) -> Tensor"
+)
+lib.define(
+    "quantized_transpose_conv2d_nhwc.out("
+    "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, "
+    "int[] output_padding, int[] dilation, int input_offset, int output_offset, "
+    "Tensor requantize_multipliers, Tensor requantize_shifts, "
+    "int activation_min, int activation_max, Tensor scratch, "
+    "Tensor output_scratch, *, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+@register_fake("cortex_m::quantized_transpose_conv2d_nhwc")  # type: ignore[misc]
+@experimental(_EXPLICIT_LAYOUT_EXPERIMENTAL)  # type: ignore[misc]
+def quantized_transpose_conv2d_nhwc_meta(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    stride: Sequence[int],
+    padding: Sequence[int],
+    output_padding: Sequence[int],
+    dilation: Sequence[int],
+    input_offset: int,
+    output_offset: int,
+    requantize_multipliers: torch.Tensor,
+    requantize_shifts: torch.Tensor,
+    activation_min: int,
+    activation_max: int,
+    scratch: torch.Tensor,
+    output_scratch: torch.Tensor,
+) -> torch.Tensor:
+    nchw = quantized_transpose_conv2d_meta(
+        input.permute(0, 3, 1, 2),
+        weight,
+        bias,
+        stride,
+        padding,
+        output_padding,
+        dilation,
+        input_offset,
+        output_offset,
+        requantize_multipliers,
+        requantize_shifts,
+        activation_min,
+        activation_max,
+        scratch,
+        output_scratch,
+    )
+    return nchw.permute(0, 2, 3, 1).contiguous()
+
+
+@impl(lib, "quantized_transpose_conv2d_nhwc", "CompositeExplicitAutograd")  # type: ignore[misc]
+def quantized_transpose_conv2d_nhwc_impl(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    stride: Sequence[int],
+    padding: Sequence[int],
+    output_padding: Sequence[int],
+    dilation: Sequence[int],
+    input_offset: int,
+    output_offset: int,
+    requantize_multipliers: torch.Tensor,
+    requantize_shifts: torch.Tensor,
+    activation_min: int,
+    activation_max: int,
+    scratch: torch.Tensor,
+    output_scratch: torch.Tensor,
+) -> torch.Tensor:
+    nchw = quantized_transpose_conv2d_impl(
+        input.permute(0, 3, 1, 2).contiguous(),
+        weight,
+        bias,
+        stride,
+        padding,
+        output_padding,
+        dilation,
+        input_offset,
+        output_offset,
+        requantize_multipliers,
+        requantize_shifts,
+        activation_min,
+        activation_max,
+        scratch,
+        output_scratch,
+    )
+    return nchw.permute(0, 2, 3, 1).contiguous()
+
+
+# ===================================================================
+# QUANTIZED AVG_POOL2D OPERATION DEFINITION
+# ===================================================================
+
+lib.define(
+    "quantized_avg_pool2d("
+    "Tensor input, "
+    "int[] kernel_size, "
+    "int[] stride, "
+    "int[] padding, "
+    "bool ceil_mode, "
+    "int zero_point, "
+    "int multiplier, "
+    "int shift, "
+    "Tensor scratch"
+    ") -> Tensor"
+)
+lib.define(
+    "quantized_avg_pool2d.out("
+    "Tensor input, "
+    "int[] kernel_size, "
+    "int[] stride, "
+    "int[] padding, "
+    "bool ceil_mode, "
+    "int zero_point, "
+    "int multiplier, "
+    "int shift, "
+    "Tensor scratch, "
+    "*, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+@register_fake("cortex_m::quantized_avg_pool2d")  # type: ignore[misc]
+def quantized_avg_pool2d_meta(
+    input: torch.Tensor,
+    kernel_size: Sequence[int],
+    stride: Sequence[int],
+    padding: Sequence[int],
+    ceil_mode: bool,
+    zero_point: int,
+    multiplier: int,
+    shift: int,
+    scratch: torch.Tensor,
+) -> torch.Tensor:
+    kernel = _ensure_tuple2(kernel_size)
+    stride_vals = _ensure_tuple2(stride)
+    padding_vals = _ensure_tuple2(padding)
+    output = F.avg_pool2d(
+        input.to(torch.float),
+        kernel,
+        stride=stride_vals,
+        padding=padding_vals,
+        ceil_mode=ceil_mode,
+        count_include_pad=False,
+    )
+    return torch.empty(
+        output.shape,
+        dtype=torch.int8,
+        device=input.device,
+        memory_format=torch.channels_last,
+    )
+
+
+@impl(lib, "quantized_avg_pool2d", "CompositeExplicitAutograd")  # type: ignore[misc]
+def quantized_avg_pool2d_impl(
+    input: torch.Tensor,
+    kernel_size: Sequence[int],
+    stride: Sequence[int],
+    padding: Sequence[int],
+    ceil_mode: bool,
+    zero_point: int,
+    multiplier: int,
+    shift: int,
+    scratch: torch.Tensor,
+) -> torch.Tensor:
+    dequant_input = dequantize_per_tensor_cmsis(input, zero_point, multiplier, shift)
+
+    kernel = _ensure_tuple2(kernel_size)
+    stride_vals = _ensure_tuple2(stride)
+    padding_vals = _ensure_tuple2(padding)
+
+    # TODO: implement dilation != 1.
+    result = F.avg_pool2d(
+        dequant_input,
+        kernel,
+        stride=stride_vals,
+        padding=padding_vals,
+        ceil_mode=ceil_mode,
+        count_include_pad=False,
+    )
+    result = quantize_per_tensor_cmsis(result, zero_point, multiplier, shift)
+    output = torch.clamp(result, -128, 127)
+    return output.to(torch.int8)
+
+
+lib.define(
+    "quantized_avg_pool2d_nhwc("
+    "Tensor input, int[] kernel_size, int[] stride, int[] padding, "
+    "bool ceil_mode, int zero_point, int multiplier, int shift, "
+    "Tensor scratch) -> Tensor"
+)
+lib.define(
+    "quantized_avg_pool2d_nhwc.out("
+    "Tensor input, int[] kernel_size, int[] stride, int[] padding, "
+    "bool ceil_mode, int zero_point, int multiplier, int shift, "
+    "Tensor scratch, *, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+@register_fake("cortex_m::quantized_avg_pool2d_nhwc")  # type: ignore[misc]
+@experimental(_EXPLICIT_LAYOUT_EXPERIMENTAL)  # type: ignore[misc]
+def quantized_avg_pool2d_nhwc_meta(
+    input: torch.Tensor,
+    kernel_size: Sequence[int],
+    stride: Sequence[int],
+    padding: Sequence[int],
+    ceil_mode: bool,
+    zero_point: int,
+    multiplier: int,
+    shift: int,
+    scratch: torch.Tensor,
+) -> torch.Tensor:
+    nchw = quantized_avg_pool2d_meta(
+        input.permute(0, 3, 1, 2),
+        kernel_size,
+        stride,
+        padding,
+        ceil_mode,
+        zero_point,
+        multiplier,
+        shift,
+        scratch,
+    )
+    return nchw.permute(0, 2, 3, 1).contiguous()
+
+
+@impl(lib, "quantized_avg_pool2d_nhwc", "CompositeExplicitAutograd")  # type: ignore[misc]
+def quantized_avg_pool2d_nhwc_impl(
+    input: torch.Tensor,
+    kernel_size: Sequence[int],
+    stride: Sequence[int],
+    padding: Sequence[int],
+    ceil_mode: bool,
+    zero_point: int,
+    multiplier: int,
+    shift: int,
+    scratch: torch.Tensor,
+) -> torch.Tensor:
+    nchw = quantized_avg_pool2d_impl(
+        input.permute(0, 3, 1, 2).contiguous(),
+        kernel_size,
+        stride,
+        padding,
+        ceil_mode,
+        zero_point,
+        multiplier,
+        shift,
+        scratch,
+    )
+    return nchw.permute(0, 2, 3, 1).contiguous()
+
+
+# ===================================================================
+# QUANTIZED MAX POOL2D OPERATION DEFINITION
+# ===================================================================
+
+lib.define(
+    "quantized_max_pool2d("
+    "Tensor input, "
+    "int[] kernel_size, "
+    "int[] stride, "
+    "int[] padding, "
+    "int[] dilation, "
+    "bool ceil_mode, "
+    "int input_zero_point, "
+    "int output_zero_point, "
+    "int activation_min, "
+    "int activation_max"
+    ") -> Tensor"
+)
+
+lib.define(
+    "quantized_max_pool2d.out("
+    "Tensor input, "
+    "int[] kernel_size, "
+    "int[] stride, "
+    "int[] padding, "
+    "int[] dilation, "
+    "bool ceil_mode, "
+    "int input_zero_point, "
+    "int output_zero_point, "
+    "int activation_min, "
+    "int activation_max, "
+    "*, Tensor(a!) out"
+    ") -> Tensor(a!)"
+)
+
+
+def _ensure_tuple2(value: Sequence[int]) -> tuple[int, int]:
+    if len(value) == 1:
+        return (int(value[0]), int(value[0]))
+    if len(value) != 2:
+        raise RuntimeError(f"Expected length-2 sequence, got {value}")
+    return (int(value[0]), int(value[1]))
+
+
+def _compute_max_pool2d_output_shape(
+    input_shape: torch.Size,
+    kernel_size: Sequence[int],
+    stride: Sequence[int],
+    padding: Sequence[int],
+    dilation: Sequence[int],
+) -> torch.Size:
+    batch = input_shape[0]
+    channels = input_shape[1]
+    in_height = input_shape[2]
+    in_width = input_shape[3]
+
+    kernel_height, kernel_width = kernel_size
+    stride_h, stride_w = stride
+    pad_h, pad_w = padding
+    dilation_h, dilation_w = dilation
+
+    out_height = (
+        in_height + 2 * pad_h - dilation_h * (kernel_height - 1) - 1
+    ) // stride_h + 1
+    out_width = (
+        in_width + 2 * pad_w - dilation_w * (kernel_width - 1) - 1
+    ) // stride_w + 1
+    return torch.Size([batch, channels, out_height, out_width])
+
+
+@register_fake("cortex_m::quantized_max_pool2d")  # type: ignore[misc]
+def quantized_max_pool2d_meta(
+    input: torch.Tensor,
+    kernel_size: Sequence[int],
+    stride: Sequence[int],
+    padding: Sequence[int],
+    dilation: Sequence[int],
+    ceil_mode: bool,
+    input_zero_point: int,
+    output_zero_point: int,
+    activation_min: int,
+    activation_max: int,
+) -> torch.Tensor:
+    kernel = _ensure_tuple2(kernel_size)
+    stride_vals = _ensure_tuple2(stride)
+    padding_vals = _ensure_tuple2(padding)
+    dilation_vals = _ensure_tuple2(dilation)
+
+    output_shape = _compute_max_pool2d_output_shape(
+        input.shape, kernel, stride_vals, padding_vals, dilation_vals
+    )
+    return torch.empty(
+        output_shape,
+        dtype=torch.int8,
+        device=input.device,
+        memory_format=torch.channels_last,
+    )
+
+
+@impl(lib, "quantized_max_pool2d", "CompositeExplicitAutograd")  # type: ignore[misc]
+def quantized_max_pool2d_impl(
+    input: torch.Tensor,
+    kernel_size: Sequence[int],
+    stride: Sequence[int],
+    padding: Sequence[int],
+    dilation: Sequence[int],
+    ceil_mode: bool,
+    input_zero_point: int,
+    output_zero_point: int,
+    activation_min: int,
+    activation_max: int,
+) -> torch.Tensor:
+    if input.dim() != 4:
+        raise RuntimeError("quantized_max_pool2d expects 4D input tensor")
+
+    if input_zero_point != output_zero_point:
+        raise RuntimeError(
+            "quantized_max_pool2d expects matching input/output zero points"
+        )
+
+    kernel = _ensure_tuple2(kernel_size)
+    stride_vals = _ensure_tuple2(stride)
+    padding_vals = _ensure_tuple2(padding)
+    dilation_vals = _ensure_tuple2(dilation)
+
+    if dilation_vals != (1, 1):
+        raise RuntimeError("quantized_max_pool2d only supports dilation == 1")
+    if ceil_mode:
+        raise RuntimeError("quantized_max_pool2d does not support ceil_mode=True")
+
+    # aten's channels-last max_pool2d caps how large an image it will take, so
+    # pool a contiguous copy instead. Pooling is layout-invariant, and the
+    # return below puts the result back in channels-last either way.
+    #
+    # The cap: cpu_max_pool_channels_last buffers each window index in
+    # vec::int_same_size_t<opmath_t> and guards it with
+    # TORCH_CHECK(input_depth * input_height * input_width <= max), so int8
+    # rejects any image with more than 127 spatial elements -- H*W, with
+    # channels not counted. int16 hits the same wall at 32767, which a future
+    # quantized_max_pool2d_s16 will need to handle the same way.
+    #
+    # .to(memory_format=...) rather than .contiguous(): for C == 1 the
+    # channels-last strides also satisfy plain contiguity, so .contiguous()
+    # returns the same tensor while aten still dispatches on the memory-format
+    # hint and raises anyway.
+    result = F.max_pool2d(
+        input.to(memory_format=torch.contiguous_format),
+        kernel,
+        stride=stride_vals,
+        padding=padding_vals,
+        dilation=dilation_vals,
+        ceil_mode=ceil_mode,
+    )
+    result = torch.clamp(result, activation_min, activation_max)
+    return result.to(torch.int8).contiguous(memory_format=torch.channels_last)
+
+
+lib.define(
+    "quantized_max_pool2d_nhwc("
+    "Tensor input, int[] kernel_size, int[] stride, int[] padding, "
+    "int[] dilation, bool ceil_mode, int input_zero_point, "
+    "int output_zero_point, int activation_min, int activation_max) -> Tensor"
+)
+lib.define(
+    "quantized_max_pool2d_nhwc.out("
+    "Tensor input, int[] kernel_size, int[] stride, int[] padding, "
+    "int[] dilation, bool ceil_mode, int input_zero_point, "
+    "int output_zero_point, int activation_min, int activation_max, "
+    "*, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+@register_fake("cortex_m::quantized_max_pool2d_nhwc")  # type: ignore[misc]
+@experimental(_EXPLICIT_LAYOUT_EXPERIMENTAL)  # type: ignore[misc]
+def quantized_max_pool2d_nhwc_meta(
+    input: torch.Tensor,
+    kernel_size: Sequence[int],
+    stride: Sequence[int],
+    padding: Sequence[int],
+    dilation: Sequence[int],
+    ceil_mode: bool,
+    input_zero_point: int,
+    output_zero_point: int,
+    activation_min: int,
+    activation_max: int,
+) -> torch.Tensor:
+    nchw = quantized_max_pool2d_meta(
+        input.permute(0, 3, 1, 2),
+        kernel_size,
+        stride,
+        padding,
+        dilation,
+        ceil_mode,
+        input_zero_point,
+        output_zero_point,
+        activation_min,
+        activation_max,
+    )
+    return nchw.permute(0, 2, 3, 1).contiguous()
+
+
+@impl(lib, "quantized_max_pool2d_nhwc", "CompositeExplicitAutograd")  # type: ignore[misc]
+def quantized_max_pool2d_nhwc_impl(
+    input: torch.Tensor,
+    kernel_size: Sequence[int],
+    stride: Sequence[int],
+    padding: Sequence[int],
+    dilation: Sequence[int],
+    ceil_mode: bool,
+    input_zero_point: int,
+    output_zero_point: int,
+    activation_min: int,
+    activation_max: int,
+) -> torch.Tensor:
+    nchw = quantized_max_pool2d_impl(
+        input.permute(0, 3, 1, 2).contiguous(),
+        kernel_size,
+        stride,
+        padding,
+        dilation,
+        ceil_mode,
+        input_zero_point,
+        output_zero_point,
+        activation_min,
+        activation_max,
+    )
+    return nchw.permute(0, 2, 3, 1).contiguous()