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

include/tilefoundry/runtime/cuda/layout/shard_layout.cuh

@@ -8,11 +8,26 @@ template <TopologyScope Scope, int Size> struct Topology {
     static constexpr int size = Size;
 };
 
-// Mesh<topology, cute_layout>: binds a topology to a MeshLayout.
-template <class TTopo, class TMeshLayout> struct Mesh {
+// Mesh<topology, cute_layout, more_topologies...>: 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 <class TTopo, class TMeshLayout, class... TMoreTopos> 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<TTopo::scope>()``.
+    CUTE_HOST_DEVICE static size_t linear_id() noexcept {
+        size_t id = program_id<TTopo::scope>();
+        ((id = id * size_t(TMoreTopos::size) + program_id<TMoreTopos::scope>()),
+         ...);
+        return id;
+    }
 };
 
 // ShardLayout<layout, attrs_tuple, mesh>: spec 003 shard layout surface.

include/tilefoundry/runtime/cuda/tensor_view/shard_tensor.cuh

@@ -44,16 +44,17 @@ template <int TRank, int MRank> struct shard_axis_projection_t {
 template <class T, class GL, class SL>
 CUTE_HOST_DEVICE auto shard_axis_projection(ShardTensor<T, GL, SL> 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;
 
     auto const &sl_layout = st.shard_layout.layout_value;
     auto const &m_layout = st.shard_layout.mesh_value.layout_value;
 
-    auto pid = program_id<scope>();
+    // Mixed-radix over the whole topology pack: with a mesh spanning cta and
+    // thread, program_id<cta>() 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<decltype(cute::shape(sl_layout_t{}))>;

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<scope, size>``.
+
+    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<TTopo, TMeshLayout, TMoreTopos...>`` 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::Shape<{shape_args}>, 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::Shape<{shape_args}>, 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::Shape<{shape_args}>, 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<tilefoundry::Topology<{scope}, {topo_size}>, "
-        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<decltype({sl_var}), cute::tuple<{attrs}>, "

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::Shape<{shape_types}>, cute::Stride<{stride_types}>>>"
+        f"{render_topology(mesh.topology)}, "
+        f"cute::Layout<cute::Shape<{shape_types}>, cute::Stride<{stride_types}>>"
+        f"{extra_topology_args(mesh)}>"
     )
 
 

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<cta>()`` 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)"
+    )