Skip to content

2D-tile ConvTranspose: dedicated output-extent solve sized to the runtime working set - #31

Merged
asteinh merged 4 commits into
developfrom
feature/convtranspose-2d
Aug 16, 2026
Merged

2D-tile ConvTranspose: dedicated output-extent solve sized to the runtime working set#31
asteinh merged 4 commits into
developfrom
feature/convtranspose-2d

Conversation

@asteinh

@asteinh asteinh commented Aug 16, 2026

Copy link
Copy Markdown
Member

Summary

Adds the compiler side of 2D (height and width) spatial tiling for ConvTranspose.
When a ConvTranspose stage's peak activation exceeds the fast SRAM budget, the
compiler now emits a 2D (HW) tile plan gridded over the expanded output extent,
so the runtime can stream the stage tile by tile.

Pairs with the runtime change in raws-labs/tigris-runtime, which executes the
tiled plan. Merge tigris-runtime first.

What changed

  • A dedicated, isolated solve (_solve_convtranspose_2d) handles ConvTranspose
    2D tiling. ConvTranspose deliberately stays UNTILEABLE in the shared op
    category table, which keeps it out of the 1D-height solve, streamable chains,
    and receptive-field accumulation; the 2D plan is emitted only through this
    branch, for a stage whose single spatial op is a ConvTranspose (optionally with
    audited pointwise ops).
  • The tile grid is normalized over the OUTPUT extent (the expanded, peak-driving
    tensor), and original_height is emitted at output scale so the runtime's tile
    count recovery stays consistent with its output-gridded execution.
  • Tile sizing replicates the runtime's per-tile working-set model term for term:
    a resident packed input tile back-computed from the output tile
    ((out_tile + eff_k + stride - 1)//stride + 2, clamped to the full input) plus
    every op's packed output tile, each aligned. The solver searches for the largest
    output tile whose modeled working set fits the budget, so every emitted tile is
    one the runtime accepts. It fails closed (a non-tileable plan) when not even a
    1x1 output tile fits.
  • Cross-repo differential gate: 2D-tiled ConvTranspose cases (float, int8, and a
    partial-edge geometry) assert the emitted plan is TILE_AXIS_HW with both axes
    split and that the runtime output matches ONNX Runtime bit-exact.

Testing

  • pytest: full suite passes, including the new ConvTranspose 2D solve tests
    (over-budget tiles 2D, infeasible budget fails closed, under-budget stays
    untiled) and a test asserting the emitted tile's working set fits the budget
    under the runtime formula.
  • Cross-repo gate: the new float_convtranspose_2d, int8_convtranspose_2d, and
    float_convtranspose_2d_partial cases match ONNX Runtime (float allclose 1e-5,
    int8 at the established 1-LSB QDQ tolerance); all existing cases pass.

@asteinh
asteinh merged commit 7ee4b3c into develop Aug 16, 2026
10 checks passed
@asteinh
asteinh deleted the feature/convtranspose-2d branch August 16, 2026 15:00
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