fix(codegen,runtime): honour every topology a Mesh spans - #67
Open
bigSheep123 wants to merge 1 commit into
Open
Conversation
`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).
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.
fix(codegen,runtime): honour every topology a Mesh spans
Fixes #66.
What was wrong
Meshmodels an ordered topology sequence —topologyis the primary level andtopologiesthe full tuple (shard §5) — but theemitter wrote only the primary one into the C++ type, and
shard_axis_projectionderived the mesh coordinate fromprogram_id<primary_scope>().A mesh over
cta×threadtherefore emitted a type whose topology claims a16-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()isthe 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_projectiontakesmesh_t::linear_id().get_hier_coordalready did the hierarchical decomposition; it was simply being handed one
scope's id.
render_topologyandextra_topology_args. They had duplicated the rendering, and the alias themesh-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, 8contiguous 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
mainbefore the fix (32640 of 32768 elements wrong) andpass after.
Verification
Every
.pre-commit-config.yamlhook clean: ruff, spec-rules, spec-refs,spec-entropy, forward-references, comment-hygiene, no-machine-paths,
english-only, and clang-format over all C++ files.