# Copyright (C) 2026 Embedl AB
"""Fused ``nn.Module`` replacements for convolution-based patterns.
Each class represents a hardware-fusible operation that replaces a multi-op
chain found by the pattern matcher. The fused module keeps the original sub-
modules (``Conv``, ``BN``, ``ReLU``) as children so that:
* Weights are trivially transferred from the original model.
* ``forward()`` is numerically identical to the original chain.
* A later compilation step can fold the ``BN`` and emit a single kernel.
"""
import torch
import torch.nn.functional as F
from torch import nn
from embedl_deploy._internal.core.modules import (
ActivationLike,
FusedModule,
PrecisionSynced,
)
from embedl_deploy._internal.core.quantize.stubs import (
QuantStub,
WeightFakeQuantize,
)
def is_depthwise_conv(conv: nn.Conv2d) -> bool:
"""Return ``True`` when *conv* is depthwise."""
return conv.groups == conv.in_channels
def _conv_weight_forward(
conv: nn.Conv2d,
weight_fake_quant: WeightFakeQuantize | None,
x: torch.Tensor,
) -> torch.Tensor:
"""Run a convolution, fake-quantizing the weight when enabled."""
weight = (
weight_fake_quant(conv.weight)
if weight_fake_quant is not None
else conv.weight
)
# pylint: disable-next=not-callable
return F.conv2d(
x,
weight,
conv.bias,
conv.stride,
conv.padding,
conv.dilation,
conv.groups,
)
[docs]
class FusedConvBNAct(FusedModule):
"""Fused ``Conv2d → [BatchNorm2d] → Act``.
Depthwise convolutions set ``precision_deferred.downstream`` so the
surround pattern can decide output precision contextually.
"""
inputs_to_quantize: set[int] = {0}
def __init__(
self,
conv: nn.Conv2d,
bn: nn.BatchNorm2d | None,
act: ActivationLike,
) -> None:
super().__init__()
self.conv = conv
self.bn = bn
self.act = act
self.weight_fake_quant = WeightFakeQuantize({self})
self.precision_deferred.upstream = True
if isinstance(act, (nn.Sigmoid, nn.Tanh)):
self.precision_synced = PrecisionSynced(output=False)
if is_depthwise_conv(conv):
self.precision_deferred.downstream = True
@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``."""
wfq = getattr(self, "weight_fake_quant", None)
x = _conv_weight_forward(self.conv, wfq, x)
if self.bn is not None:
x = self.bn(x)
return self.act(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"FusedConvBNAct("
f"{self.conv.in_channels}→{self.conv.out_channels}, "
f"k={self.conv.kernel_size}, s={self.conv.stride}"
f"{bn_info})"
)
[docs]
class FusedConvBN(FusedModule):
"""Fused ``Conv2d → [BatchNorm2d]`` (no activation).
Depthwise convolutions set ``precision_deferred.downstream`` so the
surround pattern can decide output precision contextually.
"""
inputs_to_quantize: set[int] = {0}
def __init__(
self,
conv: nn.Conv2d,
bn: nn.BatchNorm2d | None,
) -> None:
super().__init__()
self.conv = conv
self.bn = bn
self.weight_fake_quant = WeightFakeQuantize({self})
self.precision_deferred.upstream = True
if is_depthwise_conv(conv):
self.precision_deferred.downstream = True
@property
def quantized_weight(self) -> torch.Tensor | None:
return self.conv.weight
[docs]
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Apply ``conv → [bn]``."""
wfq = getattr(self, "weight_fake_quant", None)
x = _conv_weight_forward(self.conv, wfq, x)
if self.bn is not None:
x = self.bn(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"FusedConvBN("
f"{self.conv.in_channels}→{self.conv.out_channels}, "
f"k={self.conv.kernel_size}, s={self.conv.stride}"
f"{bn_info})"
)
[docs]
class FusedConvBNActMaxPool(FusedModule):
"""Fused ``Conv2d → [BatchNorm2d] → Activation → MaxPool2d``."""
inputs_to_quantize: set[int] = {0}
def __init__(
self,
conv: nn.Conv2d,
bn: nn.BatchNorm2d | None,
act: ActivationLike,
maxpool: nn.MaxPool2d,
) -> None:
super().__init__()
self.conv = conv
self.bn = bn
self.act = act
self.maxpool = maxpool
self.weight_fake_quant = WeightFakeQuantize({self})
self.precision_deferred.upstream = True
if not isinstance(act, (nn.ReLU, nn.ReLU6)):
self.precision_synced = PrecisionSynced(output=False)
@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 → maxpool``."""
wfq = getattr(self, "weight_fake_quant", None)
x = _conv_weight_forward(self.conv, wfq, x)
if self.bn is not None:
x = self.bn(x)
x = self.act(x)
return self.maxpool(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)"
mp = self.maxpool
return (
f"FusedConvBNActMaxPool("
f"{self.conv.in_channels}→{self.conv.out_channels}, "
f"k={self.conv.kernel_size}, s={self.conv.stride}"
f"{bn_info}, "
f"pool_k={mp.kernel_size}, pool_s={mp.stride})"
)
[docs]
class FusedConvBNAddAct(FusedModule):
"""Fused ``Conv2d → [BatchNorm2d] → add(·, residual) → [Activation]``.
``forward()`` accepts two inputs: the main tensor ``x`` and the
``residual`` tensor. Both the ``BatchNorm2d`` and the trailing activation
are optional so that EfficientNet-style ``Conv → BN → Add`` blocks (no
activation) are captured.
"""
inputs_to_quantize: set[int] = {0, 1}
def __init__(
self,
conv: nn.Conv2d,
bn: nn.BatchNorm2d | None,
act: ActivationLike | None,
) -> None:
super().__init__()
self.conv = conv
self.bn = bn
self.act = act
self.weight_fake_quant = WeightFakeQuantize({self})
self.precision_deferred.upstream = True
if act is not None and not isinstance(act, (nn.ReLU, nn.ReLU6)):
self.precision_synced = PrecisionSynced(output=False)
@property
def quantized_weight(self) -> torch.Tensor | None:
return self.conv.weight
[docs]
def forward(self, x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor:
"""Apply ``conv → [bn] → add(·, residual) → [act]``."""
wfq = getattr(self, "weight_fake_quant", None)
x = _conv_weight_forward(self.conv, wfq, x)
if self.bn is not None:
x = self.bn(x)
x = x + residual
if self.act is not None:
x = 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)"
act_info = ""
if self.act is not None:
act_info = f", act={type(self.act).__name__}"
return (
f"FusedConvBNAddAct("
f"{self.conv.in_channels}→{self.conv.out_channels}, "
f"k={self.conv.kernel_size}, s={self.conv.stride}"
f"{bn_info}{act_info})"
)
class FusedConvBNSigmoidMul(FusedModule):
"""Fused ``Conv2d → [BatchNorm2d] → Sigmoid → Mul(·, skip)``.
Captures the SE gate pattern where an expand convolution produces channel
attention weights via Sigmoid, then element-wise multiplies with the skip
connection. An internal
:class:`~embedl_deploy._internal.core.quantize.stubs.QuantStub` between the
conv/BN output and the Sigmoid produces the Q/DQ pair that enables
TensorRT's ``PWN(Sigmoid, Mul)`` fusion.
``forward()`` accepts two inputs: the main tensor ``x`` feeding the
convolution, and the ``skip`` tensor multiplied by the sigmoid gate.
"""
inputs_to_quantize: set[int] = {0, 1}
def __init__(
self,
conv: nn.Conv2d,
bn: nn.BatchNorm2d | None,
sigmoid: nn.Sigmoid,
) -> None:
super().__init__()
self.conv = conv
self.bn = bn
self.sigmoid = sigmoid
self.weight_fake_quant = WeightFakeQuantize({self})
self.gate_quant = QuantStub({self})
self.precision_deferred.upstream = True
@property
def quantized_weight(self) -> torch.Tensor | None:
return self.conv.weight
def forward(
self,
x: torch.Tensor,
skip: torch.Tensor,
) -> torch.Tensor:
"""Apply ``conv → [bn] → gate_quant → sigmoid → mul(·, skip)``."""
wfq = getattr(self, "weight_fake_quant", None)
x = _conv_weight_forward(self.conv, wfq, x)
if self.bn is not None:
x = self.bn(x)
x = self.gate_quant(x)
x = self.sigmoid(x)
return x * skip
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"FusedConvBNSigmoidMul("
f"{self.conv.in_channels}→{self.conv.out_channels}, "
f"k={self.conv.kernel_size}, s={self.conv.stride}"
f"{bn_info})"
)
