From f30c23ccd6bb7c58bb66cdf136d66a449ae12208 Mon Sep 17 00:00:00 2001 From: sablin39 <1020030829@qq.com> Date: Wed, 21 Oct 2026 17:49:29 +0800 Subject: [PATCH 1/3] =?UTF-8?q?fix(runtime):=20count=20reduce=20warps=20pe?= =?UTF-8?q?r=20mesh=20level=20for=20composed=20cta=C3=97thread=20meshes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ops::reduce derived its plan from a single per-mesh is_thread boolean, and codegen flattened a composed cta×thread mesh — what every multi-CTA finalize produces — into one cta-scoped Mesh. Every thread axis was therefore classified cta-level: lanes_of=1, warps_of=extent per axis, so warps_per_group became the full mesh instance count (256 instead of 8 in the RMSNorm tutorial). The cross-warp tier then indexed the shared workspace at (group_start + w) * 32 + lane, far past the allocation: OOB smem reads under compute-sanitizer. The verifier multiplied the same flattened extents and demanded 4096 workspace slots where 8 suffice. Classify each mesh axis by the level it belongs to: - render_shard_layout_value emits one shape/stride nest per topology level and names every scope in the Mesh type, so the runtime can see the levels; single-level emission is byte-identical, and a multi-level mesh that is also sliced is refused. - reduce_dispatch spends the lane budget on the thread-level axes only, fastest axis first; cta-level axes keep lanes_of=1, warps_of=extent. Mesh gains one axis_scope() member answering which scope flat axis I runs at — a single-level mesh answers its own scope for every axis, a multi-level mesh reads the nest holding I off its layout — so the levels a mesh names are read off the mesh itself, here and for any op that needs them next. The out-of-scope refusal keeps its sentence ("a reduce mesh's scope must be cta or thread"), now phrased per level. - IntraCta iterates the kept cells (fold, warp butterfly, workspace combine per cell) instead of asserting kept_cells == 1, and cta_combine_via_workspace ends with a barrier so the one-slot-per-warp workspace is safe to reuse across cells. The refusal goes away with the assert: a source keeping two cells is a supported reduce now, so the static-asserts suite drops the case that pinned the old sentence. - The verifier sizes the workspace from the thread level alone: each CTA combines its own partials in its own shared memory. Single-level thread and pure-cta plans are unchanged (compile-time probe over the old and new paths). The new regression test reduces over a cta×thread two-level mesh: it fails before this change (the verifier demands 16 slots for a 4-warp CTA, and with verification bypassed the dispatch goes cross-warp out of bounds) and passes after, with compute-sanitizer reporting 0 errors. --- docs/spec/runtime.md | 6 ++ .../tilefoundry/runtime/cuda/layout/mesh.cuh | 26 ++++++ .../cuda/ops/reduce/reduce_common_impl.h | 61 +++++++++----- .../cuda/ops/reduce/reduce_intra_cta_impl.h | 39 +++++---- .../codegen/cuda/tir/memory/tensor_view.py | 47 ++++++++--- src/tilefoundry/ir/tir/reduce.py | 18 ++++- tests/ops/tir/test_reduce.py | 81 +++++++++++++++++++ tests/runtime/cuda/static_asserts.cu | 22 ----- tests/runtime/test_cuda_static_asserts.py | 1 - 9 files changed, 224 insertions(+), 77 deletions(-) diff --git a/docs/spec/runtime.md b/docs/spec/runtime.md index 6790e129a..b50b5bd99 100644 --- a/docs/spec/runtime.md +++ b/docs/spec/runtime.md @@ -968,6 +968,12 @@ template struct Mesh { static constexpr int level_count; static constexpr TopologyScope scope; TLayout layout; + /// The scope flat axis ``I`` runs at. A mesh naming one level + /// answers its own scope for every axis. A mesh naming several + /// groups its layout one nest per level, and the nest holding + /// flat axis ``I`` states its level. + template + CUTE_HOST_DEVICE static constexpr TopologyScope axis_scope(); }; template diff --git a/include/tilefoundry/runtime/cuda/layout/mesh.cuh b/include/tilefoundry/runtime/cuda/layout/mesh.cuh index 8bc4dcd98..517eb5592 100644 --- a/include/tilefoundry/runtime/cuda/layout/mesh.cuh +++ b/include/tilefoundry/runtime/cuda/layout/mesh.cuh @@ -24,6 +24,32 @@ template struct Mesh { static constexpr int level_count = sizeof...(Topos); static constexpr TopologyScope scope = detail::first_scope::value; TLayout layout; + + /// The scope flat axis ``I`` runs at. A mesh naming one level + /// answers its own scope for every axis. A mesh naming several + /// groups its layout one nest per level, and the nest holding + /// flat axis ``I`` states its level. + template + CUTE_HOST_DEVICE static constexpr TopologyScope axis_scope() { + if constexpr (sizeof...(Topos) == 1) { + return scope; + } else { + constexpr TopologyScope scopes[] = {Topos...}; + TopologyScope found = scope; + int first = 0; + [&](std::index_sequence) { + (([&] { + constexpr int width = + int(cute::rank(cute::get(cute::shape(TLayout{})))); + if (I >= first && I < first + width) + found = scopes[Ls]; + first += width; + }()), + ...); + }(std::make_index_sequence{}); + return found; + } + } }; namespace detail { diff --git a/include/tilefoundry/runtime/cuda/ops/reduce/reduce_common_impl.h b/include/tilefoundry/runtime/cuda/ops/reduce/reduce_common_impl.h index 277b46eb1..deb3f751c 100644 --- a/include/tilefoundry/runtime/cuda/ops/reduce/reduce_common_impl.h +++ b/include/tilefoundry/runtime/cuda/ops/reduce/reduce_common_impl.h @@ -156,7 +156,9 @@ CUTE_HOST_DEVICE constexpr int reduced_span() { /// Combine one partial per warp, with the reduction's own combine. /// -/// Combine shared-workspace partials by warp group. +/// Combine shared-workspace partials by warp group. The trailing barrier +/// leaves every thread past its reads, so the workspace is safe to reuse for +/// the next group's partials. template __device__ float cta_combine_via_workspace(float warp_partial, WorkspaceT &workspace, @@ -174,6 +176,7 @@ __device__ float cta_combine_via_workspace(float warp_partial, for (int w = 0; w < warps_per_group; ++w) { acc = combine(acc, static_cast(workspace(group_start + w))); } + __syncthreads(); return acc; } @@ -205,44 +208,52 @@ struct reduce_dispatch_info { }; using plan_t = reduce_dispatch_info; -/// Derive, from the (src, dst) operand ShardLayouts, the active reduction level -/// and its ``warps_per_group``. Pure compile-time so the caller can select the -/// tier with ``if constexpr`` — otherwise the untaken tier still instantiates -/// and, e.g., ``CrossWarp`` would trip its supported-op guard. -/// Requires a static mesh layout (the reduce mesh is a thread-scoped static -/// mesh); a reduced axis on a non-thread mesh scope yields the cross-warp tier. /// Lanes of one warp this thread mesh hands out, and zero when its ids state /// no ``(lanes, warps)`` at all. template CUTE_HOST_DEVICE constexpr int lanes_of_mesh() { - if constexpr (TMesh::scope == TopologyScope::thread && - tilefoundry::is_warped(TMesh{})) + using mesh_t = cute::remove_cvref_t; + if constexpr (mesh_t::scope == TopologyScope::thread && + tilefoundry::is_warped(mesh_t{})) return kWarpSize; else return 0; } +/// Derive, from the (src, dst) operand ShardLayouts, the active reduction level +/// and its ``warps_per_group``. Pure compile-time so the caller can select the +/// tier with ``if constexpr`` — otherwise the untaken tier still instantiates +/// and, e.g., ``CrossWarp`` would trip its supported-op guard. +/// Requires a static mesh layout; each mesh axis is classified by the level it +/// belongs to, so a composed cta×thread mesh spends the lane budget only on +/// its thread-level axes. template CUTE_HOST_DEVICE constexpr reduce_dispatch_info reduce_dispatch() { using src_attrs = typename SrcSL::attrs; using dst_attrs = typename DstSL::attrs; - using mesh_t = typename SrcSL::mesh; + using mesh_t = cute::remove_cvref_t; static_assert(mesh_t::scope == TopologyScope::thread || mesh_t::scope == TopologyScope::cta, "ops::reduce: mesh must name cta or thread"); - constexpr bool is_thread = mesh_t::scope == TopologyScope::thread; using m_layout_t = typename mesh_t::layout_type; constexpr int m_rank = cute::tuple_size::value; + /// Levels order coarsest first, so a thread level is the last nest when + /// the mesh names one, and the last flat axis's scope answers for it. + constexpr bool has_thread = + mesh_t::template axis_scope() == TopologyScope::thread; static_assert(detail::shard_attrs_match_mesh() && detail::shard_attrs_match_mesh(), "ops::reduce: both operands need one attr per mesh axis"); static_assert(std::is_same_v, "ops::reduce: the two operands must name one mesh"); - static_assert(is_thread || mesh_t::scope == TopologyScope::cta, + static_assert(has_thread || mesh_t::scope == TopologyScope::cta, "ops::reduce: a reduce mesh's scope must be cta or thread"); + /// The attrs answer for the mesh's axes end to end, whatever grouping the + /// levels state, so extents and scopes read off the flattened shape. int m_ext[m_rank] = {}; bool reduced[m_rank] = {}; - auto const m_shape = cute::shape(m_layout_t{}); + bool thread_axis[m_rank] = {}; + auto const m_shape = cute::flatten(cute::shape(m_layout_t{})); [&](std::index_sequence) { ((m_ext[Is] = int(cute::get(m_shape))), ...); ((reduced[Is] = @@ -252,32 +263,44 @@ CUTE_HOST_DEVICE constexpr reduce_dispatch_info reduce_dispatch() { cute::remove_cvref_t(dst_attrs{}))>, shard::B>), ...); + ((thread_axis[Is] = + mesh_t::template axis_scope() == TopologyScope::thread), + ...); }(std::make_index_sequence{}); - /// How each mesh axis divides into lanes and warps. + /// How each mesh axis divides into lanes and warps. A thread-level axis + /// spends the lane budget, fastest axis first; a cta-level axis reduces + /// across CTAs rather than warp lanes. The thread level's own mesh, which + /// ``get`` picks out of a mesh naming several levels, states the lanes. int lanes_of[m_rank] = {}; int warps_of[m_rank] = {}; + constexpr int kLanes = [] { + if constexpr (has_thread) + return lanes_of_mesh(mesh_t{}))>>(); + else + return 0; + }(); int stride = 1; for (int i = m_rank - 1; i >= 0; --i) { - if (is_thread) { - const int room = lanes_of_mesh() / stride; + if (thread_axis[i]) { + const int room = kLanes / stride; const int lanes = room <= 1 ? 1 : (m_ext[i] < room ? m_ext[i] : room); lanes_of[i] = lanes; warps_of[i] = m_ext[i] / lanes; + stride *= m_ext[i]; } else { - /// CTA meshes reduce across CTAs rather than warp lanes. lanes_of[i] = 1; warps_of[i] = m_ext[i]; } - stride *= m_ext[i]; } bool lane_reduced = false; bool mesh_reduced = false; /// A mesh whose ids state no (lanes, warps) divides into neither, so it /// fails the same check a ragged axis does and says the same thing. - bool warp_aligned = !is_thread || lanes_of_mesh() > 0; + bool warp_aligned = !has_thread || kLanes > 0; int warps_per_group = 1; int lanes_reduced = 1; for (int i = 0; i < m_rank; ++i) { diff --git a/include/tilefoundry/runtime/cuda/ops/reduce/reduce_intra_cta_impl.h b/include/tilefoundry/runtime/cuda/ops/reduce/reduce_intra_cta_impl.h index d382b3402..ea75a1a21 100644 --- a/include/tilefoundry/runtime/cuda/ops/reduce/reduce_intra_cta_impl.h +++ b/include/tilefoundry/runtime/cuda/ops/reduce/reduce_intra_cta_impl.h @@ -8,13 +8,11 @@ namespace reduce_impl { /// ── Reduce tier-2: cross-warp within a CTA (smem workspace) ─────── -/// This tier's contract has a single group output per thread (the reduced -/// ``n_cells == 1``): each thread folds its *entire* per-thread tensor via -/// ``local_fold`` (its ``step == cute::size(s)`` whole-tensor path — correct -/// however many cute axes the per-thread layout carries, see -/// reduce_common_impl.h), then a warp butterfly + ``cta_combine_via_workspace`` -/// combine across warps; ``reduce_traits`` finalises. MEAN divides by -/// ``step × 32 × warps_per_group``. +/// This tier's contract has one group output per kept cell: each thread folds +/// its per-thread tensor cell by cell via ``local_fold``, then a warp +/// butterfly + ``cta_combine_via_workspace`` combine across warps, reusing the +/// one-slot-per-warp workspace from cell to cell; ``reduce_traits`` +/// finalises. MEAN divides by ``step × lanes_reduced × warps_per_group``. template struct IntraCta { template __device__ void operator()(SrcT const &src, DstT &dst, @@ -30,21 +28,20 @@ template struct IntraCta { auto &&ws = tilefoundry::local_tensor(workspace); using value_type = cute::remove_cvref_t; - static_assert(kept_cells() == 1, - "ops::reduce (intra-CTA tier): kept axes must form " - "one cell"); + constexpr int kCells = kept_cells(); constexpr int kSpan = reduced_span(); - const float local = local_fold(s, 0); - const float warp_partial = - tilefoundry::warp_reduce::combine_op, - plan.lanes_reduced>(local); - const float cta_partial = - cta_combine_via_workspace(warp_partial, ws, warps_per_group); - const float total_n = - float(kSpan) * float(plan.lanes_reduced) * float(warps_per_group); - const float result = reduce_traits::finalize(cta_partial, total_n); - for (int i = 0; i < static_cast(cute::size(d)); ++i) { - d(i) = static_cast(result); + CUTE_UNROLL + for (int j = 0; j < kCells; ++j) { + const float local = local_fold(s, j); + const float warp_partial = + tilefoundry::warp_reduce::combine_op, + plan.lanes_reduced>(local); + const float cta_partial = cta_combine_via_workspace( + warp_partial, ws, warps_per_group); + const float total_n = float(kSpan) * float(plan.lanes_reduced) * + float(warps_per_group); + d(j) = static_cast( + reduce_traits::finalize(cta_partial, total_n)); } } }; diff --git a/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py b/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py index 655f436b6..dfe759678 100644 --- a/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py +++ b/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py @@ -29,7 +29,7 @@ from tilefoundry.ir.tir.sync import participation from tilefoundry.ir.types import ComposedLayout from tilefoundry.ir.types.dim import DimAdd, DimFloorDiv, DimMod, DimMul, DimSub, DimVar -from tilefoundry.ir.types.layout import Layout, LayoutBase, flatten, resolve_layout +from tilefoundry.ir.types.layout import Layout, LayoutBase, flatten, get, resolve_layout from tilefoundry.ir.types.shard_layout import ( Broadcast, Dynamic, @@ -238,9 +238,19 @@ def render_shard_layout_value(var_name: str, sl: SL, dynamic_extents=None, stora ) mesh_base = sl.mesh.layout.offset participation(sl.mesh) - mesh_axes = flatten(sl.mesh.layout) - ml_shape, ml_strides, ml_base = mesh_axes.shape, mesh_axes.strides, int(mesh_base) - topo = program_topologies(sl.mesh)[0] + stated = resolve_layout(sl.mesh.layout) + if isinstance(stated, ComposedLayout): + stated = stated.outer + topos = program_topologies(sl.mesh) + if len(topos) > 1 and mesh_base: + raise NotImplementedError( + "render_shard_layout_value: a mesh naming several levels cannot also " + "be sliced; the slice and the level boundary would both be deciding " + "which positions these are" + ) + ml_base = int(mesh_base) + mesh_axes = flatten(stated) + topo = topos[0] def _static_dim(value, what): if not isinstance(value, int): @@ -267,7 +277,7 @@ def _global_dim(d): return scalar raise NotImplementedError(f"render_shard_layout_value: unsupported layout dim {d!r}") - n_dynamic = sum(1 for d in ml_shape if d is None) + n_dynamic = sum(1 for d in mesh_axes.shape if d is None) if n_dynamic > 1: raise NotImplementedError( "render_shard_layout_value: at most one dynamic (launch-provided) " @@ -288,16 +298,33 @@ def _mesh_dim(d): return "tilefoundry::program_dim()" return _static_dim(d, "mesh layout dim") + def _level_nest(topology): + """One level's arrangement as CuTe shape/stride value nests.""" + index = [t.name for t in sl.mesh.topologies].index(topology.name) + arrangement = get(stated, index) + level_shape = tuple(flatten(arrangement.shape)) + level_strides = tuple(flatten(arrangement.strides)) + if not level_shape: + level_shape, level_strides = (1,), (0,) + shape_args = ", ".join(_mesh_dim(d) for d in level_shape) + stride_args = ", ".join(_static_dim(s, "mesh layout stride") for s in level_strides) + return f"cute::make_shape({shape_args})", f"cute::make_stride({stride_args})" + sl_var = f"{var_name}__sl_layout" ml_var = f"{var_name}__mesh_layout" mesh_var = f"{var_name}__mesh" - ml_shape_args = ", ".join(_mesh_dim(d) for d in ml_shape) - ml_stride_args = ", ".join(_static_dim(s, "mesh layout stride") for s in ml_strides) - - scope = topology_scope_str(topo.name) + if len(topos) == 1: + shape_args = ", ".join(_mesh_dim(d) for d in mesh_axes.shape) + stride_args = ", ".join(_static_dim(s, "mesh layout stride") for s in mesh_axes.strides) + else: + shape_nests, stride_nests = zip(*(_level_nest(t) for t in topos)) + shape_args = ", ".join(shape_nests) + stride_args = ", ".join(stride_nests) + scope = ", ".join(topology_scope_str(t.name) for t in topos) positions = ( - f"cute::make_layout(cute::make_shape({ml_shape_args}), cute::make_stride({ml_stride_args}))" + f"cute::make_layout(cute::make_shape({shape_args}), " + f"cute::make_stride({stride_args}))" ) mesh_layout = _composed_mesh_layout(positions, ml_base) diff --git a/src/tilefoundry/ir/tir/reduce.py b/src/tilefoundry/ir/tir/reduce.py index 96661a6f5..354098c26 100644 --- a/src/tilefoundry/ir/tir/reduce.py +++ b/src/tilefoundry/ir/tir/reduce.py @@ -21,7 +21,8 @@ from tilefoundry.ir.pattern import predicates as P from tilefoundry.ir.tir.sync import WARP_SIZE from tilefoundry.ir.types import DType, ShardLayout, StorageKind, TensorType, UnitType -from tilefoundry.ir.types.layout import flatten +from tilefoundry.ir.types.layout import ComposedLayout, flatten +from tilefoundry.ir.types.mesh import _named, levels from tilefoundry.ir.types.utils import static_dim_value from tilefoundry.visitor_registry import register_typeinfer, register_verify_stmt from tilefoundry.visitor_registry.access_relation import ( @@ -103,13 +104,22 @@ def _(call: "Call", ctx: "TypeInferContext") -> UnitType: def _workspace_slots(source) -> int: - """One slot per warp of the thread mesh ``source`` is sharded over. + """One slot per warp of the thread level ``source`` is sharded over. The runtime posts each warp's partial at its warp index and reads its group - back (`cta_combine_via_workspace`); which warp takes which slot is its choice. + back (`cta_combine_via_workspace`); which warp takes which slot is its + choice. A mesh naming several levels counts only its thread level: each + CTA combines its own partials in its own shared memory. """ + mesh = source.layout.mesh + names = _named(mesh) + if "thread" in names: + level = levels(mesh)[names.index("thread")] + else: + stated = mesh.layout.outer if isinstance(mesh.layout, ComposedLayout) else mesh.layout + level = stated threads = 1 - for extent in flatten(source.layout.mesh.layout).shape: + for extent in flatten(level).shape: value = static_dim_value(extent) if value is None or value <= 0: raise ValueError( diff --git a/tests/ops/tir/test_reduce.py b/tests/ops/tir/test_reduce.py index cc47b9545..7d51608d5 100644 --- a/tests/ops/tir/test_reduce.py +++ b/tests/ops/tir/test_reduce.py @@ -340,3 +340,84 @@ def test_every_tier_answers_what_torch_answers() -> None: "is what shows the fold combining with the reduction's own operator " "rather than adding" ) + + +@module(topologies=(Topology("cta", 4), Topology("thread", 128))) +class _TwoLevel: + """A reduce sharded over a cta×thread mesh, as a multi-CTA finalize writes it. + + The CTA level takes its rows through coordinate arithmetic on the view, so + both operands leave it broadcast; the thread level alone splits a reduced + axis. That is the operand shape every nested cta/thread scope produces. + """ + + @prim_func(target=_CUDA) + def two_level(a_two: Tensor[(8, 4, 128, 4), "f32"], out_two: Tensor[(8, 4, 1, 1), "f32"]): + with Mesh( + (Topology("cta", 4), Topology("thread", 128)), + Layout(shape=(4, 128), strides=(128, 1)), + ("c", "t"), + ) as ml: + two_src = T.tensor_view( + T.ptr_of(a_two[ml.c * 2 : ml.c * 2 + 2, 0:4, 0:128, 0:4]), + layout=ShardLayout( + Layout((2, 4, 128, 4), (2048, 512, 4, 1)), (Broadcast(), Split(2)), ml + ), + shape=(2, 4, 128, 4), + ) + two_dst = T.tensor_view( + T.ptr_of(out_two[ml.c * 2 : ml.c * 2 + 2, 0:4, 0:1, 0:1]), + layout=ShardLayout( + Layout((2, 4, 1, 1), (4, 1, 1, 1)), (Broadcast(), Broadcast()), ml + ), + shape=(2, 4, 1, 1), + ) + two_ws = T.alloc_tensor(Tensor[(4,), "f32", None, "smem"]) + T.reduce( + two_src, + two_dst, + two_ws, + axes=(2, 3), + keepdim=True, + kind=ReduceKind.MEAN, + ) + + +@module(entry="two_level_host", target=_CUDA) +class TwoLevelReduce: + """Launch one CTA per two rows, 128 threads each.""" + + two_level = _TwoLevel + + @prim_func(target=CpuTarget()) + def two_level_host( + a_two: Tensor[(8, 4, 128, 4), "f32"], out_two: Tensor[(8, 4, 1, 1), "f32"] + ): + launch(two_level.two_level, a_two, out_two, grid=(4, 1, 1), block=(128, 1, 1)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_a_cta_thread_mesh_reduces_across_the_thread_level() -> None: + """On a composed cta×thread mesh, the warp count comes from the thread level. + + The dispatch used to keep one ``is_thread`` per mesh, so a cta-scoped + composite counted the thread level's instances as warps -- 128 here, not 4 + -- and the cross-warp tier indexed the workspace out of bounds, while the + verifier demanded one slot per warp of the flattened mesh (16, not 4) and + refused this program before it ran. The plan the levels state is the + intra-CTA tier: 32 lanes and 4 warps per group, one slot per warp, and + eight kept cells combined one at a time. + """ + rm = tilefoundry.compile(TwoLevelReduce, target=_CUDA) + torch.manual_seed(0) + a_two = torch.randn(8, 4, 128, 4, dtype=torch.float32, device="cuda") + out_two = torch.zeros(8, 4, 1, 1, dtype=torch.float32, device="cuda") + rm(a_two, out_two) + torch.cuda.synchronize() + assert torch.allclose(out_two, a_two.mean(dim=(2, 3), keepdim=True), rtol=1e-5, atol=1e-6), ( + "intra-CTA tier on a cta×thread mesh (warps_per_group == 4 read off the " + "thread level, 32 lanes): the mean's divisor is span 4 x lanes 32 x " + "warps 4 = 512 either way, so what a wrong warp count shows is not a " + "wrong divisor but a workspace read out of bounds -- a launch failure " + "here -- or another group's partials folded into the answer" + ) diff --git a/tests/runtime/cuda/static_asserts.cu b/tests/runtime/cuda/static_asserts.cu index dabceeea9..ae2ae990d 100644 --- a/tests/runtime/cuda/static_asserts.cu +++ b/tests/runtime/cuda/static_asserts.cu @@ -285,28 +285,6 @@ __global__ void k(float *p) { } #endif -#if CASE == 13 -/// Item 9: an intra-CTA reduce whose source keeps two cells. -__global__ void k(float *p) { - __shared__ float ws[8]; - auto mesh = tmesh<8, 32>(); - auto slay = cute::make_layout( - cute::make_shape(cute::Int<2>{}, cute::Int<8>{}, cute::Int<32>{}), - cute::make_stride(cute::Int<256>{}, cute::Int<32>{}, cute::Int<1>{})); - auto st = make_shard_tensor( - cute::make_tensor(cute::make_gmem_ptr(p), slay), slay, - make_shard_layout(slay, mesh, - cute::make_tuple(shard::S<1>{}, shard::S<2>{}))); - auto dlay = cute::make_layout(cute::make_shape(cute::Int<2>{})); - auto dt = make_shard_tensor( - cute::make_tensor(cute::make_gmem_ptr(p + 512), dlay), dlay, - make_shard_layout(dlay, mesh, - cute::make_tuple(shard::B{}, shard::B{}))); - auto slots = cute::make_tensor(cute::make_smem_ptr(ws), cute::Int<8>{}); - reduce, cute::Int<2>>>(st, dt, slots); -} -#endif - #if CASE == 14 /// Item 6: a pair of views cp.async cannot move a wide word of. __global__ void k(float *g, int n) { diff --git a/tests/runtime/test_cuda_static_asserts.py b/tests/runtime/test_cuda_static_asserts.py index 2237ea73b..2d3cbe99d 100644 --- a/tests/runtime/test_cuda_static_asserts.py +++ b/tests/runtime/test_cuda_static_asserts.py @@ -32,7 +32,6 @@ 10: "the source's projected layout must be a plain cute::Layout", 11: "the two operands must name one mesh", 12: "ops::reduce: a reduce mesh's scope must be cta or thread", - 13: "ops::reduce (intra-CTA tier): kept axes must form one cell", 14: "these two views share no run wide enough for cp.async", 15: "the two projected slices must hold the same number of elements", 16: "ops::dot (warp tier): the fastest axis of the operands' mesh", From 94ac8f07ff1f78e1d245ad734e489c4451e3afff Mon Sep 17 00:00:00 2001 From: sablin39 <1020030829@qq.com> Date: Wed, 21 Oct 2026 22:28:12 +0800 Subject: [PATCH 2/3] fix(runtime): project plain copy operands through the pairing shard layout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ops::copy projected each operand through its own shard layout alone, so a plain operand paired with a sharded one stayed whole for every instance. T.copy(fragment, window) -- what finalize emits for the RMSNorm tutorial's reshard to gmem on its cta×thread mesh -- then ran cute::copy on mismatched sizes: the whole window iterated against the fragment's modular index space, every instance tiling its 40-element fragment across the 10240-element window and racing its peers. Where the vectorizer keeps the fragment's reads contiguous instead, the same mismatch steps past the fragment -- the tutorial's "Invalid __local__ read of size 2" under compute-sanitizer. When exactly one operand is sharded, the plain side moves the slice the shard side's layout hands this instance: - local_as resolves this instance's window and origin from the shard side's mesh coordinates. Window and origin are row-major logical steps over the shard layout's shape, while CuTe decodes an integer index first-mode-fastest, so the plain tensor's layout goes through the runtime's reverse (flattened first, so nested modes turn too; a ComposedLayout overload keeps a swizzle in front of the modes): CuTe's colexicographic index then is the row-major one, and the composition keeps whatever addressing the plain tensor states. - Copy dispatches on ShardTensorLike: exactly-one-sharded pairs take the projection; shard/shard and plain/plain pairs keep the per-operand local_tensor path unchanged. The regression test round-trips a (2, 5120) bf16 window through a (2, 4, 32)-instance rmem shard and back on one 256-thread CTA; a second test copies one way into a column-pair shard, so a slice placed wrong on the way in lands in the wrong output columns directly -- the round trip cannot see a permutation the write and the read share. The round trip fails with no projection (the window comes back tiled by whichever instance writes last), the one-way test also fails with the modes handed to CuTe unreversed (the output comes back transposed), and both pass after, bit for bit, with compute-sanitizer reporting 0 errors. --- docs/spec/runtime.md | 4 + .../runtime/cuda/layout/cute_ext.cuh | 9 ++ .../runtime/cuda/ops/copy/copy_impl.h | 42 ++++- tests/ops/tir/test_copy.py | 148 ++++++++++++++++++ 4 files changed, 200 insertions(+), 3 deletions(-) diff --git a/docs/spec/runtime.md b/docs/spec/runtime.md index b50b5bd99..4d0a92a56 100644 --- a/docs/spec/runtime.md +++ b/docs/spec/runtime.md @@ -950,6 +950,10 @@ CUTE_HOST_DEVICE auto program_ids() noexcept; // include/tilefoundry/runtime/cuda/layout/cute_ext.cuh template CUTE_HOST_DEVICE constexpr auto reverse(cute::Layout const &l); + +template +CUTE_HOST_DEVICE constexpr auto +reverse(cute::ComposedLayout const &l); ``` diff --git a/include/tilefoundry/runtime/cuda/layout/cute_ext.cuh b/include/tilefoundry/runtime/cuda/layout/cute_ext.cuh index 1a9f6b55a..9d5971585 100644 --- a/include/tilefoundry/runtime/cuda/layout/cute_ext.cuh +++ b/include/tilefoundry/runtime/cuda/layout/cute_ext.cuh @@ -15,3 +15,12 @@ CUTE_HOST_DEVICE constexpr auto reverse(cute::Layout const &l) { return cute::make_layout(cute::reverse(cute::shape(l)), cute::reverse(cute::stride(l))); } + +/// A swizzled layout turns the same way, the swizzle kept in front of the +/// modes it acts on. +template +CUTE_HOST_DEVICE constexpr auto +reverse(cute::ComposedLayout const &l) { + return cute::make_composed_layout(l.layout_a(), l.offset(), + reverse(cute::flatten(l.layout_b()))); +} diff --git a/include/tilefoundry/runtime/cuda/ops/copy/copy_impl.h b/include/tilefoundry/runtime/cuda/ops/copy/copy_impl.h index 94f04296a..850df696f 100644 --- a/include/tilefoundry/runtime/cuda/ops/copy/copy_impl.h +++ b/include/tilefoundry/runtime/cuda/ops/copy/copy_impl.h @@ -4,13 +4,49 @@ namespace copy_impl { +/// The slice of a plain operand that pairs with one shard operand. +/// +/// An operand with no shard layout of its own is not therefore whole for +/// every instance: when the other operand is sharded, the plain side moves +/// exactly the slice that instance owns, projected through the shard side's +/// layout, whose shape refines the shared logical domain. The projection's +/// window and origin are row-major logical steps, and the plain layout is +/// reversed before they index it, so CuTe's first-mode-fastest decode reads +/// them as the row-major index they are. +template +CUTE_HOST_DEVICE auto local_as(T &&t, + tilefoundry::ShardLayout const &sl) { + auto logical = + cute::make_layout(cute::shape(sl.layout_value), cute::GenRowMajor{}); + tilefoundry::ShardLayout as_logical{logical, + sl.mesh_value}; + auto const [window, origin] = tilefoundry::detail::local_layout_and_offset( + as_logical, + tilefoundry::mesh_coords(sl.mesh_value, tilefoundry::program_ids())); + auto const plain = reverse(cute::flatten(t.layout())); + return cute::make_tensor(t.data() + plain(origin), + cute::composition(plain, window)); +} + /// ``copy`` is ``cute::copy`` with the operands resolved to this slice first. struct Copy { template __device__ void operator()(TSrc const &src, TDst &dst) const { - auto &&s = tilefoundry::local_tensor(src); - auto &&d = tilefoundry::local_tensor(dst); - cute::copy(s, d); + if constexpr (tilefoundry::ShardTensorLike && + !tilefoundry::ShardTensorLike) { + auto &&s = tilefoundry::local_tensor(src); + auto d = local_as(dst, src.shard_layout); + cute::copy(s, d); + } else if constexpr (!tilefoundry::ShardTensorLike && + tilefoundry::ShardTensorLike) { + auto s = local_as(src, dst.shard_layout); + auto &&d = tilefoundry::local_tensor(dst); + cute::copy(s, d); + } else { + auto &&s = tilefoundry::local_tensor(src); + auto &&d = tilefoundry::local_tensor(dst); + cute::copy(s, d); + } } }; diff --git a/tests/ops/tir/test_copy.py b/tests/ops/tir/test_copy.py index 6a28f9a01..97ac87887 100644 --- a/tests/ops/tir/test_copy.py +++ b/tests/ops/tir/test_copy.py @@ -151,3 +151,151 @@ def test_every_crossing_round_trips_the_input_bit_for_bit() -> None: "gmem->rmem at 64 bits: a sub-128-bit fragment falls back to the element " "loop, which must answer identically to the vector path" ) + + +@module( + target=_CUDA, + topologies=(Topology("cta", 1), Topology("thread", 256)), +) +class _RegisterRoundTrip: + """The optimize tutorial's register round trip, at its own mesh and shard. + + A (2, 5120) bf16 window on cta 1 x thread (2, 4, 32): the row axis over + ``x``, the 40 blocks over ``y``, the 128-wide block over ``t`` -- a + 40-element fragment per thread, and both crossings of the plain/shard + boundary, since the tutorial's output copy keeps its window plain. + """ + + @prim_func(target=_CUDA) + def round_trip(a: Tensor[(2, 5120), "bf16"], out: Tensor[(2, 5120), "bf16"]): + with Mesh((Topology("cta", 1),), layout=(1,), names=("tile",)) as _cta: + with Mesh( + (Topology("thread", 256),), layout=(2, 4, 32), names=("x", "y", "t") + ) as scope: + value = T.alloc_tensor( + tensor_type=Tensor[ + (2, 5120), + "bf16", + (2 @ scope.x, 40 @ scope.y, 128 @ scope.t), + "rmem", + ] + ) + dst_view = T.tensor_view( + T.ptr_of(value), + layout=(2 @ scope.x, 40 @ scope.y, 128 @ scope.t), + shape=(2, 5120), + ) + T.copy(a, dst_view) + src_view = T.tensor_view( + T.ptr_of(value), + layout=(2 @ scope.x, 40 @ scope.y, 128 @ scope.t), + shape=(2, 5120), + ) + T.copy(src_view, out) + + +@module(entry="round_trip_host", target=_CUDA) +class RegisterRoundTrip: + """Launch the one CTA the shard is written for.""" + + round_trip = _RegisterRoundTrip + + @prim_func(target=CpuTarget()) + def round_trip_host(a: Tensor[(2, 5120), "bf16"], out: Tensor[(2, 5120), "bf16"]): + launch( # noqa: F821 + round_trip.round_trip, # noqa: F821 + a, + out, + grid=(1, 1, 1), + block=(256, 1, 1), + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_a_plain_window_round_trips_through_thread_fragments() -> None: + """A copy with one plain and one sharded operand moves the instance's slice. + + ``local_tensor`` projects each operand through its own shard layout, so a + plain operand used to stay whole for every instance: the gmem->rmem leg + read the window's first 40 elements into every fragment, and the rmem->gmem + leg tiled the fragment across the whole window (or, vectorized, read it out + of bounds -- the tutorial's ``Invalid __local__ read of size 2``). The + plain side is the slice the shard side's layout hands this instance. + """ + rm = tilefoundry.compile(RegisterRoundTrip, target=_CUDA) + torch.manual_seed(0) + a = torch.randn(2, 5120, dtype=torch.bfloat16, device="cuda") + out = torch.zeros_like(a) + rm(a, out) + torch.cuda.synchronize() + assert torch.equal(out, a), ( + "plain->shard->plain round trip: each thread holds the 40 elements its " + "(x, y, t) coordinate names, so the window must reassemble exactly -- " + "every thread copying the window's head, or tiling its fragment across " + "the whole window, is the boundary projected wrong" + ) + + +@module( + target=_CUDA, + topologies=(Topology("thread", 32),), +) +class _OneWay: + """A plain (4, 64) tile entering a column-pair shard, one direction only. + + Lane ``t`` of 32 owns columns ``2t`` and ``2t+1`` of every row: the + plain->shard leg takes the ``local_as`` projection, the shard->shard leg + the path that predates it, so a slice placed wrong on the way in lands in + the wrong output columns directly. + """ + + @prim_func(target=_CUDA) + def one_way(a: Tensor[(4, 64), "f32"], out: Tensor[(4, 64), "f32"]): + with Mesh((Topology("thread", 32),), layout=(32,), names=("t",)) as m: + src = T.tensor_view(T.ptr_of(a), layout=((4, 64), (64, 1))) + reg = T.alloc_tensor( + tensor_type=Tensor[(4, 64), "f32", ((4, 32 @ m.t, 2), (64, 2, 1)), "rmem"] + ) + dst = T.tensor_view(T.ptr_of(out), layout=((4, 32 @ m.t, 2), (64, 2, 1))) + T.copy(src, reg) + T.copy(reg, dst) + + +@module(entry="one_way_host", target=_CUDA) +class OneWay: + """Launch the one warp the shard is written for.""" + + one_way = _OneWay + + @prim_func(target=CpuTarget()) + def one_way_host(a: Tensor[(4, 64), "f32"], out: Tensor[(4, 64), "f32"]): + launch( # noqa: F821 + one_way.one_way, # noqa: F821 + a, + out, + grid=(1, 1, 1), + block=(32, 1, 1), + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_a_plain_window_enters_thread_fragments_in_row_major_order() -> None: + """The projection's row-major steps must not come out transposed. + + CuTe decodes an integer index first-mode-fastest, so composing the + instance's row-major window and origin into the plain tensor's rank-2 + layout unreversed hands every lane a permutation of its slice: lane 1 of + a (2, 4) tile reads element (1, 0) where it owns (0, 2). A round trip + hides it -- write and read share the permutation -- so this test copies + in once and compares against the input. + """ + rm = tilefoundry.compile(OneWay, target=_CUDA) + a = torch.arange(4 * 64, dtype=torch.float32, device="cuda").view(4, 64) + out = torch.zeros_like(a) + rm(a, out) + torch.cuda.synchronize() + assert torch.equal(out, a), ( + "plain->shard with arange input: lane t must hold columns 2t and " + "2t+1 of each row, so a transposed or otherwise permuted slice shows " + "up as out != a exactly where the projection misread the index" + ) From 666aae3297a44b2676a73ca6702e63f505fa57c4 Mon Sep 17 00:00:00 2001 From: sablin39 <1020030829@qq.com> Date: Wed, 21 Oct 2026 22:28:55 +0800 Subject: [PATCH 3/3] fix(codegen): emit TupleGetItem pointers and nested layouts in tensor_view More of the #230 plumbing. _pointer_ref assumed a PtrOf source is always a named tensor: anything else fell through to _tensor_ref, which reads .name and died with "AttributeError: 'Call' object has no attribute 'name'" on the TupleGetItem stage choices WGMMA pipelines emit (tests/fixtures/schedule/tir/wgmma_rs_a_from_smem.py is the tracked reproducer) and on the TensorView-call and Constant sources nested under them. _plain_layout_value rendered shapes and strides as flat integer lists, so a layout with a nested mode died with "TypeError: int() argument must be ... not 'tuple'". Of the schedule TIR fixtures, twenty-one stopped at the first and three at the second. - TupleGetItem sources go through _tuple_pointer_ref: a constant index resolves to the chosen element's pointer; a runtime index renders through render_tuple_choice -- one cute::array of the elements subscripted by the index it renders, with .data() taking the chosen tensor's pointer, or an array of the pointers themselves when the elements name raw smem offsets. cute::array has no deduction guide, so the element type and count are spelled from the first element's rendering. A runtime choice has no single syntactic source, so no memory_var comes back for the coordinate forms to bind to. TensorView-call and Constant sources recurse instead of falling through. - One recursive renderer (codegen/cuda/tir/cute_layout.py) spells an int-tuple shape and stride in both forms -- the cute::Layout type nesting cute::Shape/cute::Stride, the make_layout value nesting make_shape/make_stride, with callables rendering the scalar entries -- and every layout emitter routes through it: tensor_view's layout type and value, its local_tile source layout, the shard layout's mesh positions, mesh_scope's levelwise types, and arith's broadcast view. Flat layouts render byte-identical to before. - Two emission fixes the executable runs pinned: fill rendered an integer constant as "0f", and a tuple index reading an rmem scalar named the whole rank-0 view where the arithmetic needs the held value; non-UMAT scalars now render subscripted. Every fixture that emitted still emits, and all twenty-four stopped ones now pass the pointer stage: eighteen reach the deliberate "CopyAsyncTensor: no emitter without host-encoded tensor maps" refusal (the missing TMA atoms, #249), two newly emit, one reaches the sliced-multi-level-mesh refusal from the reduce warps fix, and the last three refuse in their own words -- swizzled slice bases and a nested shard local shape stay out of scope here. The two newly-emitting fixtures, staged_dynamic_inner and staged_nested, are re-authored as executable tests: their stage rings compile and run on the card, and the staged tiles are checked to arrive intact. --- src/tilefoundry/codegen/cuda/context.py | 11 ++ src/tilefoundry/codegen/cuda/tir/arith.py | 11 +- .../codegen/cuda/tir/async_copy.py | 4 +- .../codegen/cuda/tir/cute_layout.py | 43 +++++ .../codegen/cuda/tir/memory/fill.py | 6 +- .../codegen/cuda/tir/memory/tensor_view.py | 140 +++++++++------ .../codegen/cuda/tir/stmts/mesh_scope.py | 40 ++--- .../codegen/cuda/tir/tuple_get_item.py | 37 +++- src/tilefoundry/ir/types/mesh.py | 16 ++ tests/codegen/test_staged_ring.py | 168 ++++++++++++++++++ 10 files changed, 381 insertions(+), 95 deletions(-) create mode 100644 src/tilefoundry/codegen/cuda/tir/cute_layout.py create mode 100644 tests/codegen/test_staged_ring.py diff --git a/src/tilefoundry/codegen/cuda/context.py b/src/tilefoundry/codegen/cuda/context.py index 52f276416..be26ae4da 100644 --- a/src/tilefoundry/codegen/cuda/context.py +++ b/src/tilefoundry/codegen/cuda/context.py @@ -48,6 +48,17 @@ def topology_scope_str(name: str) -> str: ) from None +def float_literal(value) -> str: + """A C++ float literal spelling of a Python number.""" + return repr(float(value)) + "f" + + +def scalar_read(var, ctx: "CudaCodegenContext") -> str: + """Read element zero of a one-element tensor variable, as a scalar.""" + name = ctx.name_for(var) + return f"{name}_tensor(0)" if ctx.is_kernel_param(var) else f"{name}(0)" + + class CudaCodegenContext(EmitContext): """A compile writing CUDA: the shared context plus what only CUDA states.""" diff --git a/src/tilefoundry/codegen/cuda/tir/arith.py b/src/tilefoundry/codegen/cuda/tir/arith.py index 398165fb7..fa412f779 100644 --- a/src/tilefoundry/codegen/cuda/tir/arith.py +++ b/src/tilefoundry/codegen/cuda/tir/arith.py @@ -3,6 +3,7 @@ from __future__ import annotations from tilefoundry.codegen.cuda.context import CudaCodegenContext +from tilefoundry.codegen.cuda.tir.cute_layout import render_layout_value from tilefoundry.ir.tir.arith import Binary, BinaryKind, Unary, UnaryKind from tilefoundry.ir.tir.cast import Cast from tilefoundry.ir.types.shard_layout import ShardLayout, shard_layout_local_shape @@ -87,6 +88,11 @@ def _extent(e) -> str: return f"cute::Int<{int(e)}>{{}}" if isinstance(e, int) else str(e) +def _stride(e) -> str: + """A cute stride entry: broadcast modes are always static.""" + return f"cute::Int<{int(e)}>{{}}" + + def _broadcast_modes(dst_shape, rhs_shape, n_dst_runtime) -> tuple | None: """The destination-shaped layout *rhs* is read through, or ``None``. @@ -123,12 +129,9 @@ def broadcast_view(rhs_n: str, modes, ctx: CudaCodegenContext) -> str: shape, stride = modes ctx._counter += 1 name = f"bcast_{ctx._counter}" - shape_args = ", ".join(_extent(s) for s in shape) - stride_args = ", ".join(f"cute::Int<{int(s)}>{{}}" for s in stride) ctx.emit( f"auto {name} = tilefoundry::local_tensor({rhs_n}).compose(" - f"cute::make_layout(cute::make_shape({shape_args}), " - f"cute::make_stride({stride_args})));" + f"{render_layout_value(shape, stride, _extent, _stride)});" ) return name diff --git a/src/tilefoundry/codegen/cuda/tir/async_copy.py b/src/tilefoundry/codegen/cuda/tir/async_copy.py index 985f089c4..7abb9d94e 100644 --- a/src/tilefoundry/codegen/cuda/tir/async_copy.py +++ b/src/tilefoundry/codegen/cuda/tir/async_copy.py @@ -8,7 +8,7 @@ import math -from tilefoundry.codegen.cuda.context import CudaCodegenContext +from tilefoundry.codegen.cuda.context import CudaCodegenContext, float_literal from tilefoundry.ir.tir.async_copy import ( CopyAsync, CpAsyncCommit, @@ -41,7 +41,7 @@ def _emit_copy_async(call, ctx: CudaCodegenContext) -> None: elif math.isnan(fill): literal = "__int_as_float(0x7fc00000)" else: - literal = repr(float(fill)) + "f" + literal = float_literal(fill) suffix = f", {literal}" ctx.emit( f"tilefoundry::ops::copy_async<{mesh_type}, {width}>({src}, {dst}, {index}{suffix});" diff --git a/src/tilefoundry/codegen/cuda/tir/cute_layout.py b/src/tilefoundry/codegen/cuda/tir/cute_layout.py new file mode 100644 index 000000000..2fa3f285a --- /dev/null +++ b/src/tilefoundry/codegen/cuda/tir/cute_layout.py @@ -0,0 +1,43 @@ +"""One recursive renderer for a CuTe layout, in its type and value forms. + +A layout is an int-tuple shape and an int-tuple stride; sub-tuples nest +``cute::Shape`` / ``cute::make_shape`` with them, recursively. Every emitter +spells a layout through these two functions, so a nested mode -- an operand +tile, a multi-level mesh -- renders everywhere a flat one does. +""" + +from __future__ import annotations + + +def render_layout_type(shape, stride) -> str: + """The ``cute::Layout`` type of int-tuple *shape* and *stride*.""" + + def nest(values, tag): + return ", ".join( + f"{tag}<{nest(value, tag)}>" if isinstance(value, tuple) else f"cute::Int<{value}>" + for value in values + ) + + return ( + f"cute::Layout, " + f"cute::Stride<{nest(stride, 'cute::Stride')}>>" + ) + + +def render_layout_value(shape, stride, dim, stride_dim) -> str: + """The ``cute::make_layout`` value of int-tuple *shape* and *stride*. + + *dim* and *stride_dim* render one shape entry and one stride entry, which + is where a runtime-provided extent reaches the emitted layout. + """ + + def nest(values, render, tag): + return ", ".join( + f"{tag}({nest(value, render, tag)})" if isinstance(value, tuple) else render(value) + for value in values + ) + + return ( + f"cute::make_layout(cute::make_shape({nest(shape, dim, 'cute::make_shape')}), " + f"cute::make_stride({nest(stride, stride_dim, 'cute::make_stride')}))" + ) diff --git a/src/tilefoundry/codegen/cuda/tir/memory/fill.py b/src/tilefoundry/codegen/cuda/tir/memory/fill.py index 8a2935519..693f4f4a0 100644 --- a/src/tilefoundry/codegen/cuda/tir/memory/fill.py +++ b/src/tilefoundry/codegen/cuda/tir/memory/fill.py @@ -1,7 +1,7 @@ """Codegen for ``tir.memory.Fill`` — the zero-source ``elementwise``.""" from __future__ import annotations -from tilefoundry.codegen.cuda.context import CudaCodegenContext +from tilefoundry.codegen.cuda.context import CudaCodegenContext, float_literal from tilefoundry.codegen.cuda.tir.memory.tensor_view import _tensor_ref from tilefoundry.ir.core import Constant from tilefoundry.ir.tir.memory import Fill @@ -14,6 +14,4 @@ def _emit(call, ctx: CudaCodegenContext) -> None: tensor, value = call.args[0], call.args[1] dst_n = _tensor_ref(tensor, ctx) val = value.value if isinstance(value, Constant) else 0.0 - ctx.emit( - f"tilefoundry::ops::elementwise({dst_n}, []() {{ return {val}f; }});" - ) + ctx.emit(f"tilefoundry::ops::elementwise({dst_n}, []() {{ return {float_literal(val)}; }});") diff --git a/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py b/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py index dfe759678..bdcaadefb 100644 --- a/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py +++ b/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py @@ -12,24 +12,29 @@ from tilefoundry.codegen.cuda.context import ( CudaCodegenContext, + scalar_read, topology_scope_str, ) +from tilefoundry.codegen.cuda.tir.cute_layout import render_layout_type, render_layout_value from tilefoundry.codegen.cuda.tir.reduce import REDUCE_TAG from tilefoundry.codegen.cuda.tir.stmts.mesh_scope import ( mesh_type, program_topologies, ) -from tilefoundry.ir.core import Call, Constant, Tuple +from tilefoundry.codegen.cuda.tir.tuple_get_item import constant_index, render_tuple_choice +from tilefoundry.ir.core import Call, Constant, Tuple, Var from tilefoundry.ir.core.kinds import ReduceKind from tilefoundry.ir.hir.sharding.mesh_coord import MeshCoord from tilefoundry.ir.hir.tensor.slice import Slice +from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem from tilefoundry.ir.tir.memory.ptr_of import PtrOf from tilefoundry.ir.tir.memory.tensor_view import TensorView from tilefoundry.ir.tir.stmts import LetStmt from tilefoundry.ir.tir.sync import participation from tilefoundry.ir.types import ComposedLayout from tilefoundry.ir.types.dim import DimAdd, DimFloorDiv, DimMod, DimMul, DimSub, DimVar -from tilefoundry.ir.types.layout import Layout, LayoutBase, flatten, get, resolve_layout +from tilefoundry.ir.types.layout import Layout, LayoutBase, flatten, resolve_layout +from tilefoundry.ir.types.mesh import flattened_levels from tilefoundry.ir.types.shard_layout import ( Broadcast, Dynamic, @@ -58,7 +63,8 @@ def _render_layout_type(layout: LayoutBase) -> str: A swizzled composed layout keeps its shape: CuTe states it as a ``ComposedLayout`` of the swizzle, a static offset and the layout underneath. The XOR mapping is not expressible as strides, so there is no - affine form to fall back to. + affine form to fall back to. CuTe encodes the single rank-0 address as a + singleton mode. """ layout = resolve_layout(layout) swizzle = get_swizzle_portion(layout) @@ -70,9 +76,7 @@ def _render_layout_type(layout: LayoutBase) -> str: if isinstance(layout, Layout): if not layout.shape: return "cute::Layout>, cute::Stride>>" - shape_args = ", ".join(f"cute::Int<{s}>" for s in layout.shape) - stride_args = ", ".join(f"cute::Int<{s}>" for s in layout.strides) - return f"cute::Layout, cute::Stride<{stride_args}>>" + return render_layout_type(layout.shape, layout.strides) raise NotImplementedError(f"tensor_view: no CuTe layout type for {type(layout).__name__}") @@ -95,11 +99,7 @@ def _render_layout_value(layout: LayoutBase, dim, stride) -> str: if isinstance(layout, Layout): if not layout.shape: return "cute::make_layout(cute::Shape>{}, cute::Stride>{})" - shape_args = ", ".join(dim(d) for d in layout.shape) - stride_args = ", ".join(stride(s) for s in layout.strides) - return ( - f"cute::make_layout(cute::make_shape({shape_args}), cute::make_stride({stride_args}))" - ) + return render_layout_value(layout.shape, layout.strides, dim, stride) raise NotImplementedError(f"tensor_view: no CuTe layout value for {type(layout).__name__}") @@ -298,34 +298,21 @@ def _mesh_dim(d): return "tilefoundry::program_dim()" return _static_dim(d, "mesh layout dim") - def _level_nest(topology): - """One level's arrangement as CuTe shape/stride value nests.""" - index = [t.name for t in sl.mesh.topologies].index(topology.name) - arrangement = get(stated, index) - level_shape = tuple(flatten(arrangement.shape)) - level_strides = tuple(flatten(arrangement.strides)) - if not level_shape: - level_shape, level_strides = (1,), (0,) - shape_args = ", ".join(_mesh_dim(d) for d in level_shape) - stride_args = ", ".join(_static_dim(s, "mesh layout stride") for s in level_strides) - return f"cute::make_shape({shape_args})", f"cute::make_stride({stride_args})" + def _mesh_stride(s): + return _static_dim(s, "mesh layout stride") sl_var = f"{var_name}__sl_layout" ml_var = f"{var_name}__mesh_layout" mesh_var = f"{var_name}__mesh" if len(topos) == 1: - shape_args = ", ".join(_mesh_dim(d) for d in mesh_axes.shape) - stride_args = ", ".join(_static_dim(s, "mesh layout stride") for s in mesh_axes.strides) + mesh_shape, mesh_strides = mesh_axes.shape, mesh_axes.strides else: - shape_nests, stride_nests = zip(*(_level_nest(t) for t in topos)) - shape_args = ", ".join(shape_nests) - stride_args = ", ".join(stride_nests) + arrangements = flattened_levels(replace(sl.mesh, layout=stated)) + mesh_shape = tuple(shape for shape, _ in arrangements) + mesh_strides = tuple(strides for _, strides in arrangements) scope = ", ".join(topology_scope_str(t.name) for t in topos) - positions = ( - f"cute::make_layout(cute::make_shape({shape_args}), " - f"cute::make_stride({stride_args}))" - ) + positions = render_layout_value(mesh_shape, mesh_strides, _mesh_dim, _mesh_stride) mesh_layout = _composed_mesh_layout(positions, ml_base) attrs = ", ".join(_render_attr(a) for a in sl.attrs) @@ -367,13 +354,12 @@ def visit_Call(self, expr: Call, ctx: CudaCodegenContext) -> str: ) def _leaf(self, expr, ctx: CudaCodegenContext) -> str: - name = ctx.name_for(expr) shape = getattr(getattr(expr, "type", None), "shape", ()) or () dims = tuple(getattr(d, "value", d) for d in shape) if dims == (): - return name + return ctx.name_for(expr) if dims == (1,): - return f"{name}_tensor(0)" if ctx.is_kernel_param(expr) else f"{name}(0)" + return scalar_read(expr, ctx) raise NotImplementedError( f"local_tile coordinate from a rank-{len(dims)} offset {dims} is not supported" ) @@ -400,6 +386,53 @@ def _tensor_ref(var, ctx: CudaCodegenContext) -> str: return f"{name}_tensor" if ctx.is_kernel_param(var) else name +def _tensor_value_ref(element, ctx: CudaCodegenContext) -> str: + """The tensor one tuple element names: its bound name, or the view it calls. + + An inline ``T.tensor_view(...)`` element has no binding of its own -- a + Tuple lowers at its consumers -- so the view expression is rendered in + place, in the words a plain ``TensorView`` let would have used. + """ + if isinstance(element, Call) and isinstance(element.target, TensorView): + pointer_ref, _ = _pointer_ref(element.args[0], ctx) + layout = resolve_layout(element.target.layout, element.type.shape) + if isinstance(layout, SL): + raise NotImplementedError( + "CUDA TupleGetItem pointer: an inline shard-layout tensor_view " + "element emits several statements and has no inline expression" + ) + return f"tilefoundry::ops::tensor_view({pointer_ref}, {_plain_layout_value(layout)})" + return _tensor_ref(element, ctx) + + +def _tuple_pointer_ref(call: Call, ctx: CudaCodegenContext) -> tuple[str, object | None]: + """Render a pointer chosen among a tuple's elements. + + A constant index is decided here; anything else is decided on the card as + one array subscript, over the elements' tensors with ``.data()`` taking + the chosen one's pointer, or over the pointers themselves when the + elements name raw smem offsets. A runtime choice has no single syntactic + source, so no ``memory_var`` comes back for the coordinate forms to bind + to. + """ + held, index = call.args + if isinstance(held, Var): + held = ctx.tuple_for(held) + if not isinstance(held, Tuple): + raise ValueError("CUDA TupleGetItem pointer requires a tuple value") + if isinstance(index, Constant): + selected = constant_index(index) + return _pointer_ref(Call(PtrOf(), (held.elements[selected],), type=call.type), ctx) + if any(isinstance(element, Constant) for element in held.elements): + refs = tuple( + _pointer_ref(Call(PtrOf(), (element,), type=call.type), ctx)[0] + for element in held.elements + ) + return render_tuple_choice(refs, index, ctx), None + tensors = tuple(_tensor_value_ref(element, ctx) for element in held.elements) + return f"{render_tuple_choice(tensors, index, ctx)}.data()", None + + def _pointer_ref(pointer, ctx: CudaCodegenContext) -> tuple[str, object | None]: """Render a TensorView pointer and return its syntactic source tensor, if any.""" if isinstance(pointer, Call) and isinstance(pointer.target, PtrOf): @@ -415,6 +448,12 @@ def _pointer_ref(pointer, ctx: CudaCodegenContext) -> tuple[str, object | None]: offset = f"{_plain_layout_value(layout)}(cute::make_coord({coords}))" pointer, root = _pointer_ref(Call(PtrOf(), (base,), type=pointer.type), ctx) return f"({pointer} + {offset})", root + if isinstance(source, Call) and isinstance(source.target, TupleGetItem): + return _tuple_pointer_ref(source, ctx) + if isinstance(source, Call) and isinstance(source.target, TensorView): + return _pointer_ref(source.args[0], ctx) + if isinstance(source, Constant): + return _pointer_ref(source, ctx) return f"{_tensor_ref(source, ctx)}.data()", source if isinstance(pointer, Constant): cpp_type = ctx.dtype_to_cpp(pointer.type.dtype.name) @@ -426,14 +465,18 @@ def _pointer_ref(pointer, ctx: CudaCodegenContext) -> tuple[str, object | None]: return ctx.name_for(pointer), None -def _plain_layout_value(layout: LayoutBase) -> str: - def dim(value): - return f"cute::Int<{int(upper_bound(value))}>{{}}" +def _envelope_dim(value) -> str: + """A shape entry at its static envelope bound, as a CuTe value.""" + return f"cute::Int<{int(upper_bound(value))}>{{}}" - def stride(value): - return f"cute::Int<{int(value)}>{{}}" - return _render_layout_value(layout, dim, stride) +def _static_value(value) -> str: + """A static stride entry as a CuTe value.""" + return f"cute::Int<{int(value)}>{{}}" + + +def _plain_layout_value(layout: LayoutBase) -> str: + return _render_layout_value(layout, _envelope_dim, _static_value) @register_codegen(CudaTarget, Role.EMIT, TensorView) @@ -477,21 +520,14 @@ def _emit(let: LetStmt, ctx: CudaCodegenContext) -> None: else: tensor_ref = f"{mem_name}_tensor" if ctx.is_kernel_param(memory_var) else mem_name source_shape = tuple(memory_var.type.shape) - source_shape_args = ", ".join( - f"cute::Int<{int(upper_bound(dim))}>{{}}" for dim in source_shape - ) - source_stride_args = ", ".join( - f"cute::Int<{stride}>{{}}" - for stride in resolve_layout( - None, tuple(int(upper_bound(dim)) for dim in source_shape) - ).strides - ) + source_strides = resolve_layout( + None, tuple(int(upper_bound(dim)) for dim in source_shape) + ).strides source_name = f"{var_name}__source" ctx.emit( f"auto {source_name} = cute::make_tensor(" - f"{tensor_ref}.data(), cute::make_layout(" - f"cute::make_shape({source_shape_args}), " - f"cute::make_stride({source_stride_args})));" + f"{tensor_ref}.data(), " + f"{render_layout_value(source_shape, source_strides, _envelope_dim, _static_value)});" ) tensor_ref = source_name shape = tuple(let.var.type.shape) diff --git a/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py b/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py index 725f5dcbb..ec7b547ba 100644 --- a/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py +++ b/src/tilefoundry/codegen/cuda/tir/stmts/mesh_scope.py @@ -9,9 +9,10 @@ from dataclasses import replace from tilefoundry.codegen.cuda.context import topology_scope_str +from tilefoundry.codegen.cuda.tir.cute_layout import render_layout_type from tilefoundry.ir.tir.sync import participation -from tilefoundry.ir.types.layout import ComposedLayout, Layout, flatten, get, resolve_layout -from tilefoundry.ir.types.mesh import Mesh, Topology +from tilefoundry.ir.types.layout import ComposedLayout, Layout, flatten, resolve_layout +from tilefoundry.ir.types.mesh import Mesh, Topology, flattened_levels from tilefoundry.target import validate_cuda_topology_levels @@ -40,28 +41,19 @@ def _validate_topology(mesh: Mesh, target) -> None: validate_cuda_topology_levels(target, (_resolved(t).name for t in mesh.topologies)) -def _levelwise_layout(mesh: Mesh, topos) -> str: +def _levelwise_layout(mesh: Mesh) -> str: """One nest per topology level, each in that level's own numbering. A mesh states one arrangement per level already, each in that level's own numbering, so a nest is that arrangement written out: nothing here decides where a boundary falls. The composite is a grouping, not a map: only the - nests are evaluated. A level every instance shares keeps a mode of one, so - that ``get`` still has something to pick. + nests are evaluated. """ - mesh = replace(mesh, layout=resolve_layout(mesh.layout)) - shapes, strides = [], [] - for topology in topos: - stated = mesh.layout.outer if isinstance(mesh.layout, ComposedLayout) else mesh.layout - index = [t.name for t in mesh.topologies].index(topology.name) - arrangement = get(stated, index) - level_shape = tuple(flatten(arrangement.shape)) - level_strides = tuple(flatten(arrangement.strides)) - if not level_shape: - level_shape, level_strides = (1,), (0,) - shapes.append("cute::Shape<" + ", ".join(f"cute::Int<{s}>" for s in level_shape) + ">") - strides.append("cute::Stride<" + ", ".join(f"cute::Int<{s}>" for s in level_strides) + ">") - return f"cute::Layout, cute::Stride<{', '.join(strides)}>>" + arrangements = flattened_levels(mesh) + return render_layout_type( + tuple(shape for shape, _ in arrangements), + tuple(strides for _, strides in arrangements), + ) def mesh_type(mesh: Mesh) -> str: @@ -76,16 +68,14 @@ def mesh_type(mesh: Mesh) -> str: "CUDA mesh emission: a sliced mesh needs its participating box " "as a strided Layout; this states an identity box, with no sub-box" ) - stated, base = outer, participation(mesh).base + base = participation(mesh).base else: if layout_value.strides is None: raise NotImplementedError("CUDA mesh emission: mesh layout needs strides") - stated, base = layout_value, 0 + base = 0 if len(topos) == 1: - level = get(stated, 0) - shape_types = ", ".join(f"cute::Int<{s}>" for s in flatten(level.shape)) - stride_types = ", ".join(f"cute::Int<{s}>" for s in flatten(level.strides)) - layout = f"cute::Layout, cute::Stride<{stride_types}>>" + shape, strides = flattened_levels(mesh)[0] + layout = render_layout_type(shape, strides) else: if base: raise NotImplementedError( @@ -93,7 +83,7 @@ def mesh_type(mesh: Mesh) -> str: "be sliced; the slice and the level boundary would both be " "deciding which positions these are" ) - layout = _levelwise_layout(mesh, topos) + layout = _levelwise_layout(mesh) if base: layout = f"cute::ComposedLayout, {layout}>" return f"tilefoundry::Mesh<{layout}, {', '.join(topology_scope_str(t.name) for t in topos)}>" diff --git a/src/tilefoundry/codegen/cuda/tir/tuple_get_item.py b/src/tilefoundry/codegen/cuda/tir/tuple_get_item.py index 453753af8..287d1f427 100644 --- a/src/tilefoundry/codegen/cuda/tir/tuple_get_item.py +++ b/src/tilefoundry/codegen/cuda/tir/tuple_get_item.py @@ -2,10 +2,11 @@ from __future__ import annotations -from tilefoundry.codegen.cuda.context import CudaCodegenContext +from tilefoundry.codegen.cuda.context import CudaCodegenContext, scalar_read from tilefoundry.ir.core import Call, Constant, Tuple, Var from tilefoundry.ir.hir.tensor.tuple_get_item import TupleGetItem from tilefoundry.ir.types.dim import DimAdd, DimFloorDiv, DimMod, DimMul, DimSub +from tilefoundry.ir.types.storage import StorageKind from tilefoundry.target import CudaTarget from tilefoundry.visitor_registry.registries import Role, register_codegen @@ -31,7 +32,10 @@ def _index(expr, ctx: CudaCodegenContext) -> str: if isinstance(expr, Constant) and isinstance(expr.value, int): return str(expr.value) if isinstance(expr, Var): - return ctx.name_for(expr) + type_ = getattr(expr, "type", None) + if getattr(type_, "storage", None) in (None, StorageKind.UMAT): + return ctx.name_for(expr) + return scalar_read(expr, ctx) if isinstance(expr, Call): for op_type, symbol in _DIM_SYMBOLS.items(): if isinstance(expr.target, op_type): @@ -42,6 +46,26 @@ def _index(expr, ctx: CudaCodegenContext) -> str: ) +def constant_index(index: Constant) -> int: + """The element a constant TupleGetItem index selects.""" + selected = index.value + if not isinstance(selected, int) or isinstance(selected, bool): + raise ValueError("CUDA TupleGetItem constant index must be an integer") + return selected + + +def render_tuple_choice(rendered: tuple[str, ...], index, ctx: CudaCodegenContext) -> str: + """One runtime choice among *rendered* elements, subscripted by *index*. + + The choice is decided on the card: one ``cute::array`` of the elements, + subscripted by the index this function renders. ``cute::array`` has no + deduction guide, so the element type and count come from the first + element's spelling. + """ + array = f"cute::array{{{', '.join(rendered)}}}" + return f"{array}[{_index(index, ctx)}]" + + @register_codegen(CudaTarget, Role.EMIT, TupleGetItem) def _emit(let_stmt, ctx: CudaCodegenContext) -> None: call = let_stmt.value @@ -52,10 +76,7 @@ def _emit(let_stmt, ctx: CudaCodegenContext) -> None: raise ValueError("CUDA TupleGetItem requires a tuple value") result = ctx.name_for(let_stmt.var) if isinstance(index, Constant): - selected = index.value - if not isinstance(selected, int) or isinstance(selected, bool): - raise ValueError("CUDA TupleGetItem constant index must be an integer") - ctx.emit(f"auto {result} = {_value(held.elements[selected], ctx)};") + ctx.emit(f"auto {result} = {_value(held.elements[constant_index(index)], ctx)};") return - elements = ", ".join(_value(element, ctx) for element in held.elements) - ctx.emit(f"auto {result} = cute::array{{{elements}}}[{_index(index, ctx)}];") + rendered = tuple(_value(element, ctx) for element in held.elements) + ctx.emit(f"auto {result} = {render_tuple_choice(rendered, index, ctx)};") diff --git a/src/tilefoundry/ir/types/mesh.py b/src/tilefoundry/ir/types/mesh.py index 02bc4adca..897742d27 100644 --- a/src/tilefoundry/ir/types/mesh.py +++ b/src/tilefoundry/ir/types/mesh.py @@ -240,6 +240,22 @@ def levels(mesh: Mesh) -> tuple[Layout, ...]: return tuple(get(stated, index) for index in range(_rank(stated))) +def flattened_levels(mesh: Mesh) -> tuple[tuple[tuple, tuple], ...]: + """Each level's arrangement as flat (shape, strides) int-tuples. + + A level stating no axes reads as one mode of one, so a consumer + subscripting the levels still has something to pick. + """ + arrangements = [] + for level in levels(mesh): + shape = tuple(flatten(level.shape)) + strides = tuple(flatten(level.strides)) + if not shape: + shape, strides = (1,), (0,) + arrangements.append((shape, strides)) + return tuple(arrangements) + + def starts(mesh: Mesh) -> tuple[int, ...]: """Where each level's run starts, decoded from the device-numbered offset.""" offset = mesh.layout.offset if isinstance(mesh.layout, ComposedLayout) else 0 diff --git a/tests/codegen/test_staged_ring.py b/tests/codegen/test_staged_ring.py new file mode 100644 index 000000000..995cac855 --- /dev/null +++ b/tests/codegen/test_staged_ring.py @@ -0,0 +1,168 @@ +"""The staged-ring schedule fixtures, compiled and run on the card. + +``staged_dynamic_inner`` and ``staged_nested`` reach full CUDA emission +through the tuple-choice pointer plumbing: two smem stages held as one +tuple of tensor_views, subscripted by a runtime index. Their rings are +re-authored with what execution needs and the corpus does not state -- +the dynamic-smem carve-out, the cp.async commit/wait before read-back, +and trip counts uniform across the warp. + +See [runtime §2.6.2](docs/spec/runtime.md#262-opscopycuh). +""" + +from __future__ import annotations + +import pytest +import torch + +import tilefoundry +import tilefoundry.codegen.cuda # noqa: F401 -- trigger emitter autodiscovery +from tilefoundry.dsl import * +from tilefoundry.target import CpuTarget, CudaTarget + +_CUDA = CudaTarget("nvidia.h200_sxm") + + +@module( + entry="counter_ring_host", + target=_CUDA, + topologies=(Topology("cta", 1), Topology("thread", 32)), +) +class CounterRing: + """``staged_dynamic_inner``'s ring, driven by its rmem stage counter. + + The choice is as authored: two tensor_views over raw smem offsets held in + one tuple, subscripted by the counter modulo two, so the selected stage's + base pointer is decided on the card each iteration. + """ + + @prim_func(target=_CUDA) + def counter_ring_device(x: Tensor[(1, 8), "bf16"], out: Tensor[(1, 8), "bf16"]): + with Mesh((Topology("cta", 1),), layout=(1,), names=("d0",)) as cta: # noqa: F841 + with Mesh((Topology("thread", 32),), layout=(32,), names=("d0",)) as scope: + ring_stages = ( + T.tensor_view(0, dtype="bf16", storage="smem", layout=(1, 8), shape=(1, 8)), + T.tensor_view(16, dtype="bf16", storage="smem", layout=(1, 8), shape=(1, 8)), + ) + stage_counter = T.alloc_tensor(tensor_type=Tensor[(), "i64", "rmem"]) + T.fill(stage_counter, 0) + stage_one = T.alloc_tensor(tensor_type=Tensor[(), "i64", "rmem"]) + T.fill(stage_one, 1) + for i in range(0, 2, 1): + value = T.tensor_view(T.ptr_of(x), layout=(1, 8), shape=(1, 8)) + stage = T.tensor_view( + T.ptr_of(ring_stages[stage_counter % 2]), layout=(1, 8), shape=(1, 8) + ) + T.copy_async(value, stage) + T.cp_async_commit() + T.cp_async_wait(n=0) + T.sync(scope) + out_view = T.tensor_view(T.ptr_of(out), layout=(1, 8), shape=(1, 8)) + T.copy(stage, out_view) + T.binary(stage_counter, stage_one, stage_counter, kind=BinaryKind.ADD) + + @prim_func(target=CpuTarget()) + def counter_ring_host(x: Tensor[(1, 8), "bf16"], out: Tensor[(1, 8), "bf16"]): + launch( # noqa: F821 + counter_ring_device, # noqa: F821 + x, + out, + grid=(1, 1, 1), + block=(32, 1, 1), + dynamic_smem=32, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_the_counter_driven_ring_stages_the_tile_through_both_stages() -> None: + """The runtime index picks a stage per iteration; the tile arrives intact. + + A choice rendered wrong -- an index that does not read the counter, or + stage bases spelled so the subscript cannot compile -- either fails + emission here outright or lands the tile nowhere near ``out``. + """ + rm = tilefoundry.compile(CounterRing, target=_CUDA) + x = torch.arange(8, dtype=torch.bfloat16, device="cuda").view(1, 8) + out = torch.full((1, 8), -1.0, dtype=torch.bfloat16, device="cuda") + rm(x, out) + torch.cuda.synchronize() + assert torch.equal(out, x), ( + "counter-driven ring: iteration k copies x through stage k % 2, so " + "out must hold x exactly -- a stale or uninitialized stage read " + "breaks the equality" + ) + + +@module( + entry="row_ring_host", + target=_CUDA, + topologies=(Topology("cta", 1), Topology("thread", 32)), +) +class RowRing: + """``staged_nested``'s six rows through the same two-stage ring. + + Kept as authored: row ``3i + j`` stages through slot ``(3i + j) % 2``, + the row windows are slice-based views with a stride-0 first mode, and + the ring index is dim arithmetic over the two loop counters. + """ + + @prim_func(target=_CUDA) + def row_ring_device(x: Tensor[(6, 8), "bf16"], out: Tensor[(6, 8), "bf16"]): + with Mesh((Topology("cta", 1),), layout=(1,), names=("d0",)) as cta: # noqa: F841 + with Mesh((Topology("thread", 32),), layout=(32,), names=("d0",)) as scope: + ring_stages = ( + T.tensor_view(0, dtype="bf16", storage="smem", layout=(1, 8), shape=(1, 8)), + T.tensor_view(16, dtype="bf16", storage="smem", layout=(1, 8), shape=(1, 8)), + ) + for i in range(0, 2, 1): + for j in range(0, 3, 1): + value = T.tensor_view( + T.ptr_of(x[(3 * i) + j:(3 * i) + j + 1, 0:0 + 1]), + layout=((1, 8), (0, 1)), + shape=(1, 8), + ) + stage = T.tensor_view( + T.ptr_of(ring_stages[((i * 3) + j) % 2]), + layout=(1, 8), + shape=(1, 8), + ) + T.copy_async(value, stage) + T.cp_async_commit() + T.cp_async_wait(n=0) + T.sync(scope) + out_row = T.tensor_view( + T.ptr_of(out[(3 * i) + j:(3 * i) + j + 1, 0:0 + 1]), + layout=((1, 8), (0, 1)), + shape=(1, 8), + ) + T.copy(stage, out_row) + + @prim_func(target=CpuTarget()) + def row_ring_host(x: Tensor[(6, 8), "bf16"], out: Tensor[(6, 8), "bf16"]): + launch( # noqa: F821 + row_ring_device, # noqa: F821 + x, + out, + grid=(1, 1, 1), + block=(32, 1, 1), + dynamic_smem=32, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_every_row_stages_through_the_slot_its_index_arithmetic_picks() -> None: + """Dim arithmetic over both loop counters drives the ring; rows arrive. + + The index is ``((i * 3) + j) % 2`` rendered for the card, over windows + whose bases come out of the slice path, so a row landed anywhere but its + slot's stage breaks the row-for-row equality. + """ + rm = tilefoundry.compile(RowRing, target=_CUDA) + x = torch.arange(48, dtype=torch.bfloat16, device="cuda").view(6, 8) + out = torch.full((6, 8), -1.0, dtype=torch.bfloat16, device="cuda") + rm(x, out) + torch.cuda.synchronize() + assert torch.equal(out, x), ( + "row ring: row 3i + j passes through stage (3i + j) % 2 into out's " + "row 3i + j, so out must equal x row for row" + )