From 372713c837ceda6cd5ac008b8dac81f0fe7917e7 Mon Sep 17 00:00:00 2001 From: asteinh Date: Mon, 17 Aug 2026 17:40:49 +0200 Subject: [PATCH] fix: fail closed on post-spatial external skip in 1D-height tiling to match the 2D path --- src/tigris/analysis/partition_spatial.py | 49 +++++++++++++++++------ tests/test_2d_tiling.py | 51 +++++++++++++++++++++++- 2 files changed, 87 insertions(+), 13 deletions(-) diff --git a/src/tigris/analysis/partition_spatial.py b/src/tigris/analysis/partition_spatial.py index 2808f5a..9d4ca98 100644 --- a/src/tigris/analysis/partition_spatial.py +++ b/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] diff --git a/tests/test_2d_tiling.py b/tests/test_2d_tiling.py index 2711667..307dc84 100644 --- a/tests/test_2d_tiling.py +++ b/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