Repository navigation
fix(codegen,runtime): per-level reduce planning and lowered-TIR memory views - #257
Conversation
zhen8838
left a comment
zhen8838
left a comment
There was a problem hiding this comment.
Thanks for tracking these three defects down, and for the compute-sanitizer evidence. The per-level reduce split and the TupleGetItem / nested-layout plumbing are what #255 asked for.
Two changes are needed before this can go in; the reasons are in the inline comments:
local_asbuilds a row-major index that CuTe reads in the other order, so the plain side of a one-sided copy is addressed through a permutation.tests/codegen/test_pointer_ref.pyshould be deleted in favour of executable tests.
zhen8838
left a comment
zhen8838
left a comment
There was a problem hiding this comment.
A follow-up on size: most of the new source code repeats information or logic the codebase already has. Three places can shrink; details are inline.
…ad meshes
ops::reduce derived its plan from a single per-mesh is_thread boolean,
and codegen flattened a composed cta×thread mesh — what every
multi-CTA finalize produces — into one cta-scoped Mesh. Every thread
axis was therefore classified cta-level: lanes_of=1, warps_of=extent
per axis, so warps_per_group became the full mesh instance count (256
instead of 8 in the RMSNorm tutorial). The cross-warp tier then
indexed the shared workspace at (group_start + w) * 32 + lane, far
past the allocation: OOB smem reads under compute-sanitizer. The
verifier multiplied the same flattened extents and demanded 4096
workspace slots where 8 suffice.
Classify each mesh axis by the level it belongs to:
- render_shard_layout_value emits one shape/stride nest per topology
level and names every scope in the Mesh type, so the runtime can
see the levels; single-level emission is byte-identical, and a
multi-level mesh that is also sliced is refused.
- reduce_dispatch spends the lane budget on the thread-level axes
only, fastest axis first; cta-level axes keep lanes_of=1,
warps_of=extent. Mesh gains one axis_scope<I>() member answering
which scope flat axis I runs at — a single-level mesh answers its
own scope for every axis, a multi-level mesh reads the nest
holding I off its layout — so the levels a mesh names are read off
the mesh itself, here and for any op that needs them next. The
out-of-scope refusal keeps its sentence ("a reduce mesh's scope
must be cta or thread"), now phrased per level.
- IntraCta iterates the kept cells (fold, warp butterfly, workspace
combine per cell) instead of asserting kept_cells == 1, and
cta_combine_via_workspace ends with a barrier so the
one-slot-per-warp workspace is safe to reuse across cells. The
refusal goes away with the assert: a source keeping two cells is a
supported reduce now, so the static-asserts suite drops the case
that pinned the old sentence.
- The verifier sizes the workspace from the thread level alone: each
CTA combines its own partials in its own shared memory.
Single-level thread and pure-cta plans are unchanged (compile-time
probe over the old and new paths). The new regression test reduces
over a cta×thread two-level mesh: it fails before this change (the
verifier demands 16 slots for a 4-warp CTA, and with verification
bypassed the dispatch goes cross-warp out of bounds) and passes
after, with compute-sanitizer reporting 0 errors.
…ayout ops::copy projected each operand through its own shard layout alone, so a plain operand paired with a sharded one stayed whole for every instance. T.copy(fragment, window) -- what finalize emits for the RMSNorm tutorial's reshard to gmem on its cta×thread mesh -- then ran cute::copy on mismatched sizes: the whole window iterated against the fragment's modular index space, every instance tiling its 40-element fragment across the 10240-element window and racing its peers. Where the vectorizer keeps the fragment's reads contiguous instead, the same mismatch steps past the fragment -- the tutorial's "Invalid __local__ read of size 2" under compute-sanitizer. When exactly one operand is sharded, the plain side moves the slice the shard side's layout hands this instance: - local_as resolves this instance's window and origin from the shard side's mesh coordinates. Window and origin are row-major logical steps over the shard layout's shape, while CuTe decodes an integer index first-mode-fastest, so the plain tensor's layout goes through the runtime's reverse (flattened first, so nested modes turn too; a ComposedLayout overload keeps a swizzle in front of the modes): CuTe's colexicographic index then is the row-major one, and the composition keeps whatever addressing the plain tensor states. - Copy dispatches on ShardTensorLike: exactly-one-sharded pairs take the projection; shard/shard and plain/plain pairs keep the per-operand local_tensor path unchanged. The regression test round-trips a (2, 5120) bf16 window through a (2, 4, 32)-instance rmem shard and back on one 256-thread CTA; a second test copies one way into a column-pair shard, so a slice placed wrong on the way in lands in the wrong output columns directly -- the round trip cannot see a permutation the write and the read share. The round trip fails with no projection (the window comes back tiled by whichever instance writes last), the one-way test also fails with the modes handed to CuTe unreversed (the output comes back transposed), and both pass after, bit for bit, with compute-sanitizer reporting 0 errors.
…_view More of the tile-ai#230 plumbing. _pointer_ref assumed a PtrOf source is always a named tensor: anything else fell through to _tensor_ref, which reads .name and died with "AttributeError: 'Call' object has no attribute 'name'" on the TupleGetItem stage choices WGMMA pipelines emit (tests/fixtures/schedule/tir/wgmma_rs_a_from_smem.py is the tracked reproducer) and on the TensorView-call and Constant sources nested under them. _plain_layout_value rendered shapes and strides as flat integer lists, so a layout with a nested mode died with "TypeError: int() argument must be ... not 'tuple'". Of the schedule TIR fixtures, twenty-one stopped at the first and three at the second. - TupleGetItem sources go through _tuple_pointer_ref: a constant index resolves to the chosen element's pointer; a runtime index renders through render_tuple_choice -- one cute::array of the elements subscripted by the index it renders, with .data() taking the chosen tensor's pointer, or an array of the pointers themselves when the elements name raw smem offsets. cute::array has no deduction guide, so the element type and count are spelled from the first element's rendering. A runtime choice has no single syntactic source, so no memory_var comes back for the coordinate forms to bind to. TensorView-call and Constant sources recurse instead of falling through. - One recursive renderer (codegen/cuda/tir/cute_layout.py) spells an int-tuple shape and stride in both forms -- the cute::Layout type nesting cute::Shape/cute::Stride, the make_layout value nesting make_shape/make_stride, with callables rendering the scalar entries -- and every layout emitter routes through it: tensor_view's layout type and value, its local_tile source layout, the shard layout's mesh positions, mesh_scope's levelwise types, and arith's broadcast view. Flat layouts render byte-identical to before. - Two emission fixes the executable runs pinned: fill rendered an integer constant as "0f", and a tuple index reading an rmem scalar named the whole rank-0 view where the arithmetic needs the held value; non-UMAT scalars now render subscripted. Every fixture that emitted still emits, and all twenty-four stopped ones now pass the pointer stage: eighteen reach the deliberate "CopyAsyncTensor: no emitter without host-encoded tensor maps" refusal (the missing TMA atoms, tile-ai#249), two newly emit, one reaches the sliced-multi-level-mesh refusal from the reduce warps fix, and the last three refuse in their own words -- swizzled slice bases and a nested shard local shape stay out of scope here. The two newly-emitting fixtures, staged_dynamic_inner and staged_nested, are re-authored as executable tests: their stage rings compile and run on the card, and the staged tiles are checked to arrive intact.
150984a to
666aae3
Compare
|
All five threads are addressed in the force-push (branch rewritten, still three commits; net +785/-155 vs the pre-review +617/-85 — the growth is the executable tests you asked for). Per-thread: copy_impl.h (permutation) — test_pointer_ref.py — Deleted. The two fixtures that newly emit ( reduce_common_impl.h (one Mesh helper) — One member, tensor_view.py (one recursive renderer) — New tensor_view.py (shared tuple choice) — Proactive dedup — we re-audited the diff for redundancy before replying and folded six consolidations into the same commits: Verification on the final tree (rebased onto main at 91ec3d7): scoped |
Why
composed cta x thread mesh every multi-CTA finalize produces:
ops::reduceclassified thread axes as cta-level, sowarps_per_groupbecame the full mesh instance count (256 instead of 8 in the RMSNorm
tutorial): cross-warp workspace indexing read smem out of bounds (3 x
Invalid __shared__ read of size 4under compute-sanitizer), and theverifier demanded 4096 slots where 8 suffice.
ops::copyprojected each operand through its own shard layout alone, so aplain window paired with a thread-sharded fragment stayed whole for every
instance:
cute::copyran mismatched domains, the tutorial's residualInvalid __local__ read of size 2(LDL.U16).TupleGetItemnornested layouts (fix(codegen): handle sliced pointers and nested layouts in lowered TIR #230): 24 schedule TIR fixtures (every WGMMA one among
them) stopped in
_tensor_ref/_plain_layout_value.What
reduce_dispatchclassifies each mesh axis by its topology level via one newMesh::axis_scope<I>()member (lane budget on thread-level axes only;per-level codegen layout nests, single level byte-identical,
multi-level+sliced refused;
IntraCtaiterates kept cells; verifier sizesthe workspace from the thread level alone).
ops::copy: when exactly one operand is sharded, the plain side moves theslice the shard side's layout hands this instance.
local_asresolves theinstance's window and origin from the shard side's mesh coordinates, then
reverses the plain layout's modes so CuTe's colexicographic integer index is
the row-major one — unreversed, a rank-2 plain tensor is visited through a
permutation a round trip cannot see.
codegen/cuda/tir/cute_layout.py, type and valueforms over int-tuple shape/stride) serves every
cute::make_layout(...)site (
tensor_view.py,mesh_scope.py,arith.py): nesting — multi-levelmeshes included — works everywhere, flat byte-identical.
render_tuple_choice, shared with theTupleGetItememitter, owns thechoice: a constant index resolves to its element; a runtime index emits
cute::array<decltype(e0), N>{...}[idx],.data()appended for pointers.errors -> 0), numerics unchanged; tests fail before and pass after
(two-level-mesh reduce; one-way plain->shard copy exposing the old
permutation as
out != arange; bf16 round trip through thread fragments);the two newly-emitting fixtures (
staged_dynamic_inner,staged_nested)run bit-exact on card; other stopped fixtures refuse as before — 18 at the
deliberate TMA-atom boundary (feat(codegen): emit SM90 WGMMA atoms in TiledMma with a runtime op #249), 4 in their own words.
Contract
kept_cells == 1refusal is withdrawn (two kept cells are now supported,its static-asserts case dropped); the out-of-scope mesh refusal keeps its
sentence, now phrased per level.
NotImplementedError.Risk
is the plumbing only, as scoped in refactor(parser,ir,analysis,codegen): finish the remaining work of twelve earlier issues #255.
shard local shape) — documented out of scope.
Module(topologies=())auto-launch derives grid/block only from nestedMeshScopes (single-scope two-topology mesh launches 1x1) — separate report.