Skip to content

fix(codegen,runtime): honour every topology a Mesh spans - #67

Open
bigSheep123 wants to merge 1 commit into
tile-ai:mainfrom
bigSheep123:fix/multi-topology-mesh
Open

fix(codegen,runtime): honour every topology a Mesh spans#67
bigSheep123 wants to merge 1 commit into
tile-ai:mainfrom
bigSheep123:fix/multi-topology-mesh

Conversation

@bigSheep123

Copy link
Copy Markdown
Contributor

fix(codegen,runtime): honour every topology a Mesh spans

Fixes #66.

What was wrong

Mesh models an ordered topology sequence — topology is the primary level and
topologies the full tuple (shard §5) — but the
emitter wrote only the primary one into the C++ type, and
shard_axis_projection derived the mesh coordinate from
program_id<primary_scope>().

A mesh over cta × thread therefore emitted a type whose topology claims a
16-element domain against a layout of 4096, and every thread of a CTA resolved
to the same mesh coordinate. Each CTA wrote one thread's worth of elements and
left the rest of its rows untouched — 32640 of 32768 elements wrong on the test
added here, with no diagnostic.

The change

  • Mesh<TTopo, TMeshLayout>Mesh<TTopo, TMeshLayout, TMoreTopos...>.
    The primary topology stays first, so every existing single-topology
    instantiation renders and behaves exactly as before. The new linear_id() is
    the mixed-radix position across the pack, coarsest outermost — which is the
    index the mesh layout is already built against
    (cta_id * thread_size + thread_id).
  • shard_axis_projection takes mesh_t::linear_id(). get_hier_coord
    already did the hierarchical decomposition; it was simply being handed one
    scope's id.
  • The mesh-scope and tensor-view emitters share render_topology and
    extra_topology_args.
    They had duplicated the rendering, and the alias the
    mesh-scope emitter registers is matched against the tensor-view emitter's
    inline type by string equality, so the two must not drift — that coupling
    is why this is one change rather than two.

Test

tests/integration/test_multi_topology_mesh.py — 16 CTAs × 256 threads, 8
contiguous f32 per thread, gmem → rmem → gmem.

It asserts the full output rather than sampling. That is deliberate: the
untouched elements keep whatever the output buffer already held, so a wrong mesh
coordinate still produces a plausible-looking result and a spot check passes.
The test seeds the output with NaN and the input away from zero so an untouched
element cannot coincidentally look right.

Verified to fail on main before the fix (32640 of 32768 elements wrong) and
pass after.

Verification

pytest tests/ -q     644 passed (main @ e55e4b0)  ->  645 passed (+1 new)

Every .pre-commit-config.yaml hook clean: ruff, spec-rules, spec-refs,
spec-entropy, forward-references, comment-hygiene, no-machine-paths,
english-only, and clang-format over all C++ files.

`Mesh` models an ordered topology sequence -- `topology` is the primary level
and `topologies` the full tuple (shard spec section 5) -- but the emitter wrote
only the primary one into the C++ type, and `shard_axis_projection` derived the
mesh coordinate from `program_id<primary_scope>()`.

So a mesh over cta x thread emitted a type whose topology claims a 16-element
domain while its layout has 4096, and every thread of a CTA resolved to the
same mesh coordinate. Each CTA then wrote one thread's worth of elements and
left the rest of its rows untouched. The kernel compiled and ran, so nothing
reported it: on the added test, 32640 of 32768 output elements were wrong.

- `Mesh<TTopo, TMeshLayout>` becomes `Mesh<TTopo, TMeshLayout, TMoreTopos...>`,
  the primary first so existing single-topology instantiations render and
  behave exactly as before, and gains `linear_id()` -- the mixed-radix position
  across the pack, coarsest outermost, which is the index the mesh layout is
  already built against (`cta_id * thread_size + thread_id`).
- `shard_axis_projection` takes `mesh_t::linear_id()`. `get_hier_coord` already
  did the hierarchical decomposition; it was just handed one scope's id.
- The mesh-scope and tensor-view emitters now share `render_topology` and
  `extra_topology_args`. They had duplicated the rendering, and the alias the
  mesh-scope emitter registers is matched against the tensor-view emitter's
  inline type by string equality, so the two must not drift.

The test asserts the full output rather than sampling: a wrong mesh coordinate
still produces a plausible-looking buffer, and only an exact all-elements
comparison distinguishes it.

pytest tests/ -q: 644 passed before, 645 after (+1 new, no regressions).
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.

A Mesh spanning several topologies silently addresses only the primary one

1 participant