2D (H x W) spatial tiling: axis, solver, serialization, and gate - #29
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 true 2D (height and width) spatial tiling. When a
stage's minimum 1D-height tile (a single output row) still exceeds the fast
budget, the compiler now solves a 2D
tile_h x tile_wshape with halo on bothaxes and emits it under a new
HWtile axis, so a high-resolution convolutionthat previously fail-closed becomes tileable and runs bit-exact. If even the
minimum 2D tile does not fit, compilation fails closed with a distinct
diagnostic.
Pairs with the runtime change in raws-labs/tigris-runtime#31, which executes the
2D plan. Merge tigris-runtime#31 first: the cross-repo gate below builds the
sibling runtime.
What changed
TILE_AXIS_HW = 3, a 2D receptive-field computation (both kernel axes), and atile solver that picks a feasible
(tile_h, tile_w)using the sameproportional model as the 1D path. 2D is attempted only as a fail-closed
rescue, when the 1D single-row tile is infeasible.
Binary ops (Add, Mul) and Concat are excluded: the runtime loads every stage
input with the spatial op's input rectangle, which is not guaranteed to be
co-tiled with an independent second operand, so admitting such a stage could
read the wrong region. This mirrors the existing rank-3 length-tiling guard.
axis = HWand packstile_widthinto the existingreserved field (no wire-format layout change); the reader decodes it. An
even-2D-infeasible stage produces a distinct "minimum 2D tile" compile error.
int8) compiled at a budget where 1D is infeasible but 2D fits, run through the
runtime and checked against ONNX Runtime, with the height and width chosen so
partial edge and corner tiles are exercised; plus a Conv-then-Sigmoid case for
the pointwise-wrapped path.
Testing
float_2d_tiled_conv,int8_2d_tiled_conv, andfloat_2d_tiled_conv_sigmoidcases match ONNXRuntime (float allclose 1e-5; int8 at the established 1-LSB QDQ tolerance),
with the emitted plan asserted to be
axis = HWand multi-tile on both axes.