Skip to content

A Mesh spanning several topologies silently addresses only the primary one #66

Description

@bigSheep123

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions