Co-tile pre-spatial same-resolution skip connections in 2D tiling - #32
Merged
Conversation
…s are not undercounted
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
Relaxes the 2D (height and width) tiling eligibility so a stage whose
Concat/Add/Mulis consumed BEFORE its single spatial op, withsame-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, wherepreviously 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/Mulkernels already consume the tile geometry - the cross-repodifferential gate proves the tiled result is bit-exact.
What changed
_stage_2d_eligible(analysis/partition_spatial.py): the blanketAdd/Mul/Concatexclusion is replaced by a co-tileability check. A stageis admitted only when it has one spatial op, every
Concat/Add/Mulisconsumed 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 nowcomputes 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).
scripts/crossrepo_contract.py): new co-tiledConcat-skip (float and int8) andAdd-skip 2D-tiled cases assert the emittedplan is
TILE_AXIS_HWwith both axes split and match ONNX Runtime bit-exact.Testing
float_cotiled_concat_2d,int8_cotiled_concat_2d,and
float_cotiled_add_2dcases reachTILE_AXIS_HW(tiles on both axes) andmatch ONNX Runtime (float allclose 1e-5; int8 at the established 1-LSB QDQ
tolerance); all existing cases pass.
controls (post-spatial, different-resolution, constant operand), and the
concurrent-liveness slow-pool peak for a long-lived skip and a multi-op stage.