2D-tile ConvTranspose: dedicated output-extent solve sized to the runtime working set - #31
Merged
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
_solve_convtranspose_2d) handles ConvTranspose2D tiling. ConvTranspose deliberately stays
UNTILEABLEin the shared opcategory 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).
tensor), and
original_heightis emitted at output scale so the runtime's tilecount recovery stays consistent with its output-gridded execution.
a resident packed input tile back-computed from the output tile
(
(out_tile + eff_k + stride - 1)//stride + 2, clamped to the full input) plusevery 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.
partial-edge geometry) assert the emitted plan is
TILE_AXIS_HWwith both axessplit and that the runtime output matches ONNX Runtime bit-exact.
Testing
(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.
float_convtranspose_2d,int8_convtranspose_2d, andfloat_convtranspose_2d_partialcases match ONNX Runtime (float allclose 1e-5,int8 at the established 1-LSB QDQ tolerance); all existing cases pass.