Source code for embedl_deploy._internal.lattice.modules.linear

# Copyright (C) 2026 Embedl AB

"""Lattice fused ``nn.Module`` replacement for Linear/GEMM patterns.

:class:`LatticeFusedLinear` wraps ``Linear -> [BatchNorm1d] -> [Activation]``
as a single fused block with an output quantizer. Both batch normalization and
the activation are optional — any prefix of the chain is accepted.

The ``output_quant_stub``
(:class:`~embedl_deploy._internal.lattice.modules.quant.LatticeQuant`) is
initially configured to INT8 (unfixed). The Lattice STATE_PREP patterns will
upgrade intermediate blocks to INT8 and the final block to INT16 before
calibration, choosing signedness from the activation (unsigned for ReLU, signed
otherwise).

Weight fake-quantization follows the same ONNX-export guard used by the
convolution modules: :func:`torch.onnx.is_in_onnx_export` is checked inside
:func:`_linear_weight_forward` to avoid emitting spurious
``QuantizeLinear/DequantizeLinear`` nodes for the weight path.
"""

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 _linear_weight_forward(
    linear: nn.Linear,
    weight_fake_quant: WeightFakeQuantize,
    x: torch.Tensor,
) -> torch.Tensor:
    """Run a linear layer, 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 = (
        linear.weight
        if torch.onnx.is_in_onnx_export()
        else weight_fake_quant(linear.weight)
    )
    return F.linear(x, weight, linear.bias)




class LatticeFusedLinear(FusedModule):
    """Fused ``Linear -> [BatchNorm1d] -> [Activation]`` for Lattice hardware.

    The Lattice accelerator implements a linear (GEMM) layer with optional
    batch normalization and optional activation as a single fused block. When
    an activation is present it must be ``ReLU`` or ``LatticeLeakyReLU``.

    :param linear:
        The ``nn.Linear`` from the matched chain.
    :param bn:
        Optional batch normalization layer.
    :param act:
        Optional ``ReLU`` or ``LatticeLeakyReLU`` activation.
    :param output_precision:
        Initial precision for ``output_quant_stub``. ``STATE_PREP`` will
        override this unless `output_fixed_calibration` is also supplied.
    :param output_fixed_calibration:
        If provided, ``(scale, zero_point)`` fixes the output quantizer so
        that ``STATE_PREP`` and calibration leave it unchanged.
    """

    inputs_to_quantize: set[int] = set()

    def __init__(
        self,
        linear: nn.Linear,
        bn: nn.BatchNorm1d | None = None,
        act: LatticeActivationLike | None = None,
        *,
        output_precision: Precision = Precision.INT8,
        output_fixed_calibration: tuple[float, int] | None = None,
    ) -> None:
        super().__init__()
        self.linear = linear
        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,
        )

    @property
    def quantized_weight(self) -> torch.Tensor | None:
        return self.linear.weight



    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply ``linear -> [bn] -> [act]``."""
        x = _linear_weight_forward(self.linear, 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:  # pragma: no cover
        bn_info = ""
        if self.bn is not None:
            bn_info = f", bn={self.bn.num_features} (foldable)"
        return (
            f"LatticeFusedLinear("
            f"{self.linear.in_features}→{self.linear.out_features}"
            f"{bn_info})"
        )