From ff1fb6ccbd1a9cfea0ad0b6c3a593a1e421610f1 Mon Sep 17 00:00:00 2001
From: asteinh <armin.steinhauser@raws.at>
Date: Sat, 15 Aug 2026 13:44:09 +0200
Subject: [PATCH 1/7] feature: 2D receptive field and HW tile axis constant

---
 src/tigris/__init__.py                   |  2 +
 src/tigris/analysis/partition_spatial.py | 67 ++++++++++++++++++------
 tests/test_2d_tiling.py                  | 38 ++++++++++++++
 tests/test_tiling.py                     | 47 +++++++++--------
 4 files changed, 115 insertions(+), 39 deletions(-)
 create mode 100644 tests/test_2d_tiling.py

diff --git a/src/tigris/__init__.py b/src/tigris/__init__.py
index 2941e4b..d11273d 100644
--- a/src/tigris/__init__.py
+++ b/src/tigris/__init__.py
@@ -14,6 +14,7 @@
 TILE_AXIS_NONE = 0
 TILE_AXIS_HEIGHT_OR_LENGTH = 1
 TILE_AXIS_WIDTH = 2
+TILE_AXIS_HW = 3
 
 __all__ = [
     "__version__",
@@ -25,4 +26,5 @@
     "TILE_AXIS_NONE",
     "TILE_AXIS_HEIGHT_OR_LENGTH",
     "TILE_AXIS_WIDTH",
+    "TILE_AXIS_HW",
 ]
diff --git a/src/tigris/analysis/partition_spatial.py b/src/tigris/analysis/partition_spatial.py
index 3daac08..57da144 100644
--- a/src/tigris/analysis/partition_spatial.py
+++ b/src/tigris/analysis/partition_spatial.py
@@ -76,29 +76,40 @@ def classify_op(op_type: str) -> TileCategory:
 
 
 def compute_receptive_field(ops: list[OpNode]) -> tuple[int, int]:
-    """Compute the receptive field and total stride for a sequence of ops.
+    """Compute the height and width receptive fields for a sequence of ops.
 
-    Walks the ops in reverse, accumulating RF and jump (cumulative stride).
-    Returns (receptive_field, total_jump).
+    Walks the ops in reverse once, accumulating RF and jump (cumulative
+    stride) independently for the height and width axes.
+    Returns (rf_h, rf_w).
 
-    For pointwise ops, RF and jump are unchanged.
+    For pointwise ops, RF and jump are unchanged on both axes.
     For conv/pool ops, RF grows based on effective kernel size.
     """
-    rf = 1
-    jump = 1
+    rf_h = 1
+    jump_h = 1
+    rf_w = 1
+    jump_w = 1
 
     for op in reversed(ops):
         cat = classify_op(op.op_type)
         if cat in (TileCategory.CONV, TileCategory.POOL):
-            kernel = _get_kernel_h(op)
-            stride = _get_stride_h(op)
-            dilation = _get_dilation_h(op)
+            kernel_h = _get_kernel_h(op)
+            stride_h = _get_stride_h(op)
+            dilation_h = _get_dilation_h(op)
 
-            effective_k = dilation * (kernel - 1) + 1
-            rf = rf + (effective_k - 1) * jump
-            jump = jump * stride
+            effective_kh = dilation_h * (kernel_h - 1) + 1
+            rf_h = rf_h + (effective_kh - 1) * jump_h
+            jump_h = jump_h * stride_h
 
-    return rf, jump
+            kernel_w = _get_kernel_w(op)
+            stride_w = _get_stride_w(op)
+            dilation_w = _get_dilation_w(op)
+
+            effective_kw = dilation_w * (kernel_w - 1) + 1
+            rf_w = rf_w + (effective_kw - 1) * jump_w
+            jump_w = jump_w * stride_w
+
+    return rf_h, rf_w
 
 
 def _get_kernel_h(op: OpNode) -> int:
@@ -125,6 +136,30 @@ def _get_dilation_h(op: OpNode) -> int:
     return 1
 
 
+def _get_kernel_w(op: OpNode) -> int:
+    """Get the width dimension of the kernel (second element of kernel_shape)."""
+    ks = op.attrs.get("kernel_shape")
+    if ks and len(ks) >= 2:
+        return int(ks[1])
+    return 1
+
+
+def _get_stride_w(op: OpNode) -> int:
+    """Get the width dimension of the stride."""
+    strides = op.attrs.get("strides")
+    if strides and len(strides) >= 2:
+        return int(strides[1])
+    return 1
+
+
+def _get_dilation_w(op: OpNode) -> int:
+    """Get the width dimension of the dilation."""
+    dilations = op.attrs.get("dilations")
+    if dilations and len(dilations) >= 2:
+        return int(dilations[1])
+    return 1
+
+
 # Tile solver
 
 
@@ -168,8 +203,8 @@ def partition_spatial(ag: AnalyzedGraph) -> AnalyzedGraph:
             continue
 
         # Compute receptive field
-        rf, _jump = compute_receptive_field(stage_ops)
-        halo = rf - 1
+        rf_h, _rf_w = compute_receptive_field(stage_ops)
+        halo = rf_h - 1
 
         # Axis 1 in the serialized NHWC/NLC layout maps to H/L at source dim 2.
         input_h = _find_input_extent(ag, stage, tile_axis)
@@ -208,7 +243,7 @@ def partition_spatial(ag: AnalyzedGraph) -> AnalyzedGraph:
             tile_height=tile_h,
             num_tiles=num_tiles,
             halo=halo,
-            receptive_field=rf,
+            receptive_field=rf_h,
             original_height=input_h,
             tiled_peak_bytes=tiled_peak,
             overhead_bytes=overhead,
diff --git a/tests/test_2d_tiling.py b/tests/test_2d_tiling.py
new file mode 100644
index 0000000..85faeb3
--- /dev/null
+++ b/tests/test_2d_tiling.py
@@ -0,0 +1,38 @@
+"""Tests for 2D (height + width) receptive field computation."""
+
+from tigris import TILE_AXIS_HW
+from tigris.analysis.partition_spatial import compute_receptive_field
+from tigris.graph.ir import OpNode
+
+
+def make_conv_op(kernel, stride, dilation):
+    """Build the minimal OpNode a Conv-like op needs for receptive field analysis."""
+    return OpNode(
+        name="conv",
+        op_type="Conv",
+        inputs=[],
+        outputs=[],
+        attrs={
+            "kernel_shape": list(kernel),
+            "strides": list(stride),
+            "dilations": list(dilation),
+        },
+    )
+
+
+def test_hw_axis_constant():
+    assert TILE_AXIS_HW == 3
+
+
+def test_receptive_field_returns_both_axes():
+    # A single 3x3 stride-1 conv: rf_h == rf_w == 3.
+    ops = [make_conv_op(kernel=(3, 3), stride=(1, 1), dilation=(1, 1))]
+    rf_h, rf_w = compute_receptive_field(ops)
+    assert (rf_h, rf_w) == (3, 3)
+
+
+def test_receptive_field_asymmetric_kernel():
+    # 5x3 kernel, stride 2x1: rf_h = 5, rf_w = 3.
+    ops = [make_conv_op(kernel=(5, 3), stride=(2, 1), dilation=(1, 1))]
+    rf_h, rf_w = compute_receptive_field(ops)
+    assert (rf_h, rf_w) == (5, 3)
diff --git a/tests/test_tiling.py b/tests/test_tiling.py
index abe7cba..33a29af 100644
--- a/tests/test_tiling.py
+++ b/tests/test_tiling.py
@@ -143,29 +143,30 @@ def test_gemm_is_untileable(self):
 
 class TestReceptiveField:
     def test_single_3x3_conv(self):
-        """A single 3x3 conv has RF=3."""
+        """A single 3x3 conv has RF=3 on both axes."""
         ops = [OpNode(name="c", op_type="Conv", inputs=[], outputs=[],
                       attrs={"kernel_shape": [3, 3], "strides": [1, 1]})]
-        rf, jump = compute_receptive_field(ops)
-        assert rf == 3
-        assert jump == 1
+        rf_h, rf_w = compute_receptive_field(ops)
+        assert rf_h == 3
+        assert rf_w == 3
 
     def test_two_3x3_convs(self):
-        """Two stacked 3x3 convs have RF=5."""
+        """Two stacked 3x3 convs have RF=5 on both axes."""
         ops = [
             OpNode(name="c0", op_type="Conv", inputs=[], outputs=[],
                    attrs={"kernel_shape": [3, 3], "strides": [1, 1]}),
             OpNode(name="c1", op_type="Conv", inputs=[], outputs=[],
                    attrs={"kernel_shape": [3, 3], "strides": [1, 1]}),
         ]
-        rf, jump = compute_receptive_field(ops)
-        assert rf == 5
-        assert jump == 1
+        rf_h, rf_w = compute_receptive_field(ops)
+        assert rf_h == 5
+        assert rf_w == 5
 
     def test_conv_stride2_pool(self):
-        """Conv3x3(s=1) + MaxPool2x2(s=2) + Conv3x3(s=1).
+        """Conv3x3(s=1) + MaxPool2x2(s=2) + Conv3x3(s=1), symmetric kernel/stride.
 
         reversed: conv1(k=3,s=1): rf=3, j=1 -> pool(k=2,s=2): rf=4, j=2 -> conv0(k=3,s=1): rf=8, j=2
+        Kernel and stride are symmetric across height/width, so rf_h == rf_w == 8.
         """
         ops = [
             OpNode(name="c0", op_type="Conv", inputs=[], outputs=[],
@@ -175,18 +176,18 @@ def test_conv_stride2_pool(self):
             OpNode(name="c1", op_type="Conv", inputs=[], outputs=[],
                    attrs={"kernel_shape": [3, 3], "strides": [1, 1]}),
         ]
-        rf, jump = compute_receptive_field(ops)
-        assert rf == 8
-        assert jump == 2
+        rf_h, rf_w = compute_receptive_field(ops)
+        assert rf_h == 8
+        assert rf_w == 8
 
     def test_dilated_conv(self):
-        """Conv3x3 with dilation=2: effective_k = 2*(3-1)+1 = 5, RF=5."""
+        """Conv3x3 with dilation=2: effective_k = 2*(3-1)+1 = 5, RF=5 on both axes."""
         ops = [OpNode(name="c", op_type="Conv", inputs=[], outputs=[],
                       attrs={"kernel_shape": [3, 3], "strides": [1, 1],
                              "dilations": [2, 2]})]
-        rf, jump = compute_receptive_field(ops)
-        assert rf == 5
-        assert jump == 1
+        rf_h, rf_w = compute_receptive_field(ops)
+        assert rf_h == 5
+        assert rf_w == 5
 
     def test_pointwise_passthrough(self):
         """Pointwise ops (Relu) don't change RF."""
@@ -195,19 +196,19 @@ def test_pointwise_passthrough(self):
                    attrs={"kernel_shape": [3, 3], "strides": [1, 1]}),
             OpNode(name="r", op_type="Relu", inputs=[], outputs=[], attrs={}),
         ]
