perf(runtime): vectorize the store side of copy_fragment

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

@@ -4,17 +4,73 @@
 
 namespace detail {
 
-template <bool SrcIsGmem, class SView, class DView>
+// The gmem-side precondition both vector paths need: the view's first element
+// is 16B aligned and the next N - 1 elements follow it contiguously. Only the
+// register side is statically known to be a unit-stride run, so this half is a
+// runtime check. `N` is a template parameter so the unroll stays compile-time.
+template <int N, class View>
+CUTE_HOST_DEVICE bool is_contiguous_16b(View const &v) {
+    auto const *p = &v(0);
+    bool ok = (reinterpret_cast<uintptr_t>(p) & 0xF) == 0;
+    CUTE_UNROLL
+    for (int i = 1; i < N; ++i)
+        ok = ok && (&v(i) == p + i);
+    return ok;
+}
+
+// A coalesced view is a single unit-stride run of statically known length —
+// the shape a 16B vector access needs on the register side.
+template <class View> struct StaticContigView {
+    using layout_t =
+        decltype(cute::coalesce(cute::remove_cvref_t<View>{}.layout()));
+    static constexpr bool value =
+        cute::is_static<layout_t>::value &&
+        decltype(cute::rank(layout_t{}))::value == 1 &&
+        (int(cute::size(layout_t{})) == 1 ||
+         int(cute::stride<0>(layout_t{})) == 1);
+};
+
+template <bool SrcIsGmem, bool DstIsGmem, class SView, class DView>
 CUTE_HOST_DEVICE void copy_fragment(SView const &sv, DView &dv) {
     using s_val_t = cute::remove_cvref_t<decltype(sv(0))>;
     using d_val_t = cute::remove_cvref_t<decltype(dv(0))>;
-    using dvc_layout_t =
-        decltype(cute::coalesce(cute::remove_cvref_t<DView>{}.layout()));
-    constexpr bool dst_static_contig =
-        cute::is_static<dvc_layout_t>::value &&
-        decltype(cute::rank(dvc_layout_t{}))::value == 1 &&
-        (int(cute::size(dvc_layout_t{})) == 1 ||
-         int(cute::stride<0>(dvc_layout_t{})) == 1);
+    using dvc_layout_t = typename StaticContigView<DView>::layout_t;
+    using svc_layout_t = typename StaticContigView<SView>::layout_t;
+    constexpr bool dst_static_contig = StaticContigView<DView>::value;
+    constexpr bool src_static_contig = StaticContigView<SView>::value;
+
+    // Store side: the destination is global and the register-side source is a
+    // contiguous run. Without this a reshard back to gmem writes one narrow
+    // access per element, so a kernel whose loads vectorize still stores
+    // scalar.
+    if constexpr (DstIsGmem && !SrcIsGmem && std::is_same_v<s_val_t, d_val_t> &&
+                  src_static_contig) {
+        constexpr int mcv = int(decltype(cute::max_common_vector(
+            svc_layout_t{}, svc_layout_t{}))::value);
+        constexpr int vec_bits = mcv * int(cute::sizeof_bits<d_val_t>::value);
+        if constexpr (vec_bits >= 128) {
+            constexpr int N = int(cute::size(svc_layout_t{}));
+            constexpr int C = 16 / int(sizeof(d_val_t));
+            constexpr int NV = N / C;
+            d_val_t *dp = &dv(0);
+            if (is_contiguous_16b<N>(dv)) {
+                uint4 tmp[NV];
+                d_val_t *tp = reinterpret_cast<d_val_t *>(tmp);
+                CUTE_UNROLL
+                for (int i = 0; i < NV * C; ++i)
+                    tp[i] = sv(i);
+                uint4 *vp = reinterpret_cast<uint4 *>(dp);
+                CUTE_UNROLL
+                for (int i = 0; i < NV; ++i)
+                    vp[i] = tmp[i];
+                CUTE_UNROLL
+                for (int i = NV * C; i < N; ++i)
+                    dv(i) = sv(i);
+                return;
+            }
+        }
+    }
+
     if constexpr (SrcIsGmem && std::is_same_v<s_val_t, d_val_t> &&
                   dst_static_contig) {
         constexpr int mcv = int(decltype(cute::max_common_vector(
@@ -25,11 +81,7 @@ CUTE_HOST_DEVICE void copy_fragment(SView const &sv, DView &dv) {
             constexpr int C = 16 / int(sizeof(d_val_t));
             constexpr int NV = N / C;
             s_val_t const *sp = &sv(0);
-            bool ok = (reinterpret_cast<uintptr_t>(sp) & 0xF) == 0;
-            CUTE_UNROLL
-            for (int i = 1; i < N; ++i)
-                ok = ok && (&sv(i) == sp + i);
-            if (ok) {
+            if (is_contiguous_16b<N>(sv)) {
                 uint4 const *vp = reinterpret_cast<uint4 const *>(sp);
                 uint4 tmp[NV];
                 CUTE_UNROLL
@@ -78,5 +130,7 @@ CUTE_HOST_DEVICE void copy(ShardTensor<TS, GLS, SLS> const &src,
     auto &&dv = local(dst);
     constexpr bool src_gmem =
         cute::is_gmem<cute::remove_cvref_t<decltype(src.engine)>>::value;
-    detail::copy_fragment<src_gmem>(sv, dv);
+    constexpr bool dst_gmem =
+        cute::is_gmem<cute::remove_cvref_t<decltype(dst.engine)>>::value;
+    detail::copy_fragment<src_gmem, dst_gmem>(sv, dv);
 }

tests/integration/test_vectorized_store.py

@@ -0,0 +1,90 @@
+"""A reshard back to global memory keeps its numerics when the store vectorizes.
+
+``copy_fragment`` takes a 128-bit path when the register-side source is a
+statically contiguous run and the destination is global. That path packs whole
+``uint4`` groups and copies the remainder one element at a time, so a run length
+that is *not* a multiple of four f32 is the interesting case: an off-by-one in
+the tail loop drops or duplicates the last elements of every thread's run, which
+no aligned length can reveal.
+
+Each kernel below is a distinct top-level function on purpose. The compiled
+artifact is cached per function, so a factory producing several same-named
+kernels would hand every size the first one that was compiled.
+"""
+
+from __future__ import annotations
+
+import torch
+
+import tilefoundry
+from tilefoundry import func
+from tilefoundry.dsl import Tensor, tf
+from tilefoundry.dsl.storage import gmem, rmem
+from tilefoundry.ir.types.shard import Layout, Mesh, Topology
+
+_ROWS = 64
+
+
+# 4 f32 == exactly one 128-bit group, no tail.
+@func(topologies=(Topology("cta", _ROWS),))
+def square_4(a: Tensor[(_ROWS, 4), "f32"]) -> Tensor[(_ROWS, 4), "f32"]:
+    with Mesh(topology="cta", layout=Layout(shape=(_ROWS,), strides=(1,))) as cta:
+        reg = tf.reshard(a, layout=(_ROWS @ cta, 4), storage=rmem)
+        return tf.reshard(tf.mul(reg, reg), layout=(_ROWS @ cta, 4), storage=gmem)
+
+
+# One full group plus a 2-element tail.
+@func(topologies=(Topology("cta", _ROWS),))
+def square_6(a: Tensor[(_ROWS, 6), "f32"]) -> Tensor[(_ROWS, 6), "f32"]:
+    with Mesh(topology="cta", layout=Layout(shape=(_ROWS,), strides=(1,))) as cta:
+        reg = tf.reshard(a, layout=(_ROWS @ cta, 6), storage=rmem)
+        return tf.reshard(tf.mul(reg, reg), layout=(_ROWS @ cta, 6), storage=gmem)
+
+
+# Three full groups plus a 1-element tail.
+@func(topologies=(Topology("cta", _ROWS),))
+def square_13(a: Tensor[(_ROWS, 13), "f32"]) -> Tensor[(_ROWS, 13), "f32"]:
+    with Mesh(topology="cta", layout=Layout(shape=(_ROWS,), strides=(1,))) as cta:
+        reg = tf.reshard(a, layout=(_ROWS @ cta, 13), storage=rmem)
+        return tf.reshard(tf.mul(reg, reg), layout=(_ROWS @ cta, 13), storage=gmem)
+
+
+# A long run — 64 full groups plus a 3-element tail.
+@func(topologies=(Topology("cta", _ROWS),))
+def square_259(a: Tensor[(_ROWS, 259), "f32"]) -> Tensor[(_ROWS, 259), "f32"]:
+    with Mesh(topology="cta", layout=Layout(shape=(_ROWS,), strides=(1,))) as cta:
+        reg = tf.reshard(a, layout=(_ROWS @ cta, 259), storage=rmem)
+        return tf.reshard(tf.mul(reg, reg), layout=(_ROWS @ cta, 259), storage=gmem)
+
+
+def _check(kernel, cols: int) -> None:
+    rm = tilefoundry.compile(kernel, target="cuda")
+    torch.manual_seed(0)
+    # Away from zero so a dropped tail element cannot coincidentally match, and
+    # the output starts as NaN so an unwritten element never looks right.
+    x = torch.randn(_ROWS, cols, dtype=torch.float32, device="cuda") + 2.0
+    out = torch.full_like(x, float("nan"))
+    rm(x, out)
+    torch.cuda.synchronize()
+
+    expected = x * x
+    assert torch.allclose(out, expected, rtol=0, atol=0), (
+        f"cols={cols}: {int((out != expected).sum())} of {out.numel()} wrong "
+        f"(tail of {cols % 4} elements past the last full 128-bit group)"
+    )
+
+
+def test_a_store_of_exactly_one_vector_group_is_exact() -> None:
+    _check(square_4, 4)
+
+
+def test_a_store_with_a_two_element_tail_is_exact() -> None:
+    _check(square_6, 6)
+
+
+def test_a_store_with_a_one_element_tail_is_exact() -> None:
+    _check(square_13, 13)
+
+
+def test_a_long_store_with_a_three_element_tail_is_exact() -> None:
+    _check(square_259, 259)