fix(codegen,runtime): per-level reduce planning and lowered-TIR memory views

docs/spec/runtime.md

@@ -950,6 +950,10 @@ CUTE_HOST_DEVICE auto program_ids() noexcept;
 // include/tilefoundry/runtime/cuda/layout/cute_ext.cuh
 template <class Shape, class Stride>
 CUTE_HOST_DEVICE constexpr auto reverse(cute::Layout<Shape, Stride> const &l);
+
+template <class A, class O, class B>
+CUTE_HOST_DEVICE constexpr auto
+reverse(cute::ComposedLayout<A, O, B> const &l);
 ```
 <!-- /generated -->
 
@@ -968,6 +972,12 @@ template <class TLayout, TopologyScope... Topos> 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 <int I>
+    CUTE_HOST_DEVICE static constexpr TopologyScope axis_scope();
 };
 
 template <TopologyScope S, class L, TopologyScope... Topos>

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

@@ -15,3 +15,12 @@ CUTE_HOST_DEVICE constexpr auto reverse(cute::Layout<Shape, Stride> 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 <class A, class O, class B>
+CUTE_HOST_DEVICE constexpr auto
+reverse(cute::ComposedLayout<A, O, B> const &l) {
+    return cute::make_composed_layout(l.layout_a(), l.offset(),
+                                      reverse(cute::flatten(l.layout_b())));
+}

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

@@ -24,6 +24,32 @@ template <class TLayout, TopologyScope... Topos> struct Mesh {
     static constexpr int level_count = sizeof...(Topos);
     static constexpr TopologyScope scope = detail::first_scope<Topos...>::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 <int I>
+    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;
+            [&]<size_t... Ls>(std::index_sequence<Ls...>) {
+                (([&] {
+                     constexpr int width =
+                         int(cute::rank(cute::get<Ls>(cute::shape(TLayout{}))));
+                     if (I >= first && I < first + width)
+                         found = scopes[Ls];
+                     first += width;
+                 }()),
+                 ...);
+            }(std::make_index_sequence<sizeof...(Topos)>{});
+            return found;
+        }
+    }
 };
 
 namespace detail {

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 <class T, class L, class A, class M>
+CUTE_HOST_DEVICE auto local_as(T &&t,
+                               tilefoundry::ShardLayout<L, A, M> const &sl) {
+    auto logical =
+        cute::make_layout(cute::shape(sl.layout_value), cute::GenRowMajor{});
+    tilefoundry::ShardLayout<decltype(logical), A, M> 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 <class TSrc, class TDst>
     __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<TSrc> &&
+                      !tilefoundry::ShardTensorLike<TDst>) {
+            auto &&s = tilefoundry::local_tensor(src);
+            auto d = local_as(dst, src.shard_layout);
+            cute::copy(s, d);
+        } else if constexpr (!tilefoundry::ShardTensorLike<TSrc> &&
+                             tilefoundry::ShardTensorLike<TDst>) {
+            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);
+        }
     }
 };
 

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 <class Op, class WorkspaceT>
 __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<Op>(acc, static_cast<float>(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<mean_op>`` 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 <class TMesh> 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<TMesh>;
+    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<mean_op>`` 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 <class SrcSL, class DstSL>
 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<typename SrcSL::mesh>;
     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<src_attrs>::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<m_rank - 1>() == TopologyScope::thread;
     static_assert(detail::shard_attrs_match_mesh<SrcSL>() &&
                       detail::shard_attrs_match_mesh<DstSL>(),
                   "ops::reduce: both operands need one attr per mesh axis");
     static_assert(std::is_same_v<typename SrcSL::mesh, typename DstSL::mesh>,
                   "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{}));
     [&]<size_t... Is>(std::index_sequence<Is...>) {
         ((m_ext[Is] = int(cute::get<Is>(m_shape))), ...);
         ((reduced[Is] =
@@ -252,32 +263,44 @@ CUTE_HOST_DEVICE constexpr reduce_dispatch_info reduce_dispatch() {
                   cute::remove_cvref_t<decltype(cute::get<Is>(dst_attrs{}))>,
                   shard::B>),
          ...);
+        ((thread_axis[Is] =
+              mesh_t::template axis_scope<int(Is)>() == TopologyScope::thread),
+         ...);
     }(std::make_index_sequence<m_rank>{});
 
-    /// 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<cute::remove_cvref_t<
+                decltype(tilefoundry::get<TopologyScope::thread>(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<mesh_t>() / 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<mesh_t>() > 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) {

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<Op>`` 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<Op>``
+/// finalises. MEAN divides by ``step × lanes_reduced × warps_per_group``.
 template <class Op, class Axes> struct IntraCta {
     template <class SrcT, class DstT, class WorkspaceT>
     __device__ void operator()(SrcT const &src, DstT &dst,
@@ -30,21 +28,20 @@ template <class Op, class Axes> struct IntraCta {
         auto &&ws = tilefoundry::local_tensor(workspace);
         using value_type = cute::remove_cvref_t<decltype(d(0))>;
 
-        static_assert(kept_cells<Axes, decltype(s)>() == 1,
-                      "ops::reduce (intra-CTA tier): kept axes must form "
-                      "one cell");
+        constexpr int kCells = kept_cells<Axes, decltype(s)>();
         constexpr int kSpan = reduced_span<Axes, decltype(s)>();
-        const float local = local_fold<Op, Axes>(s, 0);
-        const float warp_partial =
-            tilefoundry::warp_reduce<typename reduce_traits<Op>::combine_op,
-                                     plan.lanes_reduced>(local);
-        const float cta_partial =
-            cta_combine_via_workspace<Op>(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<Op>::finalize(cta_partial, total_n);
-        for (int i = 0; i < static_cast<int>(cute::size(d)); ++i) {
-            d(i) = static_cast<value_type>(result);
+        CUTE_UNROLL
+        for (int j = 0; j < kCells; ++j) {
+            const float local = local_fold<Op, Axes>(s, j);
+            const float warp_partial =
+                tilefoundry::warp_reduce<typename reduce_traits<Op>::combine_op,
+                                         plan.lanes_reduced>(local);
+            const float cta_partial = cta_combine_via_workspace<Op>(
+                warp_partial, ws, warps_per_group);
+            const float total_n = float(kSpan) * float(plan.lanes_reduced) *
+                                  float(warps_per_group);
+            d(j) = static_cast<value_type>(
+                reduce_traits<Op>::finalize(cta_partial, total_n));
         }
     }
 };

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."""
 

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
 

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});"