A Mesh spanning several topologies silently addresses only the primary one
Affects: main @ e55e4b0
Summary
Mesh models an ordered topology sequence — topology is the primary level and
topologies the full tuple (shard §5). A mesh over
cta × thread therefore has a domain of n_cta * n_thread, which its layout
subdivides.
The CUDA emitter writes only the primary topology into the tilefoundry::Mesh<>
type, and shard_axis_projection derives the mesh coordinate from
program_id<primary_scope>(). Every thread of a CTA therefore resolves to the
same mesh coordinate: each CTA writes one thread's worth of elements and
leaves the rest untouched.
The kernel compiles and runs, and the output buffer looks plausible, so nothing
reports the problem.
Reproduction
16 CTAs × 256 threads, 8 contiguous f32 per thread, gmem → rmem → gmem:
_CTAS, _THREADS, _PER_THREAD = 16, 256, 8
_COLS = _THREADS * _PER_THREAD
@func(topologies=(Topology("cta", _CTAS), Topology("thread", _THREADS)))
def square_grid_of_threads(
a: Tensor[(_CTAS, _COLS), "f32"],
) -> Tensor[(_CTAS, _COLS), "f32"]:
with Mesh(
topology=[Topology("cta", _CTAS), Topology("thread", _THREADS)],
layout=Layout(shape=(_CTAS, _THREADS), strides=(_THREADS, 1)),
names=("c", "t"),
) as m:
reg = tf.reshard(a, layout=(_CTAS @ m.c, _COLS @ m.t), storage=rmem)
out = tf.mul(reg, reg)
return tf.reshard(out, layout=(_CTAS @ m.c, _COLS @ m.t), storage=gmem)
Actual
AssertionError: 32640 of 32768 elements wrong
128 elements are correct — exactly 16 CTAs × 8 per-thread elements, i.e. one
thread's worth per CTA.
The emitted type shows why: a topology claiming a 16-element domain against a
layout of 4096.
tilefoundry::Mesh<tilefoundry::Topology<TopologyScope::cta, 16>,
cute::Layout<Shape<16,256>, Stride<256,1>>>
Expected
All 32768 elements written.
Root cause
include/tilefoundry/runtime/cuda/layout/shard_layout.cuh —
Mesh<TTopo, TMeshLayout> holds exactly one topology.
include/tilefoundry/runtime/cuda/tensor_view/shard_tensor.cuh —
shard_axis_projection uses program_id<topo_t::scope>().
src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py and
src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py — both render only
mesh.topology, ignoring mesh.topologies.
The Python IR is already correct; only the emission and the C++ runtime are
single-topology.
Notes
A spot check cannot catch this — the untouched elements are whatever the output
buffer already held, so a wrong mesh coordinate still yields a plausible
buffer. Only a full-output comparison distinguishes it.
PR follows.
A Mesh spanning several topologies silently addresses only the primary one
Affects:
main@e55e4b0Summary
Meshmodels an ordered topology sequence —topologyis the primary level andtopologiesthe full tuple (shard §5). A mesh overcta×threadtherefore has a domain ofn_cta * n_thread, which its layoutsubdivides.
The CUDA emitter writes only the primary topology into the
tilefoundry::Mesh<>type, and
shard_axis_projectionderives the mesh coordinate fromprogram_id<primary_scope>(). Every thread of a CTA therefore resolves to thesame mesh coordinate: each CTA writes one thread's worth of elements and
leaves the rest untouched.
The kernel compiles and runs, and the output buffer looks plausible, so nothing
reports the problem.
Reproduction
16 CTAs × 256 threads, 8 contiguous f32 per thread, gmem → rmem → gmem:
Actual
128 elements are correct — exactly
16 CTAs × 8 per-thread elements, i.e. onethread's worth per CTA.
The emitted type shows why: a topology claiming a 16-element domain against a
layout of 4096.
Expected
All 32768 elements written.
Root cause
include/tilefoundry/runtime/cuda/layout/shard_layout.cuh—Mesh<TTopo, TMeshLayout>holds exactly one topology.include/tilefoundry/runtime/cuda/tensor_view/shard_tensor.cuh—shard_axis_projectionusesprogram_id<topo_t::scope>().src/tilefoundry/codegen/cuda/tir/memory/tensor_view.pyandsrc/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py— both render onlymesh.topology, ignoringmesh.topologies.The Python IR is already correct; only the emission and the C++ runtime are
single-topology.
Notes
A spot check cannot catch this — the untouched elements are whatever the output
buffer already held, so a wrong mesh coordinate still yields a plausible
buffer. Only a full-output comparison distinguishes it.
PR follows.