# Copyright (C) 2026 Embedl AB
"""Pooling conversion patterns for Lattice."""
import copy
import logging
from collections.abc import Callable, Sequence
from functools import partial
import torch
from torch import fx, nn
from embedl_deploy._internal.core.patterns.main import Pattern, Phase
from embedl_deploy._internal.core.tree.replace import get_auto_name
from embedl_deploy._internal.core.tree.state import get_replaced_nodes
from embedl_deploy._internal.core.tree.types import (
Graft,
NodeInserter,
ReplacementMaker,
Tree,
TreeMatch,
Wildcard,
)
from embedl_deploy._internal.core.tree.utils import (
get_input_shape,
get_module,
resolve_module,
)
from embedl_deploy._internal.lattice.modules.conv import (
LatticeConv2dAdvanced,
)
from embedl_deploy._internal.lattice.modules.pool import (
LatticeAdaptiveAvgPool2d,
LatticeMaxPool2d,
)
_LOG = logging.getLogger(__name__)
_OPTIONAL_BN = Wildcard(nn.BatchNorm2d, quantifier="?")
_OPTIONAL_RELU = Wildcard(nn.ReLU, quantifier="?")
def _has_maxpool_after_chain(node: fx.Node) -> bool:
"""Return ``True`` when a ``MaxPool2d`` follows ``Conv->[BN]->[ReLU]``.
Walks from a ``Conv2d`` node through optional ``BatchNorm2d`` and ``ReLU``
successors to check whether any path ends at a ``MaxPool2d``.
"""
frontier = {node}
for mod_type in (nn.BatchNorm2d, nn.ReLU):
next_frontier: set[fx.Node] = set()
for cur in frontier:
for user in cur.users:
user_mod = get_module(user)
if isinstance(user_mod, mod_type):
next_frontier.add(user)
frontier |= next_frontier
return any(
isinstance(get_module(u), nn.MaxPool2d)
for cur in frontier
for u in cur.users
)
def _is_stride2_conv(node: fx.Node) -> bool:
"""Return ``True`` for stride-2 ``Conv2d`` in ``Conv->[BN]->[ReLU]``.
Rejects nodes that are followed by ``MaxPool2d`` (possibly via ``BN``
and/or ``ReLU``) — those are handled by
:class:`LatticeStride2AsymPadConvMaxPoolPattern` instead.
"""
mod = get_module(node)
if not isinstance(mod, nn.Conv2d):
return False
if mod.stride != (2, 2):
return False
return not _has_maxpool_after_chain(node)
def _make_stride2_conv_with_pool(
tree_match: TreeMatch,
) -> tuple[nn.Module, ...]:
"""Build ``Conv(stride=1)->[BN]->[ReLU]->MaxPool2d(2,2)`` replacements."""
conv_node = tree_match.get_node(0)
conv = resolve_module(conv_node, nn.Conv2d)
new_conv = copy.deepcopy(conv)
new_conv.stride = (1, 1)
_LOG.warning(
"%s: Conv2d(stride=%s) replaced with "
"Conv2d(stride=(1, 1)) + MaxPool2d(kernel_size=2, stride=2) "
"— output will differ; retraining recommended.",
conv_node.name,
conv.stride,
)
result: list[nn.Module] = [new_conv]
bn_wc = tree_match.get_node(1, is_wildcard=True)
result.extend(
copy.deepcopy(resolve_module(n, nn.BatchNorm2d)) for n in bn_wc.nodes
)
relu_wc = tree_match.get_node(2, is_wildcard=True)
result.extend(
copy.deepcopy(resolve_module(n, nn.ReLU)) for n in relu_wc.nodes
)
result.append(
LatticeMaxPool2d(
nn.MaxPool2d(kernel_size=2, stride=2, padding=0, ceil_mode=True)
),
)
return tuple(result)
[docs]
class LatticeStride2ConvBnReLUPattern(Pattern):
"""Rewrite ``Conv(stride=2)->[BN]->[ReLU]`` as ``Conv(stride=1)+MaxPool2d``.
This pattern preserves convolution weights by moving the downsampling step
into a following ``MaxPool2d(kernel_size=2, stride=2)``. The
``BatchNorm2d`` and ``ReLU`` are optional — standalone ``Conv(stride=2)``
and ``Conv(stride=2)->BN`` chains are also matched.
Stride-2 convolutions that are already followed by a ``MaxPool2d``
(possibly via ``BN`` and/or ``ReLU``) are excluded — those are handled by
:class:`LatticeStride2AsymPadConvMaxPoolPattern`.
"""
is_numerically_equivalent = False
phase = Phase.CONVERSION
tree: Tree = (_is_stride2_conv, _OPTIONAL_BN, _OPTIONAL_RELU)
graft: Graft = (_make_stride2_conv_with_pool,)
def _is_stride2_same_pad_conv_before_maxpool(node: fx.Node) -> bool:
"""Return ``True`` for stride-2 SAME-padded ``Conv2d`` followed by ``MaxPool2d``.
SAME padding means ``padding == (kernel_size // 2, kernel_size // 2)``.
Once converted to asymmetric padding the convolution no longer satisfies
this predicate, preventing re-application on subsequent passes.
"""
mod = get_module(node)
if not isinstance(mod, nn.Conv2d):
return False
if mod.stride != (2, 2):
return False
kh = (
mod.kernel_size[0]
if isinstance(mod.kernel_size, tuple)
else mod.kernel_size
)
kw = (
mod.kernel_size[1]
if isinstance(mod.kernel_size, tuple)
else mod.kernel_size
)
if mod.padding != (kh // 2, kw // 2):
return False
return _has_maxpool_after_chain(node)
def _make_stride2_asym_pad_conv_maxpool(
tree_match: TreeMatch,
) -> tuple[nn.Module, ...]:
"""Build ``ZeroPad2d(asym)->Conv(3x3,stride=2,pad=0)->[BN]->[ReLU]->MaxPool``.
When the original kernel is larger than 3×3 it is snapped down to 3×3
because only 3×3 supports stride 2 on Lattice hardware. Weights are
reinitialized in that case.
"""
conv_node = tree_match.get_node(0)
conv = resolve_module(conv_node, nn.Conv2d)
kh = (
conv.kernel_size[0]
if isinstance(conv.kernel_size, tuple)
else conv.kernel_size
)
kw = (
conv.kernel_size[1]
if isinstance(conv.kernel_size, tuple)
else conv.kernel_size
)
# Only 3×3 supports stride 2 on Lattice — snap larger kernels down.
if max(kh, kw) > 3:
new_conv = LatticeConv2dAdvanced(
nn.Conv2d(
conv.in_channels,
conv.out_channels,
kernel_size=3,
stride=2,
padding=0,
dilation=1,
groups=conv.groups,
bias=conv.bias is not None,
)
)
new_conv.padding = (0, 0)
new_conv.to(device=conv.weight.device, dtype=conv.weight.dtype)
pad = nn.ZeroPad2d((0, 1, 0, 1))
_LOG.warning(
"%s: Conv2d(kernel_size=%s, stride=%s, padding=%s) snapped to "
"ZeroPad2d(%s) + Conv2d(kernel_size=(3, 3), stride=%s, "
"padding=(0, 0)) — weights reinitialized, output will differ; "
"retraining recommended.",
conv_node.name,
conv.kernel_size,
conv.stride,
conv.padding,
(0, 1, 0, 1),
conv.stride,
)
else:
# Emit a LatticeConv2dAdvanced so that Conv2d-snapping patterns
# (which skip isinstance(mod, LatticeConv2dAdvanced)) leave it alone.
new_conv = LatticeConv2dAdvanced(conv)
new_conv.padding = (0, 0)
# Asymmetric: pad only on right / bottom.
pad = nn.ZeroPad2d((0, kw // 2, 0, kh // 2))
_LOG.warning(
"%s: Conv2d(stride=%s, padding=%s) replaced with "
"ZeroPad2d(%s) + Conv2d(stride=%s, padding=(0, 0)) "
"— output will differ; retraining recommended.",
conv_node.name,
conv.stride,
conv.padding,
(0, kw // 2, 0, kh // 2),
conv.stride,
)
result: list[nn.Module] = [pad, new_conv]
bn_wc = tree_match.get_node(1, is_wildcard=True)
result.extend(
copy.deepcopy(resolve_module(n, nn.BatchNorm2d)) for n in bn_wc.nodes
)
relu_wc = tree_match.get_node(2, is_wildcard=True)
result.extend(
copy.deepcopy(resolve_module(n, nn.ReLU)) for n in relu_wc.nodes
)
pool_node = tree_match.get_node(3)
result.append(copy.deepcopy(resolve_module(pool_node, nn.MaxPool2d)))
return tuple(result)
class LatticeStride2AsymPadConvMaxPoolPattern(Pattern):
"""Rewrite ``Conv(stride=2)->[BN]->[ReLU]->MaxPool`` to use asymmetric padding.
When a stride-2 convolution with SAME padding is followed by ``MaxPool2d``
(possibly via ``BatchNorm2d`` and/or ``ReLU``), this pattern replaces the
symmetric SAME padding with an explicit ``ZeroPad2d`` that pads only on the
right and bottom edges, and sets the convolution padding to zero.
Only matches convolutions whose padding equals ``kernel_size // 2`` (the
SAME convention). After conversion the padding is ``(0, 0)`` which no
longer satisfies the predicate, ensuring convergence.
"""
is_numerically_equivalent = False
phase = Phase.CONVERSION
tree: Tree = (
_is_stride2_same_pad_conv_before_maxpool,
_OPTIONAL_BN,
_OPTIONAL_RELU,
nn.MaxPool2d,
)
graft: Graft = (_make_stride2_asym_pad_conv_maxpool,)
def _is_sppf_maxpool2d(node: fx.Node) -> bool:
"""Return ``True`` when `node` is a stride-1 ``MaxPool2d`` in an SPPF block.
SPPF (Spatial Pyramid Pooling - Fast) uses cascaded
``MaxPool2d(kernel_size=5, stride=1, padding=2)`` operations whose outputs
are concatenated. Lattice hardware does not support anything other than the
canonical ``MaxPool2d(kernel_size=2, stride=2, padding=0)`` — snapping a
stride-1 pool to that canonical form would halve spatial dimensions and
break the SPPF structure.
To avoid false positives on standalone stride-1 pools that are not part of
an SPPF module, this predicate requires at least one user of the node to be
a ``torch.cat`` call — the defining structural characteristic of SPPF.
"""
mod = get_module(node)
if not isinstance(mod, nn.MaxPool2d):
return False
if isinstance(mod, LatticeMaxPool2d):
return False
stride = mod.stride
if isinstance(stride, tuple):
stride = stride[0] if stride[0] == stride[1] else -1
if stride != 1:
return False
if get_input_shape(node) is None:
return False
return any(
u.op == "call_function" and u.target is torch.cat for u in node.users
)
def _build_sppf_conv_modules(node: fx.Node) -> tuple[nn.Module, ...]:
"""Build ``Conv2d(3×3)`` + ``BatchNorm2d`` + ``ReLU`` for one SPPF pool.
Replaces a stride-1 ``MaxPool2d`` with a learnable convolution block that
preserves spatial dimensions and channel count, always using a 3×3 kernel.
"""
shape = get_input_shape(node)
assert shape is not None
channels = int(shape[1])
pool = resolve_module(node, nn.MaxPool2d)
_LOG.warning(
"%s: MaxPool2d(kernel_size=%s, stride=%s) replaced with "
"Conv2d + BatchNorm2d + ReLU — output will differ; "
"retraining required.",
node.name,
pool.kernel_size,
pool.stride,
)
conv = nn.Conv2d(channels, channels, 3, stride=1, padding=1, bias=False)
bn = nn.BatchNorm2d(channels)
bn.train(pool.training)
relu = nn.ReLU(inplace=True)
return (conv, bn, relu)
def _build_sppf_conv_modules_advanced(node: fx.Node) -> tuple[nn.Module, ...]:
"""Build kernel-matched ``Conv2d`` + ``BatchNorm2d`` + ``ReLU`` for one pool.
Uses a convolution kernel that matches the original ``MaxPool2d`` kernel
size whenever the advanced hardware supports it (``3×3``, ``5×5``, or
``7×7``), preserving the receptive field of the original pool.
"""
shape = get_input_shape(node)
assert shape is not None
channels = int(shape[1])
pool = resolve_module(node, nn.MaxPool2d)
pool_k = pool.kernel_size
if isinstance(pool_k, tuple):
pool_k = max(pool_k)
conv_k: int = min((3, 5, 7), key=lambda k: (abs(k - pool_k), k))
_LOG.warning(
"%s: MaxPool2d(kernel_size=%s, stride=%s) replaced with "
"Conv2d(kernel_size=%s) + BatchNorm2d + ReLU — output will differ; "
"retraining required.",
node.name,
pool.kernel_size,
pool.stride,
conv_k,
)
conv = nn.Conv2d(
channels, channels, conv_k, stride=1, padding=conv_k // 2, bias=False
)
bn = nn.BatchNorm2d(channels)
bn.train(pool.training)
relu = nn.ReLU(inplace=True)
return (conv, bn, relu)
def _insert_module_chain(
graph_module: fx.GraphModule,
modules: Sequence[nn.Module],
input_node: fx.Node,
) -> list[fx.Node]:
"""Insert `modules` as a ``call_module`` chain fed from `input_node`."""
created: list[fx.Node] = []
prev = input_node
for module in modules:
name = get_auto_name(graph_module, module)
graph_module.add_module(name, module)
with graph_module.graph.inserting_after(prev):
node = graph_module.graph.call_module(name, (prev,))
created.append(node)
prev = node
return created
def _sppf_replacement(
tree_match: TreeMatch,
build_fn: Callable[[fx.Node], tuple[nn.Module, ...]],
) -> tuple[NodeInserter]:
"""Build a ``NodeInserter`` that rewrites a three-pool SPPF cascade.
`build_fn` produces the ``Conv → BatchNorm2d → ReLU`` modules for one pool
(the optimized and advanced patterns differ only in this;
:func:`_sppf_maker` binds it to obtain a ``ReplacementMaker``). The
inserter grows an independent chain per pool and rewires the ``torch.cat``
slots of the first two pools — the cascade's intermediate outputs — onto
their new chains. The third pool is the matched tree's output, so the
surrounding ``replace_tree`` call rewires its ``cat`` slot to the final
chain and erases all three pools.
"""
replaced_nodes = get_replaced_nodes()
pools = [replaced_nodes.get(n, n) for n in tree_match.get_tree_nodes()]
def _insert(
graph_module: fx.GraphModule,
prev_args: tuple[fx.Node, ...],
) -> list[fx.Node]:
created: list[fx.Node] = []
source = prev_args[0]
relus: list[fx.Node] = []
for pool in pools:
chain = _insert_module_chain(graph_module, build_fn(pool), source)
created.extend(chain)
relus.append(chain[-1])
source = chain[-1]
for pool, relu in zip(pools[:-1], relus[:-1], strict=True):
pool.replace_all_uses_with(relu)
return created
return (_insert,)
def _sppf_maker(
build_fn: Callable[[fx.Node], tuple[nn.Module, ...]],
) -> ReplacementMaker:
"""Bind `build_fn` into a ``ReplacementMaker`` for an SPPF pattern."""
return partial(_sppf_replacement, build_fn=build_fn)
[docs]
class LatticeSPPFMaxPoolPattern(Pattern):
"""Replace three cascaded stride-1 ``MaxPool2d`` with ``Conv+BN+ReLU``.
SPPF (Spatial Pyramid Pooling - Fast) modules use three cascaded
``MaxPool2d`` operations with stride 1 to expand the receptive field
without reducing spatial resolution. Lattice hardware only supports
``MaxPool2d`` with kernel 2, stride 2, and zero padding — snapping a
stride-1 pool to that canonical form would halve spatial dimensions and
destroy the SPPF structure.
This pattern matches the full three-pool cascade at once and replaces all
three with ``Conv2d(c, c, 3, stride=1, padding=1)`` + ``BatchNorm2d`` +
``ReLU`` chains. The cascaded 3×3 convolutions provide receptive field
expansion (3→5→7) analogous to the original cascaded max pools (5→9→13),
but the replacement is **not** mathematically equivalent — retraining is
required.
Standalone stride-1 pools that do not feed ``torch.cat`` are left for
``LatticeMaxPool2dPattern`` to handle.
Requires shape metadata on the input node, propagated by
:class:`~torch.fx.passes.shape_prop.ShapeProp`.
"""
is_numerically_equivalent = False
phase = Phase.CONVERSION
tree: Tree = (_is_sppf_maxpool2d, _is_sppf_maxpool2d, _is_sppf_maxpool2d)
graft: Graft = (_sppf_maker(_build_sppf_conv_modules),)
[docs]
class LatticeSPPFMaxPoolAdvancedPattern(Pattern):
"""Replace three cascaded stride-1 ``MaxPool2d`` with kernel-matched convs.
Like ``LatticeSPPFMaxPoolPattern``, this matches the full three-pool SPPF
cascade and replaces it with ``Conv2d`` + ``BatchNorm2d`` + ``ReLU``
chains. The difference is the convolution kernel size: instead of always
using ``3×3``, this pattern snaps to the nearest kernel from ``{3, 5, 7}``
to match the original pool's receptive field whenever the advanced hardware
supports it.
For example, a ``MaxPool2d(kernel_size=5, stride=1, padding=2)`` is
replaced with ``Conv2d(c, c, 5, stride=1, padding=2)`` — preserving the 5×5
receptive field. The replacement is **not** mathematically equivalent —
retraining is required.
Requires shape metadata on the input node, propagated by
:class:`~torch.fx.passes.shape_prop.ShapeProp`.
"""
is_numerically_equivalent = False
phase = Phase.CONVERSION
tree: Tree = (_is_sppf_maxpool2d, _is_sppf_maxpool2d, _is_sppf_maxpool2d)
graft: Graft = (_sppf_maker(_build_sppf_conv_modules_advanced),)
def _is_nonconforming_maxpool2d(node: fx.Node) -> bool:
"""Return ``True`` when `node` is a ``MaxPool2d`` that needs rewriting."""
# SPPF stride-1 pools must be handled by the dedicated SPPF patterns.
# If the generic max-pool snap grabs them, spatial dims are halved and
# the downstream SPPF concat can break.
if _is_sppf_maxpool2d(node):
return False
mod = get_module(node)
if not isinstance(mod, nn.MaxPool2d):
return False
if isinstance(mod, LatticeMaxPool2d):
return False
return not LatticeMaxPool2d.is_compatible(mod)
def _make_lattice_maxpool(tree_match: TreeMatch) -> tuple[nn.Module, ...]:
"""Build the replacement ``MaxPool2d`` for the matched node."""
node = tree_match.get_node(0)
pool = resolve_module(node, nn.MaxPool2d)
_LOG.warning(
"%s: MaxPool2d snapped from kernel_size=%s stride=%s padding=%s "
"to kernel_size=%s stride=%s padding=%s — output will differ",
node.name,
pool.kernel_size,
pool.stride,
pool.padding,
LatticeMaxPool2d.KERNEL_SIZE,
LatticeMaxPool2d.STRIDE,
LatticeMaxPool2d.PADDING,
)
return (LatticeMaxPool2d(pool),)
[docs]
class LatticeMaxPool2dPattern(Pattern):
"""Snap an out-of-spec ``MaxPool2d`` to Lattice's supported set.
Lattice hardware supports only a ``2×2`` max-pool with stride 2 and zero
padding. Any other configuration is replaced by that single canonical form,
matching the stem rewrite performed by the Lattice export reference script.
"""
is_numerically_equivalent = False
phase = Phase.CONVERSION
tree: Tree = (_is_nonconforming_maxpool2d,)
graft: Graft = (_make_lattice_maxpool,)
def _avg_pool_output_is_global(node: fx.Node) -> bool:
"""Return ``True`` when the pool collapses spatial dims to ``1×1``."""
mod = get_module(node)
shape = get_input_shape(node)
if shape is None or len(shape) != 4:
return False
in_h, in_w = int(shape[2]), int(shape[3])
if isinstance(mod, nn.AdaptiveAvgPool2d):
if isinstance(mod.output_size, tuple):
out_h, out_w = mod.output_size
else:
out_h = out_w = mod.output_size
out_h = out_h if out_h is not None else in_h
out_w = out_w if out_w is not None else in_w
return out_h == 1 and out_w == 1
if isinstance(mod, nn.AvgPool2d):
kh, kw = (
mod.kernel_size
if isinstance(mod.kernel_size, tuple)
else (mod.kernel_size, mod.kernel_size)
)
return kh == in_h and kw == in_w
return False
def _is_nonconforming_global_avg_pool(node: fx.Node) -> bool:
"""Return ``True`` when `node` is a global pool not in canonical form."""
mod = get_module(node)
if not isinstance(mod, (nn.AdaptiveAvgPool2d, nn.AvgPool2d)):
return False
if isinstance(mod, LatticeAdaptiveAvgPool2d):
return False
return _avg_pool_output_is_global(node)
def _make_lattice_global_pool(tree_match: TreeMatch) -> tuple[nn.Module, ...]:
"""Build the replacement global average pool for the matched node."""
del tree_match
return (LatticeAdaptiveAvgPool2d(),)
[docs]
class LatticeGlobalAvgPoolPattern(Pattern):
"""Normalize any global average pool to ``AdaptiveAvgPool2d((1, 1))``.
Lattice hardware represents the classification-tail pool as an
:class:`~torch.nn.AdaptiveAvgPool2d` with ``output_size == (1, 1)``. Any
other module that effectively performs a global average pool over a 4-D
input (an :class:`~torch.nn.AdaptiveAvgPool2d` with the scalar shorthand
``1`` or with a ``None`` axis equal to the input spatial size, or an
:class:`~torch.nn.AvgPool2d` whose kernel covers the full spatial extent)
is rewritten to this canonical form.
The pattern is a no-op for inputs already in canonical form, so it
converges in a single conversion pass without infinite looping.
Requires shape metadata on the input node, propagated by
:class:`~torch.fx.passes.shape_prop.ShapeProp` (the ``transform`` driver
runs this automatically between conversion iterations).
"""
phase = Phase.CONVERSION
tree: Tree = (_is_nonconforming_global_avg_pool,)
graft: Graft = (_make_lattice_global_pool,)
