refactor(shard): state one arrangement per level a mesh names - #185
Merged
zhen8838 merged 11 commits intoSep 25, 2026
Merged
Conversation
Compare and compose mesh scopes on the positions each level states rather than on the axes standing where it does, so a scope that is part of a level -- one warp of a CTA's threads, two warpgroups of three -- is read as the positions it is. Count a loop off its own bounds instead of by walking the coordinates its bounds name, permute an unstated layout as the C order it stands for, and narrow a window of a grouped arrangement through its modes.
A loop body admits `with Mesh(...)`, carries what one binds -- including a name the region reads on its way to binding it, which a body that repeats carries -- and yields the region a rebound name reaches. A TIR loop bound reads back as the dimension arithmetic it was printed from, its coordinates are held to the scopes around the loop, and the printer writes a window as the subscript it was authored as and a grouped arrangement in full.
The four programs an author states before choosing an instruction, each read the way anything downstream reads one: parsed, checked, and measured.
Drop the restated device-domain sentence and the topology-count aside from the mesh section, the second statement of the verbose-form fallback, and the rationale that repeats a rule stated two bullets up.
An unstated layout is the C order it stands for, so the case that asserted an unstated result now asserts the same permutation the stated one gives.
Three small programs, each for one thing the stage changed: a loop whose region names a suffix of the levels in force and whose count is read at its widest, a TIR loop bound that reads a unit coordinate, and a window of an arrangement whose modes are grouped by the axis they belong to.
zhen8838
commented
Sep 24, 2026
A mesh stated one arrangement over every level's axes and left eight helpers to work out which axes were whose and what each level's own numbering was. It now states them: the levels run parallel to the topologies, each in its own numbering, and the cut is made once, where a mesh is built. Composing is CuTe's append and replace on those statements, under one merge_mesh entry. Gone with the guessing: topology_axes, grouped_layout, positions_at, topology_projection, level_positions and two private readings of the same question. unflatten moves to int_tuple, beside flatten.
Making the steps of a compact arrangement, and reading an index back into the coordinate that reaches it, are not operations on layouts. They move out of the algebra into stride.py, as CuTe keeps stride.hpp apart from layout.hpp.
A folder naming them a sublayer said nothing a reader could use: Topology, Mesh, Layout and ShardLayout are types, and they sit in the type root beside the others now, filed one module per thing the algebra is written against as CuTe files them. make_mesh moves in beside Mesh.
They were filed as helpers, which named nothing: reading a dim's static value, its upper bound, or a shape's total are operations on dims, and they sit in dim.py with the rest of them. Two cycle guards go with the move, since dim.py already holds the Constant they were dodging.
zhen8838
commented
Sep 24, 2026
A mesh's layout is now a nested Layout for every mesh, one top-level mode per level it names, a mesh naming one level included. That is what a layout already is -- a tuple -- so the CuTe operations read it: get(layout, i) is level i, flatten(layout) is every axis end to end, and repeat_like against its shape states per axis whatever was stated per level. level_axes, which worked out which flat positions each level covered, is gone; nothing needs the positions. merge_mesh and the scope-match readings move out of ir/types to ir/mesh_scope.py: which scope a statement stands inside is neither a type nor a visitor. local.py dissolves -- local_layout and local_layout_and_offset belong to ShardLayout, local_type_of to utils -- and the type root exports types and their constructors only. The seven shape readings leave dim.py for utils.py, where the helpers they came from were headed. compact_major recurses over nested shapes as CuTe's does, so a nested mesh shape can be authored directly; idx2crd and crd2idx recurse with it.
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.
Why
GEMMs: mesh scope comparison, scope composition, loop trip counts, the transpose layout
rule, windows of grouped arrangements, the HIR loop grammar, TIR loop-bound reading and
verification, and three printer forms. Each is an upstream gap, not a downstream policy.
Meshstated one arrangement over everylevel's axes and left eight helpers to work out which axes were whose.
What
Meshstates one arrangement per level it names, each in that level's own numbering.The cut is made once, where a mesh is built, from what the author wrote. Composing is
CuTe's
appendandreplaceon those statements, under onemerge_meshentry.Mesh.levels/positions/written/offset/slicedread it.Gone:
topology_axes,grouped_layout,positions_at,topology_projection,level_positions,composedand two private readings of the same question.unflattenbesideflatteninint_tuple;prefix_product,c_order_strides,try_c_order_stridesandidx2crdinto a newstride.py; and their/types/shard/folder dissolved into the type root, since these are types.IterationScope.trips()counts a loop offstop - startinstead of enumerating thecoordinates its bounds name, so a 64x68 tile grid is no longer refused by the 4096-point
limit;
trips_precisionsays whether that count is exact or widened.Transposeof a layout-less input states the permuted C order;Slicenarrows an axiswritten as a group of modes, which previously raised
TypeError.with Mesh(...)and carries what one binds; a TIR loop bound readsback as the dimension arithmetic it was printed from, its coordinates are verified against
the scopes around the loop, and the printer writes a window as the subscript it was
authored as and a grouped arrangement in full.
Contract
docs/spec/shard.md:Mesh.layoutis one arrangement per level;merge_meshreplacescomposed, and a mesh naming a suffix of the levels in force refines them where it wasrejected. No stride or offset is rescaled on composition.
docs/spec/analysis.md:trips()is read off the loop's bounds;trips_precisionis new.docs/spec/hir.md:Transposeoflayout=Nonestates aLayout;Slicenarrows an axiswritten as a group of modes.
docs/spec/parser.md,tir.md,inspection.md: loop bodies holdwith Mesh(...),Forbound coordinates are verified, grouped layouts print in full.
docs/spec/code-organization.md: the shard types are filed in the type root.Risk
cut out of it at construction. The five hand-written multi-level meshes in the tree are
untouched.
46 programs, IDENTICAL through every step of the refactor.