diff --git a/docs/spec/runtime.md b/docs/spec/runtime.md index c46081e35..4fcbc03c0 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); ``` @@ -968,6 +972,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/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/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/copy/copy_impl.h b/include/tilefoundry/runtime/cuda/ops/copy/copy_impl.h index dcf4e5761..6e91402fb 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/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/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 655f436b6..bdcaadefb 100644 --- a/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py +++ b/src/tilefoundry/codegen/cuda/tir/memory/tensor_view.py @@ -12,17 +12,21 @@ 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 @@ -30,6 +34,7 @@ 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.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__}") @@ -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,17 +298,21 @@ def _mesh_dim(d): return "tilefoundry::program_dim()" return _static_dim(d, "mesh layout dim") + 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" - 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) - positions = ( - f"cute::make_layout(cute::make_shape({ml_shape_args}), cute::make_stride({ml_stride_args}))" - ) + if len(topos) == 1: + mesh_shape, mesh_strides = mesh_axes.shape, mesh_axes.strides + else: + 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 = 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) @@ -340,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" ) @@ -373,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): @@ -388,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) @@ -399,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 _static_value(value) -> str: + """A static stride entry as a CuTe value.""" + return f"cute::Int<{int(value)}>{{}}" - def stride(value): - return f"cute::Int<{int(value)}>{{}}" - return _render_layout_value(layout, dim, stride) +def _plain_layout_value(layout: LayoutBase) -> str: + return _render_layout_value(layout, _envelope_dim, _static_value) @register_codegen(CudaTarget, Role.EMIT, TensorView) @@ -450,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/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/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" + ) 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" + ) 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 be8bb20d3..6bcf842d7 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 db4622e83..408a2a5b1 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",