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)" + )