Skip to content

refactor(shard): state one arrangement per level a mesh names - #185

Merged
zhen8838 merged 11 commits into
tile-ai:mainfrom
zhen8838:refactor/atomsched-stage1-upstream-gaps
Sep 25, 2026
Merged

zhen8838 merged 11 commits into
tile-ai:mainfrom
zhen8838:refactor/atomsched-stage1-upstream-gaps

Conversation

@zhen8838

@zhen8838 zhen8838 commented Sep 24, 2026 •

Copy link
Copy Markdown
Collaborator

Why

  • A downstream scheduler replaces nine upstream behaviours at import time to schedule real
    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.
  • Review then found the root of the first two: a Mesh stated one arrangement over every
    level's axes and left eight helpers to work out which axes were whose.

What

  • Mesh states 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 append and replace on those statements, under one merge_mesh entry.
    Mesh.levels / positions / written / offset / sliced read it.
    Gone: topology_axes, grouped_layout, positions_at, topology_projection,
    level_positions, composed and two private readings of the same question.
  • Filed as CuTe files it: unflatten beside flatten in int_tuple; prefix_product,
    c_order_strides, try_c_order_strides and idx2crd into a new stride.py; and the
    ir/types/shard/ folder dissolved into the type root, since these are types.
  • IterationScope.trips() counts a loop off stop - start instead of enumerating the
    coordinates its bounds name, so a 64x68 tile grid is no longer refused by the 4096-point
    limit; trips_precision says whether that count is exact or widened.
  • Transpose of a layout-less input states the permuted C order; Slice narrows an axis
    written as a group of modes, which previously raised TypeError.
  • A HIR loop body holds with Mesh(...) and carries what one binds; a TIR loop bound reads
    back 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.layout is one arrangement per level; merge_mesh replaces
    composed, and a mesh naming a suffix of the levels in force refines them where it was
    rejected. No stride or offset is rescaled on composition.
  • docs/spec/analysis.md: trips() is read off the loop's bounds; trips_precision is new.
  • docs/spec/hir.md: Transpose of layout=None states a Layout; Slice narrows an axis
    written as a group of modes.
  • docs/spec/parser.md, tir.md, inspection.md: loop bodies hold with Mesh(...), For
    bound coordinates are verified, grouped layouts print in full.
  • docs/spec/code-organization.md: the shard types are filed in the type root.

Risk

  • The authored surface does not change: a mesh is written the way it was, and the levels are
    cut out of it at construction. The five hand-written multi-level meshes in the tree are
    untouched.
  • Witnessed by the existing suite (1060 tests) and by a downstream byte-for-byte oracle over
    46 programs, IDENTICAL through every step of the refactor.

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.
Comment thread docs/spec/parser.md Outdated
Comment thread src/tilefoundry/analysis/footprint.py Outdated
Comment thread src/tilefoundry/ir/types/shard/__init__.py Outdated
Comment thread src/tilefoundry/ir/types/shard/mesh.py Outdated
Comment thread src/tilefoundry/parser/ast_pattern.py Outdated
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.
@zhen8838 zhen8838 changed the title feat(ir): support sub-level mesh scopes and grouped arrangements refactor(shard): state one arrangement per level a mesh names Sep 24, 2026
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.
Comment thread docs/spec/shard.md Outdated
Comment thread src/tilefoundry/analysis/compute_cost.py Outdated
Comment thread src/tilefoundry/codegen/cuda/tir/memory/copy.py Outdated
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.
@zhen8838
zhen8838 merged commit 5f417bc into tile-ai:main Sep 25, 2026
1 check passed
@zhen8838
zhen8838 deleted the refactor/atomsched-stage1-upstream-gaps branch September 25, 2026 11:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant