Skip to content

Co-tile pre-spatial same-resolution skip connections in 2D tiling - #32

Merged
asteinh merged 6 commits into
developfrom
feature/cotiled-skip-2d
Aug 16, 2026
Merged

Co-tile pre-spatial same-resolution skip connections in 2D tiling#32
asteinh merged 6 commits into
developfrom
feature/cotiled-skip-2d

Conversation

@asteinh

@asteinh asteinh commented Aug 16, 2026

Copy link
Copy Markdown
Member

Summary

Relaxes the 2D (height and width) tiling eligibility so a stage whose
Concat/Add/Mul is consumed BEFORE its single spatial op, with
same-resolution stage-external operands (a skip connection), is admitted for 2D
tiling instead of being blanket-excluded. This lets an encoder-decoder block of
the form [Concat(up, skip), Conv] 2D-tile bit-exact against ONNX Runtime, where
previously any binary/Concat stage fell back to the 1D path or failed closed.

Compiler-only: no runtime change is required. The executor already loads a
pre-spatial same-resolution operand at the correct halo rectangle, and the
Concat/Add/Mul kernels already consume the tile geometry - the cross-repo
differential gate proves the tiled result is bit-exact.

What changed

  • _stage_2d_eligible (analysis/partition_spatial.py): the blanket
    Add/Mul/Concat exclusion is replaced by a co-tileability check. A stage
    is admitted only when it has one spatial op, every Concat/Add/Mul is
    consumed strictly before that spatial op (pre-spatial), and every stage-external
    operand is a rank-4 tensor at the same H/W as the op output. Post-spatial
    operands, different-resolution operands, and constant operands stay fail-closed.
  • slow_pool_usage (analysis/validation.py): the slow/PSRAM budget check now
    computes a per-tiled-stage interval-overlap concurrent-liveness peak. A boundary
    tensor is counted for a tiled stage when its lifetime overlaps the stage's
    op-step interval, so a stage's own inputs and outputs are counted concurrently
    and a long-lived skip resident across a stage is counted too. This replaces the
    per-stage input+output max, which under-counted a skip spanning multiple stages
    (and, if left as a naive per-op sample, would under-count a multi-op stage).
  • Cross-repo differential gate (scripts/crossrepo_contract.py): new co-tiled
    Concat-skip (float and int8) and Add-skip 2D-tiled cases assert the emitted
    plan is TILE_AXIS_HW with both axes split and match ONNX Runtime bit-exact.

Testing

  • pytest full suite: 335 passed, 2 skipped.
  • Cross-repo gate: the new float_cotiled_concat_2d, int8_cotiled_concat_2d,
    and float_cotiled_add_2d cases reach TILE_AXIS_HW (tiles on both axes) and
    match ONNX Runtime (float allclose 1e-5; int8 at the established 1-LSB QDQ
    tolerance); all existing cases pass.
  • New unit tests cover the admitted pre-spatial Concat/Add skip, the fail-closed
    controls (post-spatial, different-resolution, constant operand), and the
    concurrent-liveness slow-pool peak for a long-lived skip and a multi-op stage.

@asteinh
asteinh merged commit a19f882 into develop Aug 16, 2026
10 checks passed
@asteinh
asteinh deleted the feature/cotiled-skip-2d branch August 16, 2026 20:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant