fix: fail closed on post-spatial external skip in 1D-height tiling to match the 2D path

src/tigris/analysis/partition_spatial.py

@@ -263,6 +263,31 @@ def _cotileable_skip_operands(
     return True
 
 
+def _has_post_spatial_binary(stage_ops: list[OpNode]) -> bool:
+    """True if a Concat/Add/Mul is consumed at or after the stage's spatial op.
+
+    Both tiled executors (the 1D-height exec_stage_tiled and the 2D
+    exec_stage_tiled_2d) load every stage input at the spatial op's INPUT-halo
+    rectangle. A post-spatial binary/Concat operand lives at the spatial op's
+    OUTPUT resolution, so a strided spatial op would make the executor read that
+    operand at stride*out_start rows -- the wrong rows, and past the operand's
+    height -- on interior tiles. Both tile paths must fail closed on this shape
+    and run the stage untiled. Shared by _stage_2d_eligible and _stage_tile_axis
+    so the height and HW paths reject in lockstep.
+    """
+    spatial_idx = next(
+        (i for i, o in enumerate(stage_ops)
+         if classify_op(o.op_type) in (TileCategory.CONV, TileCategory.POOL)),
+        None,
+    )
+    if spatial_idx is None:
+        return False
+    return any(
+        (op.op_type in _BINARY_OPS or op.op_type == "Concat") and i >= spatial_idx
+        for i, op in enumerate(stage_ops)
+    )
+
+
 def _stage_2d_eligible(
     ag: AnalyzedGraph, stage: Stage, stage_ops: list[OpNode]
 ) -> bool:
@@ -293,22 +318,16 @@ def _stage_2d_eligible(
     )
     if spatial_count > 1:
         return False
-    # Locate the single spatial op (CONV/POOL); spatial_count <= 1 is ensured above.
-    spatial_idx = next(
-        (i for i, o in enumerate(stage_ops)
-         if classify_op(o.op_type) in (TileCategory.CONV, TileCategory.POOL)),
-        None,
-    )
     # Concat/Add/Mul are admitted only when consumed strictly BEFORE the spatial
     # op (so their operands sit at the spatial op's input resolution = the halo
     # rectangle the executor loads every stage input at) and every stage-external
-    # operand is a same-resolution skip. Post-spatial or different-resolution
-    # operands are the deferred cases: fail closed. op_indices/stage_ops are in
-    # topological (execution) order, so list index is dependency order.
-    for i, op in enumerate(stage_ops):
+    # operand is a same-resolution skip. Post-spatial operands fail closed via the
+    # shared _has_post_spatial_binary check; different-resolution or constant
+    # pre-spatial operands fail closed via _cotileable_skip_operands.
+    if _has_post_spatial_binary(stage_ops):
+        return False
+    for op in stage_ops:
         if op.op_type in _BINARY_OPS or op.op_type == "Concat":
-            if spatial_idx is not None and i >= spatial_idx:
-                return False
             if not _cotileable_skip_operands(ag, stage, op):
                 return False
     return all(_op_supports_axis(op, TILE_AXIS_HW) for op in stage_ops)
@@ -687,6 +706,12 @@ def _stage_tile_axis(
     """Select an audited serialized activation axis for a standalone stage."""
     ranks = _stage_io_ranks(ag, stage)
     if ranks == {4} and all(op.op_type != "Conv1D" for op in stage_ops):
+        # A post-spatial binary/Concat with a stage-external operand mis-tiles on
+        # the height path exactly as on the 2D path (the skip is loaded at the
+        # spatial op's input rows, not its output rows). Fail closed so the stage
+        # runs untiled via exec_stage_normal, matching _stage_2d_eligible.
+        if _has_post_spatial_binary(stage_ops):
+            return TILE_AXIS_NONE
         return TILE_AXIS_HEIGHT_OR_LENGTH
     if ranks == {3} and stage_ops:
         op_types = [op.op_type for op in stage_ops]

tests/test_2d_tiling.py

@@ -4,10 +4,11 @@
 
 from onnx import TensorProto
 
-from tigris import TILE_AXIS_HEIGHT_OR_LENGTH, TILE_AXIS_HW
+from tigris import TILE_AXIS_HEIGHT_OR_LENGTH, TILE_AXIS_HW, TILE_AXIS_NONE
 from tigris.analysis.partition_spatial import (
     _op_supports_axis,
     _stage_2d_eligible,
+    _stage_tile_axis,
     compute_receptive_field,
     partition_spatial,
     solve_2d_tile,
@@ -449,6 +450,54 @@ def build_high_res_concat_const_skip_graph(h, w, c, mem_budget):
         budget=MemoryBudget(fast=mem_budget))
 
 
+def build_post_spatial_add_skip_graph(h, w, c, mem_budget):
+    """Post-spatial: Conv(stride 2) -> Add(conv_out, skip). The Add is consumed
+    AFTER the strided Conv, and `skip` is a stage-external tensor at the Conv's
+    OUTPUT resolution. The tiled executors load every stage input at the Conv's
+    INPUT-halo rectangle, so a strided Conv would read `skip` at stride*out_start
+    rows (and past its height) on interior tiles. Both the 1D-height and the 2D
+    tile paths must fail closed on this shape (-> exec_stage_normal)."""
+    conv = OpNode(name="conv", op_type="Conv",
+                  inputs=["input"], outputs=["conv_out"],
+                  attrs={"kernel_shape": [1, 1], "strides": [2, 2], "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 // 2, w // 2), TensorProto.INT8),
+            "skip": TensorInfo("skip", (1, c, h // 2, w // 2), TensorProto.INT8),
+            "output": TensorInfo("output", (1, c, h // 2, w // 2), TensorProto.INT8),
+        },
+        budget=MemoryBudget(fast=mem_budget))
+
+
+def test_stage_tile_axis_rejects_post_spatial_external_skip():
+    # Conv(stride 2) -> Add(conv_out, external skip): the 1D-height executor loads
+    # the skip at the Conv's input rows, mis-tiling every interior tile (and
+    # reading past the skip's height). The axis selection must fail closed to
+    # TILE_AXIS_NONE so the stage runs untiled, matching the 2D path's rejection.
+    ag = build_post_spatial_add_skip_graph(h=256, w=256, c=64, mem_budget=24 * 1024)
+    assert _stage_tile_axis(ag, ag.stages[0], ag.ops) == TILE_AXIS_NONE
+
+
+def test_stage_2d_eligible_rejects_post_spatial_external_skip():
+    # The 2D path already rejects this shape; keep both paths in lockstep.
+    ag = build_post_spatial_add_skip_graph(h=256, w=256, c=64, mem_budget=24 * 1024)
+    assert _stage_2d_eligible(ag, ag.stages[0], ag.ops) is False
+
+
+def test_stage_tile_axis_admits_pre_spatial_add_skip():
+    # A pre-spatial Add(input, skip) -> Conv is co-tileable and loaded correctly
+    # at the input resolution, so the 1D-height axis stays admitted.
+    ag = build_high_res_add_conv_graph(h=256, w=256, c=64, mem_budget=24 * 1024)
+    assert _stage_tile_axis(ag, ag.stages[0], ag.ops) == TILE_AXIS_HEIGHT_OR_LENGTH
+
+
 def test_stage_2d_eligible_admits_pre_spatial_add_skip():
     ag = build_high_res_add_conv_graph(h=256, w=256, c=256, mem_budget=24 * 1024)
     assert _stage_2d_eligible(ag, ag.stages[0], ag.ops) is True