diff --git a/include/tilefoundry/runtime/cuda/tensor_view/shard_copy.cuh b/include/tilefoundry/runtime/cuda/tensor_view/shard_copy.cuh index 5fed7377..0104f076 100644 --- a/include/tilefoundry/runtime/cuda/tensor_view/shard_copy.cuh +++ b/include/tilefoundry/runtime/cuda/tensor_view/shard_copy.cuh @@ -4,17 +4,73 @@ namespace detail { -template +// 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 +CUTE_HOST_DEVICE bool is_contiguous_16b(View const &v) { + auto const *p = &v(0); + bool ok = (reinterpret_cast(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 struct StaticContigView { + using layout_t = + decltype(cute::coalesce(cute::remove_cvref_t{}.layout())); + static constexpr bool value = + cute::is_static::value && + decltype(cute::rank(layout_t{}))::value == 1 && + (int(cute::size(layout_t{})) == 1 || + int(cute::stride<0>(layout_t{})) == 1); +}; + +template CUTE_HOST_DEVICE void copy_fragment(SView const &sv, DView &dv) { using s_val_t = cute::remove_cvref_t; using d_val_t = cute::remove_cvref_t; - using dvc_layout_t = - decltype(cute::coalesce(cute::remove_cvref_t{}.layout())); - constexpr bool dst_static_contig = - cute::is_static::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::layout_t; + using svc_layout_t = typename StaticContigView::layout_t; + constexpr bool dst_static_contig = StaticContigView::value; + constexpr bool src_static_contig = StaticContigView::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 && + 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::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(dv)) { + uint4 tmp[NV]; + d_val_t *tp = reinterpret_cast(tmp); + CUTE_UNROLL + for (int i = 0; i < NV * C; ++i) + tp[i] = sv(i); + uint4 *vp = reinterpret_cast(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 && 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(sp) & 0xF) == 0; - CUTE_UNROLL - for (int i = 1; i < N; ++i) - ok = ok && (&sv(i) == sp + i); - if (ok) { + if (is_contiguous_16b(sv)) { uint4 const *vp = reinterpret_cast(sp); uint4 tmp[NV]; CUTE_UNROLL @@ -78,5 +130,7 @@ CUTE_HOST_DEVICE void copy(ShardTensor const &src, auto &&dv = local(dst); constexpr bool src_gmem = cute::is_gmem>::value; - detail::copy_fragment(sv, dv); + constexpr bool dst_gmem = + cute::is_gmem>::value; + detail::copy_fragment(sv, dv); } diff --git a/tests/integration/test_vectorized_store.py b/tests/integration/test_vectorized_store.py new file mode 100644 index 00000000..09d31900 --- /dev/null +++ b/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)