-        rf, jump = compute_receptive_field(ops)
-        assert rf == 3  # same as single conv
-        assert jump == 1
+        rf_h, rf_w = compute_receptive_field(ops)
+        assert rf_h == 3  # same as single conv
+        assert rf_w == 3
 
     def test_only_pointwise_rf_is_1(self):
-        """A chain of only pointwise ops has RF=1."""
+        """A chain of only pointwise ops has RF=1 on both axes."""
         ops = [
             OpNode(name="r0", op_type="Relu", inputs=[], outputs=[], attrs={}),
             OpNode(name="r1", op_type="Add", inputs=[], outputs=[], attrs={}),
         ]
-        rf, jump = compute_receptive_field(ops)
-        assert rf == 1
-        assert jump == 1
+        rf_h, rf_w = compute_receptive_field(ops)
+        assert rf_h == 1
+        assert rf_w == 1
 
 
 # Integration: full pipeline with ONNX fixtures

From 3f9a4b4426b6814d534b8fa5a55707461b8da1b8 Mon Sep 17 00:00:00 2001
From: asteinh <armin.steinhauser@raws.at>
Date: Sat, 15 Aug 2026 14:01:06 +0200
Subject: [PATCH 2/7] feature: 2D tile solver and HW axis selection on
 1D-infeasible stages

---
 src/tigris/analysis/partition_spatial.py | 153 +++++++++++++++++++++--
 src/tigris/graph/ir.py                   |   1 +
 tests/test_2d_tiling.py                  | 118 ++++++++++++++++-
 tests/test_feasibility.py                |  15 ++-
 4 files changed, 268 insertions(+), 19 deletions(-)

diff --git a/src/tigris/analysis/partition_spatial.py b/src/tigris/analysis/partition_spatial.py
index 57da144..532e29d 100644
--- a/src/tigris/analysis/partition_spatial.py
+++ b/src/tigris/analysis/partition_spatial.py
@@ -11,7 +11,7 @@
 import math
 from enum import Enum
 
-from tigris import TILE_AXIS_HEIGHT_OR_LENGTH, TILE_AXIS_NONE
+from tigris import TILE_AXIS_HEIGHT_OR_LENGTH, TILE_AXIS_HW, TILE_AXIS_NONE
 from tigris.graph.ir import (
     AnalyzedGraph,
     OpNode,
@@ -163,6 +163,74 @@ def _get_dilation_w(op: OpNode) -> int:
 # Tile solver
 
 
+def solve_2d_tile(
+    budget: int,
+    peak: int,
+    input_h: int,
+    input_w: int,
+    halo_h: int,
+    halo_w: int,
+) -> tuple[int, int] | None:
+    """Largest square-ish (tile_h, tile_w) whose proportional working set fits budget.
+
+    Mirrors the 1D proportional model already used for HEIGHT_OR_LENGTH in
+    partition_spatial(): ``tile_h = floor(budget*input_h/peak) - halo`` and
+    ``tiled_peak = int(peak * (tile_h + halo) / input_h)``. The stage's whole
+    peak_bytes (weights/bias/scratch included, which stay resident whole and
+    do not scale with tile size) is scaled by the fraction of the haloed
+    input area the tile covers:
+
+        tiled_peak(th, tw) = int(peak * (th + halo_h) * (tw + halo_w)
+                                  / (input_h * input_w))
+
+    Returns None if even a 1x1 core tile does not fit.
+    """
+    if input_h <= 0 or input_w <= 0:
+        return None
+
+    def tiled_peak(th: int, tw: int) -> int:
+        return int(peak * (th + halo_h) * (tw + halo_w) / (input_h * input_w))
+
+    if tiled_peak(1, 1) > budget:
+        return None
+
+    th = tw = max(min(input_h, input_w), 1)
+
+    # Shrink the larger side first, keeping the core roughly square, until it fits.
+    while tiled_peak(th, tw) > budget:
+        if th >= tw and th > 1:
+            th -= 1
+        elif tw > 1:
+            tw -= 1
+        else:
+            th = tw = 1
+            break
+
+    th = min(th, input_h)
+    tw = min(tw, input_w)
+    return (max(th, 1), max(tw, 1))
+
+
+def _stage_2d_eligible(
+    ag: AnalyzedGraph, stage: Stage, stage_ops: list[OpNode]
+) -> bool:
+    """A stage may attempt HW tiling only if it is a standalone rank-4 stage
+    with at most one spatial op, all of whose ops implement the HW tile
+    contract. Multi-spatial-op stages are excluded: the runtime executor's
+    2D contract is audited only for a single composed spatial op per stage.
+    """
+    if _stage_io_ranks(ag, stage) != {4}:
+        return False
+    spatial_count = sum(
+        1
+        for op in stage_ops
+        if classify_op(op.op_type) in (TileCategory.CONV, TileCategory.POOL)
+    )
+    if spatial_count > 1:
+        return False
+    return all(_op_supports_axis(op, TILE_AXIS_HW) for op in stage_ops)
+
+
 def partition_spatial(ag: AnalyzedGraph) -> AnalyzedGraph:
     """Analyze each stage and attach a TilePlan where needed.
 
@@ -203,7 +271,7 @@ def partition_spatial(ag: AnalyzedGraph) -> AnalyzedGraph:
             continue
 
         # Compute receptive field
-        rf_h, _rf_w = compute_receptive_field(stage_ops)
+        rf_h, rf_w = compute_receptive_field(stage_ops)
         halo = rf_h - 1
 
         # Axis 1 in the serialized NHWC/NLC layout maps to H/L at source dim 2.
@@ -225,6 +293,51 @@ def partition_spatial(ag: AnalyzedGraph) -> AnalyzedGraph:
         # Estimate tiled peak memory
         tiled_peak = int(peak * (tile_h + halo) / input_h)
 
+        # A single-row height tile that still overflows the budget cannot be
+        # rescued by any smaller height-only tile: height is already at its
+        # floor. If the stage is eligible for 2D (HW) tiling, try shrinking
+        # both axes together before falling back to the 1D infeasible
+        # warning below.
+        if (
+            tile_h == 1
+            and tiled_peak > budget
+            and _stage_2d_eligible(ag, stage, stage_ops)
+        ):
+            input_w = _find_input_extent_width(ag, stage)
+            if input_w > 0:
+                halo_w = rf_w - 1
+                shape = solve_2d_tile(
+                    budget=budget,
+                    peak=peak,
+                    input_h=input_h,
+                    input_w=input_w,
+                    halo_h=halo,
+                    halo_w=halo_w,
+                )
+                if shape is not None:
+                    tile_h_2d, tile_w_2d = shape
+                    tiled_peak_2d = int(
+                        peak
+                        * (tile_h_2d + halo)
+                        * (tile_w_2d + halo_w)
+                        / (input_h * input_w)
+                    )
+                    stage.tile_plan = TilePlan(
+                        tileable=True,
+                        axis=TILE_AXIS_HW,
+                        tile_height=tile_h_2d,
+                        tile_width=tile_w_2d,
+                        num_tiles=math.ceil(input_h / tile_h_2d)
+                        * math.ceil(input_w / tile_w_2d),
+                        halo=halo,
+                        receptive_field=rf_h,
+                        original_height=input_h,
+                        tiled_peak_bytes=tiled_peak_2d,
+                        overhead_bytes=0,
+                        warnings=[],
+                    )
+                    continue
+
         # Overhead: extra halo reads per tile boundary
         # Each internal tile boundary reads halo rows extra from the input
         halo_tensor_bytes = _estimate_halo_bytes(ag, stage, halo, input_h)
@@ -294,17 +407,37 @@ def _stage_tile_axis(
 
 def _op_supports_axis(op: OpNode, axis: int) -> bool:
     """Fail closed unless an operator implements the selected tile contract."""
-    if axis != TILE_AXIS_HEIGHT_OR_LENGTH:
-        return False
-    if op.op_type == "Conv1D":
-        return True
-    return op.op_type in _OP_CATEGORY
+    if axis == TILE_AXIS_HEIGHT_OR_LENGTH:
+        if op.op_type == "Conv1D":
+            return True
+        return op.op_type in _OP_CATEGORY
+    if axis == TILE_AXIS_HW:
+        # Rank-4 spatial/pointwise/channel-Concat set only; Conv1D is rank-3
+        # and has no width axis to tile.
+        return op.op_type in _OP_CATEGORY
+    return False
 
 
 def _find_input_extent(ag: AnalyzedGraph, stage: Stage, axis: int) -> int:
     """Find the H/L extent that serializes as axis 1 (source NCHW/NCL dim 2)."""
-    if axis != TILE_AXIS_HEIGHT_OR_LENGTH:
+    if axis not in (TILE_AXIS_HEIGHT_OR_LENGTH, TILE_AXIS_HW):
         return 0
+    return _find_source_dim_extent(ag, stage, dim=2, ranks={3, 4})
+
+
+def _find_input_extent_width(ag: AnalyzedGraph, stage: Stage) -> int:
+    """Find the W extent that serializes as axis 2 (source NCHW dim 3).
+
+    Only rank-4 tensors carry a width dimension; HW tiling never applies to
+    the rank-3 NCL layout.
+    """
+    return _find_source_dim_extent(ag, stage, dim=3, ranks={4})
+
+
+def _find_source_dim_extent(
+    ag: AnalyzedGraph, stage: Stage, dim: int, ranks: set[int]
+) -> int:
+    """Find a stage's source-shape extent at ``dim`` among candidate tensors."""
     # Check stage input tensors first, then look at first op's inputs
     candidates = stage.input_tensors.copy()
     if not candidates:
@@ -313,8 +446,8 @@ def _find_input_extent(ag: AnalyzedGraph, stage: Stage, axis: int) -> int:
 
     for name in candidates:
         info = ag.tensors.get(name)
-        if info and len(info.shape) in {3, 4}:
-            return int(info.shape[2])  # NCHW/NCL -> serialized H/L is dim 1
+        if info and len(info.shape) in ranks:
+            return int(info.shape[dim])
 
     return 0
 
diff --git a/src/tigris/graph/ir.py b/src/tigris/graph/ir.py
index b65accc..e306f36 100644
--- a/src/tigris/graph/ir.py
+++ b/src/tigris/graph/ir.py
@@ -99,6 +99,7 @@ class TilePlan:
     tileable: bool
     axis: int = TILE_AXIS_NONE
     tile_height: int = 0
+    tile_width: int = 0
     num_tiles: int = 0
     halo: int = 0
     receptive_field: int = 1
diff --git a/tests/test_2d_tiling.py b/tests/test_2d_tiling.py
index 85faeb3..ce6f996 100644
--- a/tests/test_2d_tiling.py
+++ b/tests/test_2d_tiling.py
@@ -1,8 +1,10 @@
-"""Tests for 2D (height + width) receptive field computation."""
+"""Tests for 2D (height + width) receptive field computation and tile solving."""
 
-from tigris import TILE_AXIS_HW
-from tigris.analysis.partition_spatial import compute_receptive_field
-from tigris.graph.ir import OpNode
+from onnx import TensorProto
+
+from tigris import TILE_AXIS_HEIGHT_OR_LENGTH, TILE_AXIS_HW
+from tigris.analysis.partition_spatial import compute_receptive_field, partition_spatial, solve_2d_tile
+from tigris.graph.ir import AnalyzedGraph, MemoryBudget, OpNode, Stage, TensorInfo
 
 
 def make_conv_op(kernel, stride, dilation):
@@ -36,3 +38,111 @@ def test_receptive_field_asymmetric_kernel():
     ops = [make_conv_op(kernel=(5, 3), stride=(2, 1), dilation=(1, 1))]
     rf_h, rf_w = compute_receptive_field(ops)
     assert (rf_h, rf_w) == (5, 3)
+
+
+# 2D tile solver
+#
+# solve_2d_tile mirrors the 1D proportional model already used by
+# partition_spatial() for HEIGHT_OR_LENGTH: tiled_peak(th, tw) scales the
+# stage's whole peak_bytes (which includes weights/bias/scratch that stay
+# resident regardless of tile shape) by the fraction of the haloed input
+# area the tile covers. This is deliberately conservative rather than the
+# activation-only "(th+hh)(tw+hw)*C_in + th*tw*C_out" formula, because that
+# formula omits resident non-activation memory and would let a 2D shape
+# through that overflows the fast arena at runtime.
+
+
+def test_solve_2d_tile_fits_budget():
+    # peak_bytes models a [512,512] int8 activation with C=256 baked in
+    # (512*512*256 = 67,108,864). Halo 2x2. A full 1-row tile does not fit
+    # 32K, but a small square core does.
+    peak_bytes = 512 * 512 * 256
+    shape = solve_2d_tile(
+        budget=32 * 1024,
+        peak=peak_bytes,
+        input_h=512,
+        input_w=512,
+        halo_h=2,
+        halo_w=2,
+    )
+    assert shape is not None
+    th, tw = shape
+    assert th >= 1 and tw >= 1
+    tiled_peak = int(peak_bytes * (th + 2) * (tw + 2) / (512 * 512))
+    assert tiled_peak <= 32 * 1024
+
+
+def test_solve_2d_tile_infeasible_returns_none():
+    # peak_bytes models a [64,64] int8 activation with C=4096 baked in.
+    # Even the 1x1 core (with 2x2 halo) overflows a 1K budget.
+    peak_bytes = 64 * 64 * 4096
+    assert (
+        solve_2d_tile(
+            budget=1024,
+            peak=peak_bytes,
+            input_h=64,
+            input_w=64,
+            halo_h=2,
+            halo_w=2,
+        )
+        is None
+    )
+
+
+# Axis selection: HW only kicks in once the 1D height solve is infeasible
+# even at a single-row tile.
+
+
+def build_high_res_conv_graph(h: int, w: int, c: int, mem_budget: int) -> AnalyzedGraph:
+    """Build a single-stage AnalyzedGraph for one 3x3 stride-1 Conv on an
+    int8 [1,c,h,w] activation.
+
+    peak_bytes is the whole-tensor working set (channels included), so a
+    single-row height tile costs exactly c*w bytes, matching the 1D
+    solver's proportional row-byte model and the comments below.
+    """
+    op = OpNode(
+        name="conv",
+        op_type="Conv",
+        inputs=["input"],
+        outputs=["output"],
+        attrs={"kernel_shape": [3, 3], "strides": [1, 1], "dilations": [1, 1]},
+    )
+    stage = Stage(
+        stage_id=0,
+        op_indices=[0],
+        input_tensors=["input"],
+        output_tensors=["output"],
+        peak_bytes=c * h * w,
+    )
+    return AnalyzedGraph(
+        ops=[op],
+        stages=[stage],
+        tensors={
+            "input": TensorInfo("input", (1, c, h, w), TensorProto.INT8),
+            "output": TensorInfo("output", (1, c, h, w), TensorProto.INT8),
+        },
+        budget=MemoryBudget(fast=mem_budget),
+    )
+
+
+def plan_for_single_stage(ag: AnalyzedGraph):
+    assert len(ag.stages) == 1
+    return ag.stages[0].tile_plan
+
+
+def test_axis_is_hw_only_when_height_infeasible():
+    # [1,256,256] row = 64K; a single row does not fit a 24K budget.
+    ag = build_high_res_conv_graph(h=256, w=256, c=256, mem_budget=24 * 1024)
+    ag = partition_spatial(ag)
+    stage_plan = plan_for_single_stage(ag)
+    assert stage_plan.axis == TILE_AXIS_HW
+    assert stage_plan.tileable is True
+
+
+def test_axis_stays_height_when_1d_feasible():
+    # [1,64,32] row = 2K, fits a 64K budget without needing HW tiling.
+    ag = build_high_res_conv_graph(h=256, w=64, c=32, mem_budget=64 * 1024)
+    ag = partition_spatial(ag)
+    stage_plan = plan_for_single_stage(ag)
+    assert stage_plan.axis == TILE_AXIS_HEIGHT_OR_LENGTH
diff --git a/tests/test_feasibility.py b/tests/test_feasibility.py
index 1aa65b5..d45b041 100644
--- a/tests/test_feasibility.py
+++ b/tests/test_feasibility.py
@@ -102,7 +102,10 @@ def test_compile_refuses_unsafe_height_tiling_but_keeps_untiled_support(
 
 
 def test_minimum_tile_that_exceeds_budget_is_infeasible(conv_relu_chain_path):
-    graph, _ = _run_pipeline(str(conv_relu_chain_path), ("1K",))
+    # 300 bytes is below the 1x1 2D core for either stage (halo 2x2, so a
+    # 2D tile solve is attempted once the 1D single-row tile also overflows;
+    # it correctly reports infeasible here too, not just the 1D fallback).
+    graph, _ = _run_pipeline(str(conv_relu_chain_path), ("300",))
 
     validation = validate_memory_plan(graph)
 
@@ -112,7 +115,7 @@ def test_minimum_tile_that_exceeds_budget_is_infeasible(conv_relu_chain_path):
 
 
 def test_findings_never_pass_when_scheduled_peak_exceeds_budget(conv_relu_chain_path):
-    graph, _ = _run_pipeline(str(conv_relu_chain_path), ("1K",))
+    graph, _ = _run_pipeline(str(conv_relu_chain_path), ("300",))
 
     findings = compute_findings(graph)
 
@@ -128,7 +131,9 @@ def test_compile_refuses_infeasible_plan_without_creating_output(
 
     result = CliRunner().invoke(
         cli,
-        ["compile", str(conv_relu_chain_path), "-m", "1K", "-o", str(output)],
+        # 300 bytes stays below the 1x1 2D core too; see
+        # test_minimum_tile_that_exceeds_budget_is_infeasible.
+        ["compile", str(conv_relu_chain_path), "-m", "300", "-o", str(output)],
     )
 
     assert result.exit_code != 0
@@ -175,7 +180,7 @@ def test_compile_rejects_budget_above_plan_format_limit(
 
 
 def test_writer_defensively_rejects_infeasible_graph(conv_relu_chain_path):
-    graph, _ = _run_pipeline(str(conv_relu_chain_path), ("1K",))
+    graph, _ = _run_pipeline(str(conv_relu_chain_path), ("300",))
 
     with pytest.raises(ValueError, match="Cannot emit an infeasible memory plan"):
         emit_binary_bytes(graph)
@@ -206,7 +211,7 @@ def test_feasible_plan_still_compiles(conv_relu_chain_path, tmp_path):
 def test_analyze_displays_failing_verdict(conv_relu_chain_path):
     result = CliRunner().invoke(
         cli,
-        ["analyze", str(conv_relu_chain_path), "-m", "1K"],
+        ["analyze", str(conv_relu_chain_path), "-m", "300"],
     )
 
     assert result.exit_code == 0

From 9fd61ff043ed929ba22a7a0e0bedeb923bf14231 Mon Sep 17 00:00:00 2001
From: asteinh <armin.steinhauser@raws.at>
Date: Sat, 15 Aug 2026 14:13:45 +0200
Subject: [PATCH 3/7] fix: exclude Conv1D from HW tile axis and correct
 peak-bytes docstring

---
 src/tigris/analysis/partition_spatial.py | 19 ++++++++++----
 tests/test_2d_tiling.py                  | 33 +++++++++++++++++++-----
 2 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/src/tigris/analysis/partition_spatial.py b/src/tigris/analysis/partition_spatial.py
index 532e29d..611760d 100644
--- a/src/tigris/analysis/partition_spatial.py
+++ b/src/tigris/analysis/partition_spatial.py
@@ -175,14 +175,19 @@ def solve_2d_tile(
 
     Mirrors the 1D proportional model already used for HEIGHT_OR_LENGTH in
     partition_spatial(): ``tile_h = floor(budget*input_h/peak) - halo`` and
-    ``tiled_peak = int(peak * (tile_h + halo) / input_h)``. The stage's whole
-    peak_bytes (weights/bias/scratch included, which stay resident whole and
-    do not scale with tile size) is scaled by the fraction of the haloed
-    input area the tile covers:
+    ``tiled_peak = int(peak * (tile_h + halo) / input_h)``. ``peak`` is the
+    stage's activation-only peak_bytes (compute_lifetimes skips constant
+    tensors, so weights/bias are never counted there); it is scaled by the
+    fraction of the haloed input area the tile covers:
 
         tiled_peak(th, tw) = int(peak * (th + halo_h) * (tw + halo_w)
                                   / (input_h * input_w))
 
+    Neither this solver nor the 1D one models resident weight/scratch bytes
+    against the tiled budget; that is a known pre-existing gap, and the
+    runtime backstops it by validating the emitted tile shape against the
+    real fast arena and failing closed.
+
     Returns None if even a 1x1 core tile does not fit.
     """
     if input_h <= 0 or input_w <= 0:
@@ -413,7 +418,11 @@ def _op_supports_axis(op: OpNode, axis: int) -> bool:
         return op.op_type in _OP_CATEGORY
     if axis == TILE_AXIS_HW:
         # Rank-4 spatial/pointwise/channel-Concat set only; Conv1D is rank-3
-        # and has no width axis to tile.
+        # and has no width axis to tile, even though it shares the CONV
+        # category with Conv in _OP_CATEGORY, so it must be excluded here
+        # explicitly rather than relying on the membership check alone.
+        if op.op_type == "Conv1D":
+            return False
         return op.op_type in _OP_CATEGORY
     return False
 
diff --git a/tests/test_2d_tiling.py b/tests/test_2d_tiling.py
index ce6f996..5c6b663 100644
--- a/tests/test_2d_tiling.py
+++ b/tests/test_2d_tiling.py
@@ -3,7 +3,12 @@
 from onnx import TensorProto
 
 from tigris import TILE_AXIS_HEIGHT_OR_LENGTH, TILE_AXIS_HW
-from tigris.analysis.partition_spatial import compute_receptive_field, partition_spatial, solve_2d_tile
+from tigris.analysis.partition_spatial import (
+    _op_supports_axis,
+    compute_receptive_field,
+    partition_spatial,
+    solve_2d_tile,
+)
 from tigris.graph.ir import AnalyzedGraph, MemoryBudget, OpNode, Stage, TensorInfo
 
 
@@ -44,12 +49,13 @@ def test_receptive_field_asymmetric_kernel():
 #
 # solve_2d_tile mirrors the 1D proportional model already used by
 # partition_spatial() for HEIGHT_OR_LENGTH: tiled_peak(th, tw) scales the
-# stage's whole peak_bytes (which includes weights/bias/scratch that stay
-# resident regardless of tile shape) by the fraction of the haloed input
-# area the tile covers. This is deliberately conservative rather than the
-# activation-only "(th+hh)(tw+hw)*C_in + th*tw*C_out" formula, because that
-# formula omits resident non-activation memory and would let a 2D shape
-# through that overflows the fast arena at runtime.
+# stage's activation-only peak_bytes (compute_lifetimes skips constant
+# tensors, so this never includes weights/bias) by the fraction of the
+# haloed input area the tile covers, instead of the brief's
+# "(th+hh)(tw+hw)*C_in + th*tw*C_out" formula. Resident weight/scratch bytes
+# are not modeled against the tiled budget by either the 1D or 2D solver (a
+# known pre-existing gap); the runtime backstops this by validating the
+# emitted tile shape against the real fast arena and failing closed.
 
 
 def test_solve_2d_tile_fits_budget():
@@ -89,6 +95,19 @@ def test_solve_2d_tile_infeasible_returns_none():
     )
 
 
+# Op eligibility for the HW axis. Conv1D shares the CONV category with Conv
+# in _OP_CATEGORY, but it is rank-3 and has no width axis to tile, so it
+# must be excluded from TILE_AXIS_HW explicitly rather than relying on the
+# _OP_CATEGORY membership check alone.
+
+
+def test_op_supports_axis_excludes_conv1d_for_hw():
+    conv1d = OpNode(name="c1d", op_type="Conv1D", inputs=[], outputs=[], attrs={})
+    conv2d = OpNode(name="c2d", op_type="Conv", inputs=[], outputs=[], attrs={})
+    assert _op_supports_axis(conv1d, TILE_AXIS_HW) is False
+    assert _op_supports_axis(conv2d, TILE_AXIS_HW) is True
+
+
 # Axis selection: HW only kicks in once the 1D height solve is infeasible
 # even at a single-row tile.
 

From 12410f8d49cf9e0b609e80000fb40525ca445707 Mon Sep 17 00:00:00 2001
From: asteinh <armin.steinhauser@raws.at>
Date: Sat, 15 Aug 2026 14:32:29 +0200
Subject: [PATCH 4/7] feature: serialize HW tile axis and tile_width, 2D
 feasibility diagnostic

---
 src/tigris/analysis/partition_spatial.py |  15 ++-
 src/tigris/analysis/validation.py        |   2 +
 src/tigris/emitters/binary/reader.py     |  19 ++--
 src/tigris/emitters/binary/writer.py     |   7 +-
 src/tigris/graph/ir.py                   |   5 +
 tests/test_2d_tiling_plan.py             | 120 +++++++++++++++++++++++
 tests/test_feasibility.py                |   9 +-
 7 files changed, 162 insertions(+), 15 deletions(-)
 create mode 100644 tests/test_2d_tiling_plan.py

diff --git a/src/tigris/analysis/partition_spatial.py b/src/tigris/analysis/partition_spatial.py
index 611760d..0d9505a 100644
--- a/src/tigris/analysis/partition_spatial.py
+++ b/src/tigris/analysis/partition_spatial.py
@@ -303,6 +303,7 @@ def partition_spatial(ag: AnalyzedGraph) -> AnalyzedGraph:
         # floor. If the stage is eligible for 2D (HW) tiling, try shrinking
         # both axes together before falling back to the 1D infeasible
         # warning below.
+        min_2d_tile_infeasible = False
         if (
             tile_h == 1
             and tiled_peak > budget
@@ -343,13 +344,24 @@ def partition_spatial(ag: AnalyzedGraph) -> AnalyzedGraph:
                     )
                     continue
 
+                # solve_2d_tile was attempted and even a 1x1 core tile does
+                # not fit the budget. Mark this stage distinctly so the
+                # surfaced diagnostic names the 2D tile instead of falling
+                # back to the generic 1D minimum-tile message below.
+                min_2d_tile_infeasible = True
+
         # Overhead: extra halo reads per tile boundary
         # Each internal tile boundary reads halo rows extra from the input
         halo_tensor_bytes = _estimate_halo_bytes(ag, stage, halo, input_h)
         overhead = halo_tensor_bytes * max(num_tiles - 1, 0)
 
         warnings: list[str] = []
-        if tiled_peak > budget:
+        if min_2d_tile_infeasible:
+            warnings.append(
+                f"Stage {stage.stage_id} minimum 2D tile still exceeds "
+                f"budget ({budget:,} bytes)"
+            )
+        elif tiled_peak > budget:
             warnings.append(
                 f"Stage {stage.stage_id} tiled peak ({tiled_peak:,} bytes) "
                 f"still exceeds budget ({budget:,} bytes)"
@@ -366,6 +378,7 @@ def partition_spatial(ag: AnalyzedGraph) -> AnalyzedGraph:
             tiled_peak_bytes=tiled_peak,
             overhead_bytes=overhead,
             warnings=warnings,
+            min_2d_tile_infeasible=min_2d_tile_infeasible,
         )
 
     return ag
diff --git a/src/tigris/analysis/validation.py b/src/tigris/analysis/validation.py
index 269c34e..6605444 100644
--- a/src/tigris/analysis/validation.py
+++ b/src/tigris/analysis/validation.py
@@ -423,6 +423,8 @@ def _execution_unit_requirement(
             return stage.peak_bytes, reason, True
         if tile_plan.tiled_peak_bytes <= 0:
             return stage.peak_bytes, "tile solver produced no positive working set", True
+        if tile_plan.min_2d_tile_infeasible:
+            return tile_plan.tiled_peak_bytes, "minimum 2D tile", False
         return tile_plan.tiled_peak_bytes, "minimum spatial tile", False
 
     return stage.peak_bytes, "untiled stage", False
diff --git a/src/tigris/emitters/binary/reader.py b/src/tigris/emitters/binary/reader.py
index d439838..0259033 100644
--- a/src/tigris/emitters/binary/reader.py
+++ b/src/tigris/emitters/binary/reader.py
@@ -7,6 +7,7 @@
     SCHEMA_VERSION_STAGE_TABLE_AUTHORITY,
     SUPPORTED_SCHEMA_VERSIONS,
     TILE_AXIS_HEIGHT_OR_LENGTH,
+    TILE_AXIS_HW,
     TILE_AXIS_NONE,
 )
 
@@ -250,17 +251,21 @@ def _read_shape(shape_off: int, ndim: int) -> list[int]:
             tileable, axis, tile_height,
             n_tiles, halo,
             rf, orig_h,
-            tiled_peak, overhead, _reserved,
+            tiled_peak, overhead, reserved,
         ) = TILE_PLAN_STRUCT.unpack_from(data, pos)
+        decoded_axis = (
+            axis
+            if version >= SCHEMA_VERSION_TILE_AXIS
+            # Schema v2-v4 used zero here and implicitly meant NHWC height.
+            else (TILE_AXIS_HEIGHT_OR_LENGTH if tileable else TILE_AXIS_NONE)
+        )
         tile_plans.append({
             "tileable": bool(tileable),
-            # Schema v2-v4 used zero here and implicitly meant NHWC height.
-            "axis": (
-                axis
-                if version >= SCHEMA_VERSION_TILE_AXIS
-                else (TILE_AXIS_HEIGHT_OR_LENGTH if tileable else TILE_AXIS_NONE)
-            ),
+            "axis": decoded_axis,
             "tile_height": tile_height,
+            # Packed into the low 16 bits of the trailing reserved u32; only
+            # meaningful for the HW axis, where the writer populates it.
+            "tile_width": reserved & 0xFFFF if decoded_axis == TILE_AXIS_HW else 0,
             "num_tiles": n_tiles,
             "halo": halo,
             "receptive_field": rf,
diff --git a/src/tigris/emitters/binary/writer.py b/src/tigris/emitters/binary/writer.py
index 1bdc8cb..609089d 100644
--- a/src/tigris/emitters/binary/writer.py
+++ b/src/tigris/emitters/binary/writer.py
@@ -9,6 +9,7 @@
 from tigris import (
     SCHEMA_VERSION,
     TILE_AXIS_HEIGHT_OR_LENGTH,
+    TILE_AXIS_HW,
     TILE_AXIS_NONE,
 )
 from tigris.graph.ir import AnalyzedGraph, OpNode
@@ -971,7 +972,7 @@ def _build_tile_plans(ag: AnalyzedGraph) -> tuple[bytes, dict[int, int]]:
 
         stage_to_tile[stage.stage_id] = idx
         idx += 1
-        if tp.tileable and tp.axis != TILE_AXIS_HEIGHT_OR_LENGTH:
+        if tp.tileable and tp.axis not in (TILE_AXIS_HEIGHT_OR_LENGTH, TILE_AXIS_HW):
             raise ValueError(
                 f"stage {stage.stage_id} has unsupported tile axis {tp.axis}"
             )
@@ -986,7 +987,7 @@ def _build_tile_plans(ag: AnalyzedGraph) -> tuple[bytes, dict[int, int]]:
         # receptive_field(u16) original_height(u16)
         # tiled_peak_bytes(u32)
         # overhead_bytes(u32)
-        # reserved(u32)
+        # reserved(u32): tile_width packed into the low 16 bits
         buf.extend(TILE_PLAN_STRUCT.pack(
             1 if tp.tileable else 0,
             tp.axis,
@@ -997,7 +998,7 @@ def _build_tile_plans(ag: AnalyzedGraph) -> tuple[bytes, dict[int, int]]:
             tp.original_height,
             tp.tiled_peak_bytes,
             tp.overhead_bytes,
-            0,  # reserved
+            tp.tile_width & 0xFFFF,
         ))
 
     return bytes(buf), stage_to_tile
diff --git a/src/tigris/graph/ir.py b/src/tigris/graph/ir.py
index e306f36..798a251 100644
--- a/src/tigris/graph/ir.py
+++ b/src/tigris/graph/ir.py
@@ -108,6 +108,11 @@ class TilePlan:
     overhead_bytes: int = 0
     untileable_ops: list[str] = field(default_factory=list)
     warnings: list[str] = field(default_factory=list)
+    # Set when a stage was eligible for 2D (HW) tiling and solve_2d_tile
+    # attempted the search but even a 1x1 core tile did not fit the budget.
+    # Distinguishes this case from the generic 1D minimum-tile shortfall so
+    # the compiler can surface a diagnostic naming the 2D tile explicitly.
+    min_2d_tile_infeasible: bool = False
 
 
 @dataclass
diff --git a/tests/test_2d_tiling_plan.py b/tests/test_2d_tiling_plan.py
new file mode 100644
index 0000000..4b79e3e
--- /dev/null
+++ b/tests/test_2d_tiling_plan.py
@@ -0,0 +1,120 @@
+"""Compiler serializes the HW tile axis and tile_width, and refuses to emit
+a plan when even a 1x1 2D core tile does not fit the fast-memory budget.
+
+Byte-level, mirrors test_linebuffer_plan.py: builds a real ONNX model, runs
+the shared CLI pipeline, and inspects the emitted binary tile-plan record.
+"""
+
+import struct
+
+import numpy as np
+import onnx
+import pytest
+from onnx import TensorProto, helper
+
+from tigris import TILE_AXIS_HW
+from tigris.analysis.validation import validate_memory_plan
+from tigris.cli import _run_pipeline
+from tigris.emitters.binary.defs import (
+    HEADER_STRUCT,
+    SEC_TILE_PLANS,
+    SECTION_ENTRY_STRUCT,
+)
+from tigris.emitters.binary.writer import emit_binary_bytes
+
+TILE_PLAN_STRUCT = struct.Struct("<BBHHHHHIII")  # matches tigris_tile_plan_t (24 bytes)
+
+
+def _build_high_res_conv(h: int, w: int, c: int) -> onnx.ModelProto:
+    """A single 3x3 stride-1 pad-1 Conv on a [1,c,h,w] float32 activation.
+
+    Same channel count in and out so the stage's live-tensor peak is exactly
+    two copies of the [1,c,h,w] tensor, matching the proportional tile model
+    used by solve_2d_tile.
+    """
+    X = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, c, h, w])
+    Y = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, c, h, w])
+    w0 = helper.make_tensor(
+        "w0", TensorProto.FLOAT, [c, c, 3, 3],
+        np.zeros((c, c, 3, 3), dtype=np.float32).flatten().tolist(),
+    )
+    b0 = helper.make_tensor("b0", TensorProto.FLOAT, [c], np.zeros(c, dtype=np.float32).tolist())
+    conv0 = helper.make_node(
+        "Conv", ["input", "w0", "b0"], ["output"], name="conv0",
+        kernel_shape=[3, 3], strides=[1, 1], pads=[1, 1, 1, 1],
+    )
+    graph = helper.make_graph([conv0], "high_res_conv", [X], [Y], initializer=[w0, b0])
+    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)])
+    model.ir_version = 8
+    return model
+
+
+def compile_high_res_conv_to_bytes(h: int, w: int, c: int, budget: str, tmp_path) -> bytes:
+    """Build the graph, run the shared pipeline, and emit the binary plan bytes.
+
+    Follows the crossrepo _compile_plan pattern: validate feasibility before
+    emitting so an infeasible plan raises with a descriptive reason instead
+    of being silently written out.
+    """
+    model_path = tmp_path / "high_res_conv.onnx"
+    onnx.save(_build_high_res_conv(h, w, c), str(model_path))
+    ag, _ = _run_pipeline(str(model_path), (budget,))
+    validation = validate_memory_plan(ag)
+    if not validation.feasible:
+        details = "; ".join(issue.describe() for issue in validation.issues)
+        raise AssertionError(f"compiler produced an infeasible graph: {details}")
+    return emit_binary_bytes(ag)
+
+
+def decode_first_tile_plan(plan_bytes: bytes):
+    """Locate the tile-plan section and unpack the first record."""
+    header = HEADER_STRUCT.unpack_from(plan_bytes, 0)
+    section_dir_off = header[3]
+
+    sections: dict[int, int] = {}
+    off = section_dir_off
+    while off + SECTION_ENTRY_STRUCT.size <= len(plan_bytes):
+        sec_type, sec_off = SECTION_ENTRY_STRUCT.unpack_from(plan_bytes, off)
+        off += SECTION_ENTRY_STRUCT.size
+        if sec_type == 0:
+            break
+        sections[sec_type] = sec_off
+
+    tp_base = sections[SEC_TILE_PLANS]
+    (
+        tileable, axis, tile_height,
+        num_tiles, halo,
+        receptive_field, original_height,
+        tiled_peak_bytes, overhead_bytes, reserved,
+    ) = TILE_PLAN_STRUCT.unpack_from(plan_bytes, tp_base)
+
+    class _TilePlan:
+        pass
+
+    tp = _TilePlan()
+    tp.tileable = tileable
+    tp.axis = axis
+    tp.tile_height = tile_height
+    tp.tile_width = reserved & 0xFFFF
+    tp.num_tiles = num_tiles
+    tp.halo = halo
+    tp.receptive_field = receptive_field
+    tp.original_height = original_height
+    tp.tiled_peak_bytes = tiled_peak_bytes
+    tp.overhead_bytes = overhead_bytes
+    return tp
+
+
+def test_2d_plan_encodes_hw_axis_and_tile_width(tmp_path):
+    plan_bytes = compile_high_res_conv_to_bytes(h=256, w=256, c=256, budget="24K", tmp_path=tmp_path)
+    tp = decode_first_tile_plan(plan_bytes)
+    assert tp.tileable == 1
+    assert tp.axis == TILE_AXIS_HW
+    assert tp.tile_height >= 1 and tp.tile_width >= 1
+
+
+def test_even_2d_min_tile_infeasible_fails_closed(tmp_path):
+    # Budget so small even a 1x1 2D core will not fit -> compile refuses.
+    with pytest.raises(Exception) as exc:
+        compile_high_res_conv_to_bytes(h=64, w=64, c=4096, budget="1K", tmp_path=tmp_path)
+    assert "2D tile" in str(exc.value)
diff --git a/tests/test_feasibility.py b/tests/test_feasibility.py
index d45b041..050cbad 100644
--- a/tests/test_feasibility.py
+++ b/tests/test_feasibility.py
@@ -102,16 +102,17 @@ def test_compile_refuses_unsafe_height_tiling_but_keeps_untiled_support(
 
 
 def test_minimum_tile_that_exceeds_budget_is_infeasible(conv_relu_chain_path):
-    # 300 bytes is below the 1x1 2D core for either stage (halo 2x2, so a
-    # 2D tile solve is attempted once the 1D single-row tile also overflows;
-    # it correctly reports infeasible here too, not just the 1D fallback).
+    # 300 bytes is below the 1x1 2D core for either stage (halo 2x2), so a
+    # 2D tile solve is attempted once the 1D single-row tile also overflows
+    # and reports its own distinct "minimum 2D tile" reason, not the generic
+    # 1D "minimum spatial tile" fallback.
     graph, _ = _run_pipeline(str(conv_relu_chain_path), ("300",))
 
     validation = validate_memory_plan(graph)
 
     assert not validation.feasible
     assert validation.scheduled_peak_bytes > graph.mem_budget
-    assert any(issue.reason == "minimum spatial tile" for issue in validation.issues)
+    assert any(issue.reason == "minimum 2D tile" for issue in validation.issues)
 
 
 def test_findings_never_pass_when_scheduled_peak_exceeds_budget(conv_relu_chain_path):

From f2e5a21456ebe068abeaa73719fd83ca4223ece4 Mon Sep 17 00:00:00 2001
From: asteinh <armin.steinhauser@raws.at>
Date: Sat, 15 Aug 2026 16:07:39 +0200
Subject: [PATCH 5/7] test: 2D tiled conv bit-exact vs ORT gate (float, int8,
 pointwise-wrapped)

---
 scripts/crossrepo_contract.py | 301 ++++++++++++++++++++++++++++++++++
 1 file changed, 301 insertions(+)

diff --git a/scripts/crossrepo_contract.py b/scripts/crossrepo_contract.py
index faf7abb..971f20c 100644
--- a/scripts/crossrepo_contract.py
+++ b/scripts/crossrepo_contract.py
@@ -11,6 +11,7 @@
 
 import argparse
 import copy
+import math
 import re
 import subprocess
 import sys
@@ -26,6 +27,7 @@
 from numpy.typing import NDArray
 from onnx import TensorProto, helper, numpy_helper
 
+from tigris import TILE_AXIS_HW
 from tigris.analysis.validation import validate_memory_plan
 from tigris.capabilities import KERNEL_CAPABILITIES, OP_TYPE_BY_CODE
 from tigris.cli import _run_pipeline
@@ -43,6 +45,7 @@
 # plan test-suite; reuse it rather than duplicating the stage-record parser.
 sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "tests"))
 from test_linebuffer_plan import _stage_reserved1  # noqa: E402
+from test_2d_tiling_plan import decode_first_tile_plan  # noqa: E402
 
 
 Array = NDArray[np.generic]
@@ -68,6 +71,7 @@ class ContractCase:
     expect_tiled: bool = False
     expect_chain: bool = False
     expect_line_buffered: bool = False
+    expect_2d: bool = False
     recompute_metric: bool = False
     force_one_op_stages: bool = False
 
@@ -1153,6 +1157,275 @@ def _qdq_conv_chain_case() -> ContractCase:
     )
 
 
+def build_conv(
+    *,
+    n: int,
+    c_in: int,
+    c_out: int,
+    h: int,
+    w: int,
+    kernel: int,
+    stride: int,
+    pad: int,
+    seed: int = 0,
+) -> tuple[onnx.ModelProto, onnx.ModelProto, dict[str, Array]]:
+    """A single float32 Conv on an NCHW activation.
+
+    The compile model and the ORT reference model are identical (as in
+    _tiled_pool_case and _dilated_conv_case); only random weights and a
+    uniform input distinguish instances at different resolutions.
+    """
+    out_h = (h + 2 * pad - kernel) // stride + 1
+    out_w = (w + 2 * pad - kernel) // stride + 1
+    model_input = helper.make_tensor_value_info(
+        "input", TensorProto.FLOAT, [n, c_in, h, w]
+    )
+    model_output = helper.make_tensor_value_info(
+        "output", TensorProto.FLOAT, [n, c_out, out_h, out_w]
+    )
+    rng = np.random.default_rng(seed)
+    weights = numpy_helper.from_array(
+        rng.normal(0.0, 0.05, size=(c_out, c_in, kernel, kernel)).astype(
+            np.float32
+        ),
+        "weights",
+    )
+    bias = numpy_helper.from_array(
+        rng.normal(0.0, 0.05, size=(c_out,)).astype(np.float32), "bias"
+    )
+    model = _model(
+        "conv2d",
+        [
+            helper.make_node(
+                "Conv",
+                ["input", "weights", "bias"],
+                ["output"],
+                name="conv0",
+                kernel_shape=[kernel, kernel],
+                strides=[stride, stride],
+                pads=[pad, pad, pad, pad],
+            )
+        ],
+        [model_input],
+        [model_output],
+        [weights, bias],
+    )
+    inputs = {
+        "input": rng.uniform(-1.0, 1.0, size=(n, c_in, h, w)).astype(
+            np.float32
+        )
+    }
+    return model, model, inputs
+
+
+def _2d_tiled_conv_case() -> ContractCase:
+    """A high-res Conv whose 1D height-only tile is infeasible at the budget
+    but a 2D (H and W) tile fits.
+
+    Input NCHW [1, 64, 66, 66], a 3x3 stride-1 pad-1 Conv to [1, 64, 66, 66],
+    at a 24K fast budget. 64 float32 channels give the same 256 bytes per
+    pixel as the int8 sibling's 256 channels, so the compiler solves the same
+    4x5 core tile. 66 is not divisible by 4 or 5, so the last row, the last
+    column, and the bottom-right corner tile are all partial.
+
+    Height-only tiling is infeasible first: partition_spatial only attempts
+    the 2D solve after its 1D tile_h == 1 candidate still exceeds budget, so
+    an emitted axis == TILE_AXIS_HW plan is itself proof the 1D path failed
+    closed at this budget.
+    """
+    compile_model, reference_model, inputs = build_conv(
+        n=1, c_in=64, c_out=64, h=66, w=66, kernel=3, stride=1, pad=1
+    )
+    return ContractCase(
+        "float_2d_tiled_conv",
+        compile_model,
+        reference_model,
+        inputs,
+        expected_operators=("Conv",),
+        mem_budget="24K",
+        expect_tiled=True,
+        expect_2d=True,
+    )
+
+
+def _qdq_2d_tiled_conv_case() -> ContractCase:
+    """The int8 sibling of _2d_tiled_conv_case, built via the _qdq_case QDQ
+    pattern: input NCHW [1, 256, 66, 66], a 3x3 stride-1 pad-1 Conv, 24K
+    budget. 256 int8 channels give the same per-pixel byte footprint as the
+    float case's 64 float32 channels, so the compiler solves the same 4x5
+    2D core tile with the same partial last row, column, and corner.
+    """
+    h = w = 66
+    c = 256
+    kernel, stride, pad = 3, 1, 1
+    out_h = (h + 2 * pad - kernel) // stride + 1
+    out_w = (w + 2 * pad - kernel) // stride + 1
+    model_input = helper.make_tensor_value_info(
+        "input", TensorProto.FLOAT, [1, c, h, w]
+    )
+    model_output = helper.make_tensor_value_info(
+        "output", TensorProto.FLOAT, [1, c, out_h, out_w]
+    )
+    int8_output = helper.make_tensor_value_info(
+        "output_q", TensorProto.INT8, [1, c, out_h, out_w]
+    )
+
+    rng = np.random.default_rng(3)
+    input_scale = numpy_helper.from_array(
+        np.array([0.02], dtype=np.float32), "input_scale"
+    )
+    input_zero_point = numpy_helper.from_array(
+        np.array([0], dtype=np.int8), "input_zero_point"
+    )
+    output_scale = numpy_helper.from_array(
+        np.array([0.05], dtype=np.float32), "output_scale"
+    )
+    output_zero_point = numpy_helper.from_array(
+        np.array([0], dtype=np.int8), "output_zero_point"
+    )
+    weight = numpy_helper.from_array(
+        rng.normal(0.0, 0.05, size=(c, c, kernel, kernel)).astype(
+            np.float32
+        ),
+        "weight",
+    )
+    weight_scale = numpy_helper.from_array(
+        np.array([0.01], dtype=np.float32), "weight_scale"
+    )
+    weight_zero_point = numpy_helper.from_array(
+        np.array([0], dtype=np.int8), "weight_zero_point"
+    )
+    initializers = [
+        input_scale,
+        input_zero_point,
+        output_scale,
+        output_zero_point,
+        weight,
+        weight_scale,
+        weight_zero_point,
+    ]
+    nodes = [
+        helper.make_node(
+            "QuantizeLinear",
+            ["input", "input_scale", "input_zero_point"],
+            ["input_q"],
+        ),
+        helper.make_node(
+            "DequantizeLinear",
+            ["input_q", "input_scale", "input_zero_point"],
+            ["input_dq"],
+        ),
+        helper.make_node(
+            "QuantizeLinear",
+            ["weight", "weight_scale", "weight_zero_point"],
+            ["weight_q"],
+        ),
+        helper.make_node(
+            "DequantizeLinear",
+            ["weight_q", "weight_scale", "weight_zero_point"],
+            ["weight_dq"],
+        ),
+        helper.make_node(
+            "Conv",
+            ["input_dq", "weight_dq"],
+            ["raw"],
+            name="conv0",
+            kernel_shape=[kernel, kernel],
+            strides=[stride, stride],
+            pads=[pad, pad, pad, pad],
+        ),
+        helper.make_node(
+            "QuantizeLinear",
+            ["raw", "output_scale", "output_zero_point"],
+            ["output_q"],
+        ),
+        helper.make_node(
+            "DequantizeLinear",
+            ["output_q", "output_scale", "output_zero_point"],
+            ["output"],
+        ),
+    ]
+    compile_model = _model(
+        "qdq_2d_tiled_conv", nodes, [model_input], [model_output], initializers
+    )
+    reference_model = copy.deepcopy(compile_model)
+    del reference_model.graph.output[:]
+    reference_model.graph.output.extend([int8_output])
+    onnx.checker.check_model(reference_model)
+
+    input_data = rng.uniform(-1.0, 1.0, size=(1, c, h, w)).astype(np.float32)
+    return ContractCase(
+        "int8_2d_tiled_conv",
+        compile_model,
+        reference_model,
+        {"input": input_data},
+        ("Conv",),
+        mem_budget="24K",
+        expect_tiled=True,
+        expect_2d=True,
+    )
+
+
+def _2d_tiled_conv_sigmoid_case() -> ContractCase:
+    """Conv followed by a non-fused pointwise Sigmoid, both forced 2D at the
+    same 24K/66x66/64-channel geometry as _2d_tiled_conv_case.
+
+    Relu/Relu6 fuse into the Conv at compile time, so this uses Sigmoid to
+    keep the pointwise op a standalone stage. At this budget the row-based
+    (full-width) chain streamer cannot fit even one row, so Conv and Sigmoid
+    stay as two independent stages, each solving its own 4x5 HW tile; the
+    Sigmoid stage exercises the 2D executor running a pointwise op on a
+    packed (non-full-width) tile, closing the Task 7 coverage gap.
+    """
+    h = w = 66
+    c = 64
+    kernel, stride, pad = 3, 1, 1
+    rng = np.random.default_rng(5)
+    model_input = helper.make_tensor_value_info(
+        "input", TensorProto.FLOAT, [1, c, h, w]
+    )
+    model_output = helper.make_tensor_value_info(
+        "output", TensorProto.FLOAT, [1, c, h, w]
+    )
+    weights = numpy_helper.from_array(
+        rng.normal(0.0, 0.05, size=(c, c, kernel, kernel)).astype(
+            np.float32
+        ),
+        "weights",
+    )
+    bias = numpy_helper.from_array(
+        rng.normal(0.0, 0.05, size=(c,)).astype(np.float32), "bias"
+    )
+    nodes = [
+        helper.make_node(
+            "Conv",
+            ["input", "weights", "bias"],
+            ["conv_out"],
+            name="conv0",
+            kernel_shape=[kernel, kernel],
+            strides=[stride, stride],
+            pads=[pad, pad, pad, pad],
+        ),
+        helper.make_node("Sigmoid", ["conv_out"], ["output"]),
+    ]
+    model = _model(
+        "conv_sigmoid_2d", nodes, [model_input], [model_output], [weights, bias]
+    )
+    inputs = {
+        "input": rng.uniform(-1.0, 1.0, size=(1, c, h, w)).astype(np.float32)
+    }
+    return ContractCase(
+        "float_2d_tiled_conv_sigmoid",
+        model,
+        model,
+        inputs,
+        expected_operators=("Conv", "Sigmoid"),
+        mem_budget="24K",
+        expect_tiled=True,
+        expect_2d=True,
+    )
+
+
 def _to_runtime_layout(value: Array) -> Array:
     if value.ndim == 4:
         return np.ascontiguousarray(value.transpose(0, 2, 3, 1))
@@ -1593,6 +1866,31 @@ def _assert_plan_mode(case: ContractCase, plan: dict) -> None:
             f"{case.name}: line_buffered={line_buffered}, expected "
             f"{case.expect_line_buffered}"
         )
+
+    # A 2D case decodes the plan's first tile-plan record straight from the
+    # emitted bytes (reusing test_2d_tiling_plan.py's decoder) and confirms
+    # both axes actually split into more than one tile. num_tiles is the
+    # solver's ceil(H/tile_h) * ceil(W/tile_w) product, so dividing it by the
+    # H-axis tile count derived from the decoded original_height/tile_height
+    # recovers the W-axis tile count without needing a separate width field
+    # in the plan format.
+    if case.expect_2d:
+        tile_plan = decode_first_tile_plan(plan_bytes)
+        if tile_plan.axis != TILE_AXIS_HW:
+            raise AssertionError(
+                f"{case.name}: tile plan axis {tile_plan.axis}, expected "
+                f"TILE_AXIS_HW ({TILE_AXIS_HW})"
+            )
+        tiles_h = math.ceil(tile_plan.original_height / tile_plan.tile_height)
+        tiles_w = tile_plan.num_tiles // tiles_h
+        if not (tiles_h > 1 and tiles_w > 1):
+            raise AssertionError(
+                f"{case.name}: expected multi-tile on both axes, got "
+                f"tiles_h={tiles_h} tiles_w={tiles_w} "
+                f"(num_tiles={tile_plan.num_tiles}, "
+                f"tile_h={tile_plan.tile_height}, tile_w={tile_plan.tile_width})"
+            )
+
     if case.compression == "lz4":
         if plan["weight_blocks_compression"] != COMPRESS_LZ4:
             raise AssertionError(f"{case.name}: expected LZ4 weight blocks")
@@ -1763,6 +2061,9 @@ def _run_gate(runtime: Path, work_dir: Path) -> None:
         _qdq_case("AveragePool"),
         _linebuffer_conv_chain_case(),
         _qdq_conv_chain_case(),
+        _2d_tiled_conv_case(),
+        _qdq_2d_tiled_conv_case(),
+        _2d_tiled_conv_sigmoid_case(),
     ]
     covered_operators = {
         operator for case in cases for operator in case.expected_operators

From 45fefd0a41f11679516951f8da898bbe8469f710 Mon Sep 17 00:00:00 2001
From: asteinh <armin.steinhauser@raws.at>
Date: Sat, 15 Aug 2026 16:13:57 +0200
Subject: [PATCH 6/7] fix: exclude binary-op stages from 2D tile eligibility

---
 src/tigris/analysis/partition_spatial.py |  35 ++++--
 tests/test_2d_tiling.py                  | 130 +++++++++++++++++++++++
 2 files changed, 159 insertions(+), 6 deletions(-)

diff --git a/src/tigris/analysis/partition_spatial.py b/src/tigris/analysis/partition_spatial.py
index 0d9505a..a07f8d8 100644
--- a/src/tigris/analysis/partition_spatial.py
+++ b/src/tigris/analysis/partition_spatial.py
@@ -52,6 +52,16 @@ class TileCategory(Enum):
 }
 
 
+# Dynamic binary ops (Add, Mul) whose second operand is an independent
+# tensor. Neither the rank-3 axis-1 executor nor the rank-4 HW (2D) executor
+# guarantees that operand is co-tiled with a spatial op's output: both load
+# every stage input using the spatial op's own tile geometry (length stripe
+# or HW halo rectangle), so a binary op combined with a spatial op can read
+# the wrong region or size from its second operand. Safe only in stages with
+# no spatial op. Shared between the rank-3 and rank-4 eligibility checks
+# below since the hazard is the same in both.
+_BINARY_OPS = frozenset({"Add", "Mul"})
+
 # Rank-3 NLC stages have a deliberately narrower axis-1 contract than rank-4
 # NHWC stages.  Unary pointwise operators preserve the current length and may
 # surround one Conv1D.  Dynamic binary operators are safe only in a
@@ -63,8 +73,7 @@ class TileCategory(Enum):
     "Sigmoid",
     "Tanh",
 })
-_RANK3_AXIS1_BINARY_OPS = frozenset({"Add", "Mul"})
-_RANK3_AXIS1_OPS = _RANK3_AXIS1_UNARY_OPS | _RANK3_AXIS1_BINARY_OPS | {"Conv1D"}
+_RANK3_AXIS1_OPS = _RANK3_AXIS1_UNARY_OPS | _BINARY_OPS | {"Conv1D"}
 
 
 def classify_op(op_type: str) -> TileCategory:
@@ -220,9 +229,21 @@ def _stage_2d_eligible(
     ag: AnalyzedGraph, stage: Stage, stage_ops: list[OpNode]
 ) -> bool:
     """A stage may attempt HW tiling only if it is a standalone rank-4 stage
-    with at most one spatial op, all of whose ops implement the HW tile
-    contract. Multi-spatial-op stages are excluded: the runtime executor's
-    2D contract is audited only for a single composed spatial op per stage.
+    with at most one spatial op, no binary op, and all of whose ops implement
+    the HW tile contract. Multi-spatial-op stages are excluded: the runtime
+    executor's 2D contract is audited only for a single composed spatial op
+    per stage.
+
+    Binary ops (Add, Mul) are excluded even though they pass the per-op HW
+    contract check below: exec_stage_tiled_2d loads every stage input using
+    the same conv input-halo rectangle, so a binary op's second operand
+    (e.g. a residual skip tensor) is not guaranteed to be co-tiled with the
+    spatial op's output at that rectangle. Admitting a stage like
+    [Conv, Add(conv_out, skip)] as 2D would load the skip operand with the
+    wrong region and size and silently produce a wrong result. This is the
+    conservative, fail-closed choice: such a stage falls back to the 1D path
+    (or fails closed if 1D is also infeasible). Co-tiled-binary 2D tiling is
+    a future refinement, not attempted here.
     """
     if _stage_io_ranks(ag, stage) != {4}:
         return False
@@ -233,6 +254,8 @@ def _stage_2d_eligible(
     )
     if spatial_count > 1:
         return False
+    if any(op.op_type in _BINARY_OPS for op in stage_ops):
+        return False
     return all(_op_supports_axis(op, TILE_AXIS_HW) for op in stage_ops)
 
 
@@ -414,7 +437,7 @@ def _stage_tile_axis(
             and not (
                 conv_count == 1
                 and any(
-                    op_type in _RANK3_AXIS1_BINARY_OPS
+                    op_type in _BINARY_OPS
                     for op_type in op_types
                 )
             )
diff --git a/tests/test_2d_tiling.py b/tests/test_2d_tiling.py
index 5c6b663..37963b5 100644
--- a/tests/test_2d_tiling.py
+++ b/tests/test_2d_tiling.py
@@ -5,6 +5,7 @@
 from tigris import TILE_AXIS_HEIGHT_OR_LENGTH, TILE_AXIS_HW
 from tigris.analysis.partition_spatial import (
     _op_supports_axis,
+    _stage_2d_eligible,
     compute_receptive_field,
     partition_spatial,
     solve_2d_tile,
@@ -165,3 +166,132 @@ def test_axis_stays_height_when_1d_feasible():
     ag = partition_spatial(ag)
     stage_plan = plan_for_single_stage(ag)
     assert stage_plan.axis == TILE_AXIS_HEIGHT_OR_LENGTH
+
+
+# 2D eligibility must exclude binary ops (Add, Mul).
+#
+# exec_stage_tiled_2d loads every stage input with the conv's own
+# input-halo rectangle. A binary op's second operand (e.g. a residual skip
+# tensor) is a separate full-resolution tensor that is not guaranteed to be
+# co-tiled with the conv's output at that rectangle, so a stage like
+# [Conv, Add(conv_out, skip)] admitted as 2D would load the skip operand
+# with the wrong region and size and silently produce a wrong result. This
+# mirrors the precedent already enforced on the rank-3 axis-1 path
+# (_BINARY_OPS / _stage_tile_axis), which forbids combining a Conv1D with
+# any binary op for the same co-tiling reason.
+
+
+def build_high_res_conv_add_graph(h: int, w: int, c: int, mem_budget: int) -> AnalyzedGraph:
+    """Single-stage graph: Conv followed by Add against a separate
+    full-resolution skip tensor, at the same [1,c,h,w] shape and peak_bytes
+    model as build_high_res_conv_graph, so the same budget that forces 2D
+    tiling for a bare Conv applies here too.
+    """
+    conv = OpNode(
+        name="conv",
+        op_type="Conv",
+        inputs=["input"],
+        outputs=["conv_out"],
+        attrs={"kernel_shape": [3, 3], "strides": [1, 1], "dilations": [1, 1]},
+    )
+    add = OpNode(
+        name="add",
+        op_type="Add",
+        inputs=["conv_out", "skip"],
+        outputs=["output"],
+        attrs={},
+    )
+    stage = Stage(
+        stage_id=0,
+        op_indices=[0, 1],
+        input_tensors=["input", "skip"],
+        output_tensors=["output"],
+        peak_bytes=c * h * w,
+    )
+    return AnalyzedGraph(
+        ops=[conv, add],
+        stages=[stage],
+        tensors={
+            "input": TensorInfo("input", (1, c, h, w), TensorProto.INT8),
+            "conv_out": TensorInfo("conv_out", (1, c, h, w), TensorProto.INT8),
+            "skip": TensorInfo("skip", (1, c, h, w), TensorProto.INT8),
+            "output": TensorInfo("output", (1, c, h, w), TensorProto.INT8),
+        },
+        budget=MemoryBudget(fast=mem_budget),
+    )
+
+
+def build_high_res_conv_sigmoid_graph(h: int, w: int, c: int, mem_budget: int) -> AnalyzedGraph:
+    """Single-stage graph: Conv followed by a unary Sigmoid wrapper, at the
+    same [1,c,h,w] shape and peak_bytes model as build_high_res_conv_graph.
+
+    Used as the control: a spatial op plus a UNARY pointwise wrapper must
+    still be admitted for 2D tiling, proving the binary-op guard is scoped
+    to binary ops only and does not over-exclude.
+    """
+    conv = OpNode(
+        name="conv",
+        op_type="Conv",
+        inputs=["input"],
+        outputs=["conv_out"],
+        attrs={"kernel_shape": [3, 3], "strides": [1, 1], "dilations": [1, 1]},
+    )
+    sigmoid = OpNode(
+        name="sigmoid",
+        op_type="Sigmoid",
+        inputs=["conv_out"],
+        outputs=["output"],
+        attrs={},
+    )
+    stage = Stage(
+        stage_id=0,
+        op_indices=[0, 1],
+        input_tensors=["input"],
+        output_tensors=["output"],
+        peak_bytes=c * h * w,
+    )
+    return AnalyzedGraph(
+        ops=[conv, sigmoid],
+        stages=[stage],
+        tensors={
+            "input": TensorInfo("input", (1, c, h, w), TensorProto.INT8),
+            "conv_out": TensorInfo("conv_out", (1, c, h, w), TensorProto.INT8),
+            "output": TensorInfo("output", (1, c, h, w), TensorProto.INT8),
+        },
+        budget=MemoryBudget(fast=mem_budget),
+    )
+
+
+def test_stage_2d_eligible_excludes_binary_op():
+    # Hand-built stage with a binary op: must not be 2D eligible.
+    ag = build_high_res_conv_add_graph(h=256, w=256, c=256, mem_budget=24 * 1024)
+    stage = ag.stages[0]
+    assert _stage_2d_eligible(ag, stage, ag.ops) is False
+
+
+def test_stage_2d_eligible_allows_spatial_plus_unary_op():
+    # Hand-built stage with only a spatial op plus a unary wrapper: must
+    # still be 2D eligible, proving the guard does not over-exclude.
+    ag = build_high_res_conv_sigmoid_graph(h=256, w=256, c=256, mem_budget=24 * 1024)
+    stage = ag.stages[0]
+    assert _stage_2d_eligible(ag, stage, ag.ops) is True
+
+
+def test_binary_op_stage_does_not_go_hw():
+    # End to end through partition_spatial(): same budget that forces a bare
+    # Conv stage to TILE_AXIS_HW (see test_axis_is_hw_only_when_height_infeasible)
+    # must NOT do so once a binary Add against an uncotiled skip tensor is
+    # in the stage. It falls back to the 1D height axis instead.
+    ag = build_high_res_conv_add_graph(h=256, w=256, c=256, mem_budget=24 * 1024)
+    ag = partition_spatial(ag)
+    stage_plan = plan_for_single_stage(ag)
+    assert stage_plan.axis != TILE_AXIS_HW
+
+
+def test_conv_plus_unary_stage_still_goes_hw():
+    # Control: the same budget still drives a spatial-plus-unary stage to
+    # TILE_AXIS_HW, proving the binary-op guard is scoped to binary ops only.
+    ag = build_high_res_conv_sigmoid_graph(h=256, w=256, c=256, mem_budget=24 * 1024)
+    ag = partition_spatial(ag)
+    stage_plan = plan_for_single_stage(ag)
+    assert stage_plan.axis == TILE_AXIS_HW

From e816b8a21b9114906f5130fae05d8b9d2253ed34 Mon Sep 17 00:00:00 2001
From: asteinh <armin.steinhauser@raws.at>
Date: Sat, 15 Aug 2026 16:32:46 +0200
Subject: [PATCH 7/7] fix: exclude Concat stages from 2D tile eligibility

---
 src/tigris/analysis/partition_spatial.py |  9 +++++
 tests/test_2d_tiling.py                  | 48 ++++++++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/src/tigris/analysis/partition_spatial.py b/src/tigris/analysis/partition_spatial.py
index a07f8d8..5bab278 100644
--- a/src/tigris/analysis/partition_spatial.py
+++ b/src/tigris/analysis/partition_spatial.py
@@ -256,6 +256,15 @@ def _stage_2d_eligible(
         return False
     if any(op.op_type in _BINARY_OPS for op in stage_ops):
         return False
+    # Concat carries the same hazard as the binary ops above: it takes an
+    # independent second operand, and exec_stage_tiled_2d loads every stage
+    # input with the spatial op's own input-halo rectangle, which is not
+    # guaranteed to be co-tiled with a distinct Concat operand at output
+    # resolution. Exclude it from HW eligibility for symmetry with _BINARY_OPS;
+    # such a stage falls back to the 1D path or fails closed. A Concat that is
+    # a stage head/fan-in is unaffected (it is not a single-spatial-op stage).
+    if any(op.op_type == "Concat" for op in stage_ops):
+        return False
     return all(_op_supports_axis(op, TILE_AXIS_HW) for op in stage_ops)
 
 
diff --git a/tests/test_2d_tiling.py b/tests/test_2d_tiling.py
index 37963b5..f9b8f3c 100644
--- a/tests/test_2d_tiling.py
+++ b/tests/test_2d_tiling.py
@@ -221,6 +221,46 @@ def build_high_res_conv_add_graph(h: int, w: int, c: int, mem_budget: int) -> An
     )
 
 
+def build_high_res_conv_concat_graph(h: int, w: int, c: int, mem_budget: int) -> AnalyzedGraph:
+    """Single-stage graph: Conv followed by a channel-axis Concat against a
+    separate full-resolution operand. Concat takes an independent second
+    operand the 2D executor's shared-rectangle load does not co-tile, so the
+    stage must not be 2D eligible (same hazard as the binary-op case).
+    """
+    conv = OpNode(
+        name="conv",
+        op_type="Conv",
+        inputs=["input"],
+        outputs=["conv_out"],
+        attrs={"kernel_shape": [3, 3], "strides": [1, 1], "dilations": [1, 1]},
+    )
+    concat = OpNode(
+        name="concat",
+        op_type="Concat",
+        inputs=["conv_out", "other"],
+        outputs=["output"],
+        attrs={"axis": 1},
+    )
+    stage = Stage(
+        stage_id=0,
+        op_indices=[0, 1],
+        input_tensors=["input", "other"],
+        output_tensors=["output"],
+        peak_bytes=c * h * w,
+    )
+    return AnalyzedGraph(
+        ops=[conv, concat],
+        stages=[stage],
+        tensors={
+            "input": TensorInfo("input", (1, c, h, w), TensorProto.INT8),
+            "conv_out": TensorInfo("conv_out", (1, c, h, w), TensorProto.INT8),
+            "other": TensorInfo("other", (1, c, h, w), TensorProto.INT8),
+            "output": TensorInfo("output", (1, 2 * c, h, w), TensorProto.INT8),
+        },
+        budget=MemoryBudget(fast=mem_budget),
+    )
+
+
 def build_high_res_conv_sigmoid_graph(h: int, w: int, c: int, mem_budget: int) -> AnalyzedGraph:
     """Single-stage graph: Conv followed by a unary Sigmoid wrapper, at the
     same [1,c,h,w] shape and peak_bytes model as build_high_res_conv_graph.
@@ -269,6 +309,14 @@ def test_stage_2d_eligible_excludes_binary_op():
     assert _stage_2d_eligible(ag, stage, ag.ops) is False
 
 
+def test_stage_2d_eligible_excludes_concat():
+    # A Conv+Concat stage carries the same non-co-tiled-operand hazard as a
+    # binary op, so it must not be 2D eligible either.
+    ag = build_high_res_conv_concat_graph(h=256, w=256, c=256, mem_budget=24 * 1024)
+    stage = ag.stages[0]
+    assert _stage_2d_eligible(ag, stage, ag.ops) is False
+
+
 def test_stage_2d_eligible_allows_spatial_plus_unary_op():
     # Hand-built stage with only a spatial op plus a unary wrapper: must
     # still be 2D eligible, proving the guard does not over-exclude.