From 661f8d988063b26938230c9d0b85056a056d827e Mon Sep 17 00:00:00 2001 From: bigSheep123 <1554869970@qq.com> Date: Fri, 7 Aug 2026 13:10:10 +0800 Subject: [PATCH] fix(codegen,runtime): honour every topology a Mesh spans `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()`. 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` becomes `Mesh`, 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). --- .../runtime/cuda/layout/shard_layout.cuh | 19 +++++- .../runtime/cuda/tensor_view/shard_tensor.cuh | 7 +- .../codegen/cuda/tir/memory/tensor_view.py | 59 +++++++++-------- .../codegen/cuda/tir/stmts/mesh_scope.py | 14 ++-- tests/integration/test_multi_topology_mesh.py | 64 +++++++++++++++++++ 5 files changed, 128 insertions(+), 35 deletions(-) create mode 100644 tests/integration/test_multi_topology_mesh.py diff --git a/include/tilefoundry/runtime/cuda/layout/shard_layout.cuh b/include/tilefoundry/runtime/cuda/layout/shard_layout.cuh index bb8f9596..5928cf6c 100644 --- a/include/tilefoundry/runtime/cuda/layout/shard_layout.cuh +++ b/include/tilefoundry/runtime/cuda/layout/shard_layout.cuh @@ -8,11 +8,26 @@ template struct Topology { static constexpr int size = Size; }; -// Mesh: binds a topology to a MeshLayout. -template struct Mesh { +// Mesh: binds one or more +// topologies to a MeshLayout. The topology product is the mesh domain, which +// the MeshLayout then subdivides into logical axes; ``topology`` stays the +// primary (coarsest) one so single-topology users are unaffected. +template struct Mesh { using topology = TTopo; using layout = TMeshLayout; TMeshLayout layout_value; + + // Linearized position of the current execution instance within the mesh + // domain, coarsest topology outermost. This is the index the MeshLayout is + // built against: a cta x thread mesh addresses as + // ``cta_id * thread_size + thread_id``. With no extra topologies this is + // just ``program_id()``. + CUTE_HOST_DEVICE static size_t linear_id() noexcept { + size_t id = program_id(); + ((id = id * size_t(TMoreTopos::size) + program_id()), + ...); + return id; + } }; // ShardLayout: spec 003 shard layout surface. diff --git a/include/tilefoundry/runtime/cuda/tensor_view/shard_tensor.cuh b/include/tilefoundry/runtime/cuda/tensor_view/shard_tensor.cuh index f8d9fe36..d451deb4 100644 --- a/include/tilefoundry/runtime/cuda/tensor_view/shard_tensor.cuh +++ b/include/tilefoundry/runtime/cuda/tensor_view/shard_tensor.cuh @@ -44,8 +44,6 @@ template struct shard_axis_projection_t { template CUTE_HOST_DEVICE auto shard_axis_projection(ShardTensor const &st) { using mesh_t = typename SL::mesh; - using topo_t = typename mesh_t::topology; - constexpr auto scope = topo_t::scope; using attrs_t = typename SL::attrs; using sl_layout_t = typename SL::layout; using m_layout_t = typename mesh_t::layout; @@ -53,7 +51,10 @@ CUTE_HOST_DEVICE auto shard_axis_projection(ShardTensor const &st) { auto const &sl_layout = st.shard_layout.layout_value; auto const &m_layout = st.shard_layout.mesh_value.layout_value; - auto pid = program_id(); + // Mixed-radix over the whole topology pack: with a mesh spanning cta and + // thread, program_id() alone maps every thread of a CTA onto one + // mesh coordinate. + auto pid = mesh_t::linear_id(); auto crd = m_layout.get_hier_coord(pid); using sl_shape_t = cute::remove_cvref_t; diff --git a/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py b/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py index d8f5ef45..e5328193 100644 --- a/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py +++ b/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py @@ -38,37 +38,48 @@ def _render_layout(shape, strides) -> str: ) +def render_topology(topo) -> str: + """``tilefoundry::Topology``. + + A launch-provided (dynamic) extent has no compile-time size; the mesh value + carries the real extent, so the type parameter is an inert 0 placeholder. + """ + size = topo.size if isinstance(topo.size, int) else 0 + return f"tilefoundry::Topology<{topology_scope_str(topo.name)}, {size}>" + + +def extra_topology_args(mesh) -> str: + """Trailing template args for a mesh spanning several topologies. + + ``Mesh`` keeps the primary topology + first so an ordinary single-topology mesh renders exactly as before, and + the remaining levels follow the layout. Empty when there is only one. + """ + extra = tuple(mesh.topologies or ())[1:] + return "".join(f", {render_topology(t)}" for t in extra) + + def _render_mesh_type(mesh, ctx=None) -> str: """tilefoundry::Mesh<...> — uses scope alias if registered in ctx.""" + ml = mesh.layout + shape_args = ", ".join(f"cute::Int<{s}>" for s in ml.shape) + stride_args = ", ".join(f"cute::Int<{s}>" for s in ml.strides) + inline = ( + f"tilefoundry::Mesh<" + f"{render_topology(mesh.topology)}, " + f"cute::Layout, cute::Stride<{stride_args}>>" + f"{extra_topology_args(mesh)}>" + ) if ctx and hasattr(ctx, '_mesh_aliases'): # try exact id match entry = ctx._mesh_aliases.get(id(mesh)) if entry: return entry[0] # alias name # try structural fallback: compare inline type string - topo = mesh.topology - scope = topology_scope_str(topo.name) - ml = mesh.layout - shape_args = ", ".join(f"cute::Int<{s}>" for s in ml.shape) - stride_args = ", ".join(f"cute::Int<{s}>" for s in ml.strides) - inline = ( - f"tilefoundry::Mesh<" - f"tilefoundry::Topology<{scope}, {topo.size}>, " - f"cute::Layout, cute::Stride<{stride_args}>>>" - ) for alias_name, type_str in ctx._mesh_aliases.values(): if type_str == inline: return alias_name - topo = mesh.topology - scope = topology_scope_str(topo.name) - ml = mesh.layout - shape_args = ", ".join(f"cute::Int<{s}>" for s in ml.shape) - stride_args = ", ".join(f"cute::Int<{s}>" for s in ml.strides) - return ( - f"tilefoundry::Mesh<" - f"tilefoundry::Topology<{scope}, {topo.size}>, " - f"cute::Layout, cute::Stride<{stride_args}>>>" - ) + return inline def _render_attr(a) -> str: @@ -181,18 +192,14 @@ def _mesh_dim(d): ml_shape = ", ".join(_mesh_dim(d) for d in ml.shape) ml_stride = ", ".join(_static_dim(s, "mesh layout stride") for s in ml.strides) - scope = topology_scope_str(topo.name) - # A launch-provided CTA extent has no compile-time size; the value carries - # the real extent, so the Topology type parameter is an inert placeholder. - topo_size = topo.size if isinstance(topo.size, int) else 0 attrs = ", ".join(_render_attr(a) for a in sl.attrs) preamble = [ f"auto {sl_var} = cute::make_layout(" f"cute::make_shape({sl_shape}), cute::make_stride({sl_stride}));", f"auto {ml_var} = cute::make_layout(" f"cute::make_shape({ml_shape}), cute::make_stride({ml_stride}));", - f"tilefoundry::Mesh, " - f"decltype({ml_var})> {mesh_var}{{{ml_var}}};", + f"tilefoundry::Mesh<{render_topology(topo)}, " + f"decltype({ml_var}){extra_topology_args(sl.mesh)}> {mesh_var}{{{ml_var}}};", ] value_expr = ( f"tilefoundry::ShardLayout, " diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py b/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py index a5ac480b..048e3236 100644 --- a/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py +++ b/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py @@ -6,7 +6,10 @@ from tilefoundry.codegen.cuda.context import ( CodegenContext, register_codegen_cuda, - topology_scope_str, +) +from tilefoundry.codegen.cuda.tir.memory.tensor_view import ( + extra_topology_args, + render_topology, ) from tilefoundry.ir.tir.stmts import MeshScope from tilefoundry.target import validate_cuda_topology_levels @@ -23,13 +26,16 @@ def _validate_topology(mesh) -> None: def _mesh_type(mesh) -> str: - topo = mesh.topology + # Shares the topology renderers with the tensor_view emitter: the alias + # registered here is matched against that emitter's inline type by string + # equality, so the two must not drift. shape_types = ", ".join(f"cute::Int<{s}>" for s in mesh.layout.shape) stride_types = ", ".join(f"cute::Int<{s}>" for s in mesh.layout.strides) return ( f"tilefoundry::Mesh<" - f"tilefoundry::Topology<{topology_scope_str(topo.name)}, {topo.size}>, " - f"cute::Layout, cute::Stride<{stride_types}>>>" + f"{render_topology(mesh.topology)}, " + f"cute::Layout, cute::Stride<{stride_types}>>" + f"{extra_topology_args(mesh)}>" ) diff --git a/tests/integration/test_multi_topology_mesh.py b/tests/integration/test_multi_topology_mesh.py new file mode 100644 index 00000000..615f3cc0 --- /dev/null +++ b/tests/integration/test_multi_topology_mesh.py @@ -0,0 +1,64 @@ +"""A Mesh spanning several topologies addresses the whole topology product. + +``Mesh`` models an ordered topology sequence ([shard §5](docs/spec/shard.md#5-mesh)): +``topology`` is the primary level and ``topologies`` the full tuple. A mesh over +``cta`` x ``thread`` therefore has a domain of ``n_cta * n_thread``, and its +layout subdivides that product. + +The emitted C++ used to carry only the primary topology, so the mesh coordinate +came from ``program_id()`` alone and every thread of a CTA resolved to the +same coordinate: each CTA wrote one thread's worth of elements and left the rest +untouched. The kernel still compiled and ran, so only a full-output assertion +catches it — hence the exact all-elements comparison below rather than a spot +check. +""" + +from __future__ import annotations + +import torch + +import tilefoundry +from tilefoundry import func +from tilefoundry.dsl import Tensor, tf +from tilefoundry.dsl.storage import gmem, rmem +from tilefoundry.ir.types.shard import Layout, Mesh, Topology + +_CTAS = 16 +_THREADS = 256 +_PER_THREAD = 8 +# One row per CTA; every thread of the CTA owns a contiguous run of the row. +_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"]: + # One mesh over both levels: the layout's leading axis spans the ctas and + # the trailing one the threads, so the mesh domain is 16 * 256 = 4096. + 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) + + +def test_a_cta_by_thread_mesh_writes_every_element() -> None: + rm = tilefoundry.compile(square_grid_of_threads, target="cuda") + torch.manual_seed(0) + # Non-uniform data so a coordinate collision cannot pass by coincidence, + # and no zeros so an untouched output element cannot look correct. + x = torch.randn(_CTAS, _COLS, dtype=torch.float32, device="cuda") + 2.0 + out = torch.full_like(x, float("nan")) + rm(x, out) + torch.cuda.synchronize() + + expected = x * x + assert torch.allclose(out, expected, rtol=0, atol=0), ( + f"{int((out != expected).sum())} of {out.numel()} elements wrong " + f"(every thread of a CTA resolving to one mesh coordinate leaves " + f"most of each row untouched)" + )