Source code for embedl_deploy._internal.lattice.patterns.fusions

# 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.
"""

from torch import fx, nn

from embedl_deploy._internal.core.patterns.main import Pattern, Phase
from embedl_deploy._internal.core.tree.replace import make_fused
from embedl_deploy._internal.core.tree.types import (
    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,
)

#: Wildcard entry for an optional ``BatchNorm2d`` between convolution and
#: activation.
_OPTIONAL_BN = Wildcard(nn.BatchNorm2d, 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




class LatticeCBSRPattern(Pattern):
    """Match ``Conv2d -> [BatchNorm2d] -> Activation`` and fuse into CBSR.

    The ``BatchNorm2d`` is optional — both ``Conv -> BN -> Act`` and ``Conv ->
    Act`` chains are matched. The activation 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, LatticeActivationLike)
    graft: Graft = (make_fused(LatticeCBSR),)





class LatticeCBSRAdvancedPattern(Pattern):
    """Match ``Conv2d -> [BatchNorm2d] -> Activation`` and fuse into CBSR.

    The ``BatchNorm2d`` is optional — both ``Conv -> BN -> Act`` and ``Conv ->
    Act`` chains are matched. The activation 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, LatticeActivationLike)
    graft: Graft = (make_fused(LatticeCBSRAdvanced),)