Support padded shared-memory contraction stages - #538
Open
SamuelBelanger wants to merge 7 commits into
Open
Conversation
Gathered lane fan-out unrolls line reads into scalar rank-1 updates indexed by `line * lw + lane`. This linear indexing assumes `lw` divides the innermost contracted axis extent so that lane stepping never crosses reduction coordinate boundaries. Gate lane fan-out on `innermost_divisible` (`reduce_extents[last].is_multiple_of(lw)`), preserving the fan-out path for multi-axis reductions whose innermost axis divides `lw` (such as vectorized convolutions) while falling back to the coordinate-decoded flat walk when the innermost axis is not a multiple of `lw` (such as padded stages).
Contraction operands that cannot vectorize in global memory (such as NHWC tensors with C = 3 where rows are unaligned) can now be staged into padded shared-memory lines. Staging rounds the innermost physical axis up to whole lines, widens during the straight fill from scalar source reads, and spreads the vector lanes back across scalar sink cells upon committing accumulator blocks.
A stage rounds its innermost extent up to whole lines so a padded one can hold a partial last line, while the gmem side still counts its own lines by truncation. Where no padding is intended the two boxes disagree silently. Refuse it in fill_straight, where both are in scope: the innermost extent has to be a whole number of source lines, and only the destination may hold a partial one. An equal-width stage is back to the divisibility rule Compaction::line_extents used to carry, and a direct stage gains one it never had. Also narrow the Dynamic-extent check to the padded case that reads it, and assert widen_line's lane count against the line width it writes into.
A spread block column holds `spread` scalar sink cells in its lanes, and the N-D nest rounds its column count up, so the last column addresses cells past the sink whenever `spread` does not divide its innermost extent. Nothing masked them: an unchecked AccumulateView writes straight through, so the block wrote past the accumulator's window. The values survived (a padding lane seeds a cell, accumulates the stage's zero, and commits the same value back) but the write itself was out of bounds, and inf inputs poison the round trip. Pass the sink's extent to block::seed and block::commit and mask the lanes that fall past it. Only the spread path rounds its column count up now: an `aw`-wide column has no per-lane handle on a partial cell, so it keeps counting whole lines as it did. Skipping a lane puts commit's plane fold under divergent control flow, so refuse a lane-split accumulator there; AccumulateView::lane_share is what asks. Also assert what the spread arithmetic already assumed, that a spread block's accumulator is served scalar, next to the code that relies on it rather than two files away in step_served.
Two entry points set the width and only one checked it: the launcher's StridedTileSource::stage_width, while TileSpec::stage_width validated nothing at all, which is the path the conv tests take. Move the checks onto TileSpec::validate_stage_width and call it from both, at launch for a host backtrace and again in Tile::of for hand-built specs, beside the Boundary::Clamp assert that catches them for the same reason. Only the caller knows the served width and whether the operand is quantized, so both are passed in. Tighten the width rule while it moves. `width >= v && width.is_multiple_of(v)` admitted a vectorized operand, which then panicked deep in fill_from where widen_line asks for scalar source cells; a padded stage is for operands global memory cannot vectorize, so require v == 1 and a width that widens it. Fold Resize1d::check_padded back into the check helper it was copied from.
The divisibility rule Compaction::line_extents used to carry at construction moved to fill_straight, which is one of two fills. A windowed or masked stage takes scan_transparent instead, where a storage extent rounded up left a last line the source never fills and nothing noticed. Extract the check as fill_extent and ask it from both. Two more boxes were left disagreeing with the rounding: projected::line_extents still truncated, so a padded stage's bounds-check box excluded the partial line it really holds (inert only because smem is Overhang::Never, and silently zeroing every read the day it is not), and MemData::at rounded a region edge up for gmem too, where Tile::of still truncates the shape it indexes. The edge now has to be whole lines or the whole axis, since a partial tail line has no sibling to start after it. Bundle the padded fill's facts into one Padding value rather than three comptime numbers passed alongside each other. fill_lines and read_stage_line lose three parameters each, the 1:1 path no longer carries a rank it never reads, and widen_line asserts on its own width type instead of a number it was handed. The rank agreement Padding needs is asserted where it is built: a storage-tiled stage splits each axis in two and cannot be one.
The fan-out names its lane with a comptime extract, and `lane_component` decodes `last_k % lw` on the flat walk; what the gate has to ask is that the two agree. Divisibility of the fastest contracted extent is one way to get there, but not the only one: a single contracted axis makes `last_k` the flat step itself, so the two are the same expression whatever the extent, tail included. Gating on divisibility alone therefore demoted every single-axis gathered contraction whose K is not a whole number of lines to the flat walk. That is the case the fan-out's tail loop exists for, and the loop became unreachable along with it: divisibility of the fastest extent implies divisibility of their product, so `k_tail` was provably zero wherever it was read. Ask `reduce.len() == 1 || fastest.is_multiple_of(lw)` instead, which revives both. `assert_operand_shapes` has already refused an empty reduce by this point, so index the extents rather than carrying an arm for a case `k_axis_idx` would have underflowed on two lines later. Also round `batch_matrix`'s column count up, the last of the three line counts still truncating. A padded lhs reaches the 2-D nest (a scalar rhs and sink leave `rw == aw`, so `memory` routes to `direct`), where `block::contract`'s tail walks its partial line correctly but the view's box claimed zero columns. `cols` is a shape there, never a stride, so this only widens the bound.
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
1 participant
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.
Problem
A contraction operand whose innermost axis cannot be vectorized in global memory is stuck reading scalar all the way to the leaf. The motivating case is an
NHWCtensor withC = 3: no 4-wide line lands on a row boundary, so the operand serves one element at a time even though every level below global memory could handle whole lines.Approach
A shared-memory stage owns its own layout, so it does not have to be served at the width its source is read at.
stage_width(w)asks for a stage padded out tow-wide lines:Three pieces make that work.
Staging rounds the innermost physical axis up to whole lines, so
C = 3allocates one 4-wide line with the fourth lane as padding.The fill (
MemData::fill_straight) assembles each destination line from adjacent scalar source reads. Lanes at or past the innermost extent hold zero rather than reading the next row's first live value. Where the extent fills whole lines the guard is dropped at comptime.The register block spreads those lanes back across a scalar sink. Block column
ncovers sink columnsn·spread … n·spread + spread - 1, so a padded rhs can run 4-wide against an accumulator addressed one element at a time.served(lanes are K-partials of one cell) andspread(lanes are neighbouring sink cells) are mutually exclusive and asserted so.Padded stages are refused for quantized operands, and the padding does not leak past the level that asked for it:
StagePlan::descendclears it.Two fixes this uncovered
Both are pre-existing and land as their own commits.
Lane fan-out over multi-axis reductions. The fan-out walk indexes by
line * lw + lane, which assumeslwdivides the innermost contracted extent. It is now gated on that. Multi-axis reductions whose innermost axis divideslw(vectorized convolutions) keep the fan-out; the rest take the coordinate-decoded flat walk.Lane resolution on the flat walk.
lane_componentresolved theKcomponent from the flat stepp, which agrees with the real coordinate only whenlwdivides the innermost reduce extent. It now reads that coordinate directly.A ragged innermost extent on an unpadded stage. Rounding the stage up to whole lines is only sound when padding is intended; the global-memory side still counts its lines by truncation.
fill_straight, where both boxes are in scope, now refuses an extent that is not a whole number of source lines. An equal-width stage is back to the divisibility ruleCompaction::line_extentsused to carry, and a direct stage gains one it never had.Tests
resize1d_staged_padded_stage_widthCI = 3padded into 4-wide linesconv1d_staged_padded_multi_axis_reduce_lane_indexinglw ∤ CIconv1d_staged_padded_multi_axis_reduce_lane_fanoutmatmul_padded_rhs_stage_into_scalar_sinkN = 3padded rhs, scalarM×Nsinkmatmul_padded_rhs_stage_multi_lineN = 5, two lines with a 3-lane tailstage_width_*(4)cubek-tileis green apart fromtest_reduce_axis_sum_spatial_unit_lanes, which fails onmainfor unrelated reasons.