# Copyright (C) 2026 Embedl AB
"""Fusion ``Pattern`` subclasses for the Lattice backend.
Each class declares a ``tree`` (what to match) and a ``graft`` (what to replace
with). The base ``Pattern`` class handles matching and replacement.
"""
import operator
from torch import fx, nn
from embedl_deploy._internal.core.patterns.main import Pattern, Phase
from embedl_deploy._internal.core.tree.match import node_check
from embedl_deploy._internal.core.tree.replace import make_fused
from embedl_deploy._internal.core.tree.types import (
Fork,
Graft,
Tree,
Wildcard,
)
from embedl_deploy._internal.core.tree.utils import get_module
from embedl_deploy._internal.lattice.modules.activation import (
LatticeActivationLike,
)
from embedl_deploy._internal.lattice.modules.conv import (
LatticeCBSR,
LatticeCBSRAdvanced,
validate_conv,
validate_conv_advanced,
)
from embedl_deploy._internal.lattice.modules.linear import LatticeFusedLinear
from embedl_deploy._internal.lattice.modules.pointwise import (
LatticeFusedAddAct,
)
#: Wildcard entry for an optional ``BatchNorm2d`` between convolution and
#: activation.
_OPTIONAL_BN = Wildcard(nn.BatchNorm2d, quantifier="?")
#: Wildcard entry for an optional ``BatchNorm1d`` between linear and
#: activation.
_OPTIONAL_LINEAR_BN = Wildcard(nn.BatchNorm1d, quantifier="?")
#: Wildcard entry for an optional activation after fused operation.
_OPTIONAL_ACT = Wildcard(LatticeActivationLike, quantifier="?")
def _is_cbsr_conv(node: fx.Node) -> bool:
"""Return ``True`` when the node is a ``Conv2d`` within the CBSR set.
Checks kernel size ∈ {1, 3}, stride ∈ {1, 2}, dilation == 1, and the
1×1-stride-1 constraint. Convolutions outside this set are skipped so that
``transform()`` does not crash.
"""
module = get_module(node)
if not isinstance(module, nn.Conv2d):
return False
try:
validate_conv(module)
except ValueError:
return False
return True
def _is_cbsr_conv_advanced(node: fx.Node) -> bool:
"""Return ``True`` when the node is a ``Conv2d`` within the advanced set.
Checks kernel size ∈ {1, 3, 5, 7}, stride ∈ {1, 2}, dilation == 1, and that
only 3×3 may use stride 2.
"""
module = get_module(node)
if not isinstance(module, nn.Conv2d):
return False
try:
validate_conv_advanced(module)
except ValueError:
return False
return True
[docs]
class LatticeCBSRPattern(Pattern):
"""Match ``Conv2d -> [BatchNorm2d] -> [Activation]`` and fuse into CBSR.
Both ``BatchNorm2d`` and the activation are optional — any prefix of the
``Conv -> BN -> Act`` chain is matched. The activation, when present, must
be ``ReLU`` or ``LatticeLeakyReLU``. The convolution must have
Lattice-supported parameters (kernel in {1, 3}, stride in {1, 2}, dilation
== 1); convolutions outside the supported set are silently skipped.
"""
phase = Phase.FUSION
tree: Tree = (_is_cbsr_conv, _OPTIONAL_BN, _OPTIONAL_ACT)
graft: Graft = (make_fused(LatticeCBSR),)
[docs]
class LatticeCBSRAdvancedPattern(Pattern):
"""Match ``Conv2d -> [BatchNorm2d] -> [Activation]`` and fuse into CBSR.
Both ``BatchNorm2d`` and the activation are optional — any prefix of the
``Conv -> BN -> Act`` chain is matched. The activation, when present, must
be ``ReLU`` or ``LatticeLeakyReLU``. The convolution must have
Lattice-supported parameters (kernel in {1, 3, 5, 7}, stride in {1, 2},
dilation == 1); convolutions outside the supported set are silently
skipped.
"""
phase = Phase.FUSION
tree: Tree = (_is_cbsr_conv_advanced, _OPTIONAL_BN, _OPTIONAL_ACT)
graft: Graft = (make_fused(LatticeCBSRAdvanced),)
class LatticeFusedLinearPattern(Pattern):
"""Match ``Linear -> [BN] -> [Act]`` and fuse into ``LatticeFusedLinear``.
Both ``BatchNorm1d`` and the activation are optional — any prefix of the
``Linear -> BN -> Act`` chain is matched. The activation, when present,
must be ``ReLU`` or ``LatticeLeakyReLU``.
"""
phase = Phase.FUSION
tree: Tree = (nn.Linear, _OPTIONAL_LINEAR_BN, _OPTIONAL_ACT)
graft: Graft = (make_fused(LatticeFusedLinear),)
@node_check
def _is_add(node: fx.Node) -> bool:
"""Match ``operator.add`` or ``operator.iadd``.
``symbolic_trace`` emits ``operator.add`` for both ``+`` and ``+=``.
``torch.export`` followed by recomposition emits ``operator.iadd`` for
in-place additions (e.g. ResNet residual shortcuts).
"""
return node.op == "call_function" and node.target in (
operator.add,
operator.iadd,
)
class LatticeFusedAddActPattern(Pattern):
"""Match ``add(·, ·) -> [Activation]`` and fuse into a ``LatticeFusedAddAct``.
Captures the element-wise addition at the merge point of residual
connections, with an optional trailing activation (``ReLU`` or
``LatticeLeakyReLU``). Both input branches are unconstrained — either may
come from any upstream module.
"""
phase = Phase.FUSION
tree: Tree = Fork(
((), ()),
_is_add,
(_OPTIONAL_ACT,),
)
graft: Graft = (make_fused(LatticeFusedAddAct),)
