# Copyright (C) 2026 Embedl AB
"""Lattice-conforming convolution modules.
Lattice hardware accelerators support a restricted set of convolution
parameters. ``LatticeConv2d`` subclasses ``Conv2d`` and snaps an arbitrary
source convolution to the closest configuration allowed by the hardware.
"""
from typing import cast
import torch
import torch.nn.functional as F
from torch import nn
from embedl_deploy._internal.core.modules import FusedModule
from embedl_deploy._internal.core.quantize.config import Precision
from embedl_deploy._internal.core.quantize.stubs import WeightFakeQuantize
from embedl_deploy._internal.lattice.modules.activation import (
LatticeActivationLike,
)
from embedl_deploy._internal.lattice.modules.quant import LatticeQuant
def _as_pair(value: int | tuple[int, ...]) -> tuple[int, int]:
"""Return `value` as a 2-tuple, broadcasting scalars.
:param value:
Value to convert. If an integer, it is broadcast to both
spatial dimensions; if a tuple, it must have two entries.
:returns:
A 2-tuple of integers.
"""
if isinstance(value, tuple):
return (value[0], value[1])
return (value, value) # pragma: no cover
def _snap(value: int, allowed: tuple[int, ...]) -> int:
"""Return the entry in `allowed` closest to `value`.
:param value:
Value to snap.
:param allowed:
Allowed values to snap to.
:returns:
The entry in `allowed` closest to `value`. In case of ties, the
smaller value is preferred.
"""
return min(allowed, key=lambda v: (abs(v - value), v))
def _conv_weight_forward(
conv: nn.Conv2d,
weight_fake_quant: WeightFakeQuantize,
x: torch.Tensor,
) -> torch.Tensor:
"""Run a convolution, fake-quantizing the weight when enabled."""
# Skip fake-quantize during ONNX export: torch.fake_quantize_* would be
# lowered to QuantizeLinear/DequantizeLinear, which are incorrect for the
# Lattice ONNX format. The toolchain quantizes weights independently.
weight = (
conv.weight
if torch.onnx.is_in_onnx_export()
else weight_fake_quant(conv.weight)
)
# pylint: disable-next=not-callable
return F.conv2d(
x,
weight,
conv.bias,
conv.stride,
conv.padding,
conv.dilation,
conv.groups,
)
[docs]
class LatticeConv2d(nn.Conv2d):
"""``Conv2d`` snapped to Lattice's supported set.
Lattice hardware accepts only ``1×1`` and ``3×3`` convolutions with stride
1 or 2 (and stride 1 is mandatory for the ``1×1`` kernel). The constructor
takes an arbitrary source :class:`~torch.nn.Conv2d` and forwards its
``in_channels``, ``out_channels``, ``dilation``, ``groups``, and bias
presence; kernel size and stride are snapped to the nearest support values.
A ``1×1`` kernel is promoted to ``3×3`` whenever its stride exceeds
1. Padding is set to ``kernel_size // 2`` on each spatial axis to
preserve the output shape under the common ``"same"``-style convention used
by ResNet-family stems and downsamples.
Weights and bias are copied from the source convolution whenever the
snapped weight tensor has the same shape as the source's (i.e., only stride
and/or padding changed). If the kernel size was snapped — and the weight
tensor shape therefore changed — the instance keeps freshly initialized
weights, since there is no well-defined way to reuse the original kernel
values.
"""
#: Permitted spatial kernel sizes.
KERNEL_SIZES: tuple[int, ...] = (1, 3)
#: Permitted spatial strides.
STRIDES: tuple[int, ...] = (1, 2)
[docs]
@classmethod
def snapped_params(
cls, conv: nn.Conv2d
) -> tuple[tuple[int, int], tuple[int, int], tuple[int, int]]:
"""Return ``(kernel_size, stride, padding)`` after Lattice snapping.
:param conv:
Source convolution whose parameters are snapped to the
nearest values accepted by Lattice hardware.
:returns:
A three-tuple ``(kernel_size, stride, padding)`` of the snapped
parameters where each element is itself an ``(h, w)`` pair.
"""
kh, kw = _as_pair(conv.kernel_size)
sh, sw = _as_pair(conv.stride)
new_sh = _snap(sh, cls.STRIDES)
new_sw = _snap(sw, cls.STRIDES)
new_kh = _snap(kh, cls.KERNEL_SIZES)
new_kw = _snap(kw, cls.KERNEL_SIZES)
if new_kh == 1 and new_sh != 1:
new_kh = 3
if new_kw == 1 and new_sw != 1:
new_kw = 3
return (
(new_kh, new_kw),
(new_sh, new_sw),
(new_kh // 2, new_kw // 2),
)
[docs]
@classmethod
def is_compatible(cls, conv: nn.Conv2d) -> bool:
"""Return ``True`` when `conv` already matches Lattice's supported set.
A convolution is compatible when its kernel size, stride, and padding
equal what ``snapped_params`` would return for it.
:param conv:
Convolution to check.
:returns:
``True`` when `conv` already conforms to Lattice
constraints; ``False`` otherwise.
"""
kernel_size, stride, padding = cls.snapped_params(conv)
return (
_as_pair(conv.kernel_size) == kernel_size
and _as_pair(conv.stride) == stride
and _as_pair(conv.padding) == padding # type: ignore[arg-type]
)
def __init__(self, conv: nn.Conv2d) -> None:
"""Create a ``LatticeConv2d`` from an arbitrary ``Conv2d``.
:param conv:
Source convolution. Its ``in_channels``,
``out_channels``, ``dilation``, ``groups``, and bias
presence are forwarded unchanged; kernel size, stride,
and padding are snapped to Lattice's supported set.
"""
kernel_size, stride, padding = self.snapped_params(conv)
super().__init__(
conv.in_channels,
conv.out_channels,
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=_as_pair(conv.dilation),
groups=conv.groups,
bias=conv.bias is not None,
)
# Move to the same device/dtype as the source conv before copying.
self.to(device=conv.weight.device, dtype=conv.weight.dtype)
# Preserve weights/bias when the snapped weight tensor has the
# same shape as the source's (i.e., only stride/padding
# changed). If the kernel size was snapped, weight shapes
# differ and we cannot meaningfully reuse the source kernel.
if self.weight.shape == conv.weight.shape:
with torch.no_grad():
self.weight.copy_(conv.weight)
if conv.bias is not None and self.bias is not None:
self.bias.copy_(conv.bias)
[docs]
class LatticeConv2dAdvanced(LatticeConv2d):
"""``Conv2d`` snapped to Lattice's advanced-kernel set.
Lattice advanced hardware supports ``1×1``, ``3×3``, ``5×5``, and ``7×7``
convolutions. ``3×3`` allows stride 1 or 2; all other kernel sizes require
stride 1. Padding is always ``kernel_size // 2``.
Like :class:`LatticeConv2d`, weights and bias are preserved when only
stride and/or padding changed, and freshly initialized when the kernel size
was snapped.
"""
#: Permitted spatial kernel sizes.
KERNEL_SIZES: tuple[int, ...] = (1, 3, 5, 7)
[docs]
@classmethod
def snapped_params(
cls, conv: nn.Conv2d
) -> tuple[tuple[int, int], tuple[int, int], tuple[int, int]]:
"""Return ``(kernel_size, stride, padding)`` after snapping.
Kernel size is snapped to the nearest square kernel from the supported
set (``1×1``, ``3×3``, ``5×5``, ``7×7``) using the larger spatial
dimension. Only ``3×3`` allows stride 2; all other kernel sizes are
restricted to stride 1.
"""
kh, kw = _as_pair(conv.kernel_size)
sh, sw = _as_pair(conv.stride)
new_k = _snap(max(kh, kw), cls.KERNEL_SIZES)
# Only 3×3 allows stride 2; all other kernels require stride 1.
allowed_strides = (1, 2) if new_k == 3 else (1,)
new_s = _snap(max(sh, sw), allowed_strides)
return (
(new_k, new_k),
(new_s, new_s),
(new_k // 2, new_k // 2),
)
def _validate_conv_params(
conv: nn.Conv2d,
name: str,
allowed_kernels: tuple[int, ...],
allowed_strides: tuple[int, ...],
) -> None:
"""Shared validation for Lattice CBSR convolution parameters."""
kh, kw = _as_pair(conv.kernel_size)
sh, sw = _as_pair(conv.stride)
dh, dw = _as_pair(conv.dilation)
if kh != kw:
raise ValueError(
f"{name}: unsupported non-square kernel_size ({kh}, {kw}). "
f"Supported kernels are square with sizes: {allowed_kernels}."
)
if sh != sw:
raise ValueError(
f"{name}: unsupported non-square stride ({sh}, {sw}). "
f"Supported strides are square with sizes: {allowed_strides}."
)
if kh not in allowed_kernels or kw not in allowed_kernels:
raise ValueError(
f"{name}: unsupported kernel_size ({kh}, {kw}). "
f"Supported sizes: {allowed_kernels}."
)
if sh not in allowed_strides or sw not in allowed_strides:
raise ValueError(
f"{name}: unsupported stride ({sh}, {sw}). "
f"Supported strides: {allowed_strides}."
)
if kh != 3 and sh != 1:
raise ValueError(
f"{name}: kernel {kh}x{kw} requires stride 1, got ({sh}, {sw})."
)
if kw != 3 and sw != 1: # pragma: no cover — redundant with kh check
raise ValueError(
f"{name}: kernel {kh}x{kw} requires stride 1, got ({sh}, {sw})."
)
if dh != 1 or dw != 1:
raise ValueError(
f"{name}: unsupported dilation ({dh}, {dw}). "
f"Only dilation=1 is supported."
)
def validate_conv(conv: nn.Conv2d) -> None:
"""Raise ``ValueError`` when `conv` is outside the supported set.
Lattice CBSR blocks support only ``1×1`` and ``3×3`` kernels with stride 1
or 2, and dilation must be 1. A ``1×1`` kernel with stride != 1 is not
allowed.
"""
_validate_conv_params(
conv, "LatticeCBSR", LatticeConv2d.KERNEL_SIZES, LatticeConv2d.STRIDES
)
[docs]
class LatticeCBSR(FusedModule):
"""Fused ``Conv2d -> [BatchNorm2d] -> [Activation]`` for Lattice hardware.
The Lattice accelerator implements convolution, optional batch
normalization, and optional activation as a single fused CBSR block. When
an activation is present it must be ``ReLU`` or ``LatticeLeakyReLU``
(negative slope fixed at 1/16). This module validates that the convolution
parameters are within the hardware's supported set and raises
:class:`ValueError` otherwise.
Supported convolution parameters:
- Kernel sizes: ``1×1``, ``3×3``
- Strides: ``1``, ``2`` (stride > 1 requires ``3×3`` kernel)
- Dilation: ``1`` only
:param conv:
The convolution layer to fuse.
:param bn:
Optional batch normalization layer.
:param act:
Optional ``ReLU`` or ``LatticeLeakyReLU`` activation.
:param output_precision:
Precision for the output quantizer.
:param output_fixed_calibration:
Optional fixed ``(scale, zero_point)`` pair for the output
quantizer.
:raises ValueError:
If `conv` has unsupported kernel size, stride, or dilation.
"""
inputs_to_quantize: set[int] = set()
def __init__(
self,
conv: LatticeConv2d,
bn: nn.BatchNorm2d | None = None,
act: LatticeActivationLike | None = None,
*,
output_precision: Precision = Precision.INT8,
output_fixed_calibration: tuple[float, int] | None = None,
) -> None:
super().__init__()
self._validate_conv(conv)
self.conv = conv
self.bn = bn
self.act = act
self.weight_fake_quant = WeightFakeQuantize({self})
self.output_quant_stub = LatticeQuant(
{self},
precision=output_precision,
fixed_calibration=output_fixed_calibration,
)
def _validate_conv(self, conv: nn.Conv2d) -> None:
validate_conv(conv)
@property
def quantized_weight(self) -> torch.Tensor | None:
return self.conv.weight
[docs]
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Apply ``conv -> [bn] -> [act]``."""
x = _conv_weight_forward(self.conv, self.weight_fake_quant, x)
if self.bn is not None:
x = cast(torch.Tensor, self.bn(x))
if self.act is not None:
x = cast(torch.Tensor, self.act(x))
return x
def __repr__(self) -> str:
bn_info = ""
if self.bn is not None:
bn_info = f", bn={self.bn.num_features} (foldable)"
return (
f"{type(self).__name__}("
f"{self.conv.in_channels}->{self.conv.out_channels}, "
f"k={self.conv.kernel_size}, s={self.conv.stride}"
f"{bn_info})"
)
def validate_conv_advanced(conv: nn.Conv2d) -> None:
"""Raise ``ValueError`` when `conv` is outside the advanced set.
Lattice advanced CBSR blocks support ``1×1``, ``3×3``, ``5×5``, and ``7×7``
kernels. Only ``3×3`` allows stride 2; all other kernel sizes require
stride 1. Dilation must be 1.
"""
_validate_conv_params(
conv,
"LatticeCBSRAdvanced",
LatticeConv2dAdvanced.KERNEL_SIZES,
LatticeConv2dAdvanced.STRIDES,
)
[docs]
class LatticeCBSRAdvanced(LatticeCBSR):
"""Fused ``Conv2d -> [BatchNorm2d] -> [Activation]`` for advanced CNN IP.
Like :class:`~embedl_deploy._internal.lattice.modules.conv.LatticeCBSR` but
accepts the broader kernel set supported by advanced Lattice accelerators:
``1×1``, ``3×3``, ``5×5``, and ``7×7``. Only ``3×3`` allows stride 2; all
other kernel sizes require stride 1. Dilation must be 1.
:param conv:
The convolution layer to fuse.
:param bn:
Optional batch normalization layer.
:param act:
Optional ``ReLU`` or ``LatticeLeakyReLU`` activation.
:param output_precision:
Precision for the output quantizer.
:param output_fixed_calibration:
Optional fixed ``(scale, zero_point)`` pair for the output
quantizer.
:raises ValueError:
If `conv` has unsupported kernel size, stride, or dilation.
"""
def _validate_conv(self, conv: nn.Conv2d) -> None:
validate_conv_advanced(conv)
