Function quantized_conv2d_impl

backends/cortex_m/ops/operators.py:752–812

(
    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,
)

@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,
) -> 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)
    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)