Vectorized copy is load-only: every reshard back to global memory stores scalar
Affects: main @ e55e4b0
Summary
copy_fragment (include/tilefoundry/runtime/cuda/tensor_view/shard_copy.cuh)
gates its 128-bit path on the source being global:
if constexpr (SrcIsGmem && std::is_same_v<s_val_t, d_val_t> && dst_static_contig)
So a gmem -> rmem load vectorizes, but an rmem -> gmem store never can — the
mirror condition is not implemented. Every reshard back to global memory emits
one narrow access per element, even when the register-side run is a statically
contiguous, correctly aligned block. A kernel whose loads are already wide still
stores scalar.
Evidence
Generated PTX for a store-heavy gmem -> rmem -> gmem reshard, compiled for
sm_90:
st.global.f32 4 <-- scalar stores only
There is no st.global.v4.* anywhere in the kernel, while the load side of the
same kernel does use wide accesses.
Cost
A bandwidth-bound gmem -> rmem -> gmem square, 1024 CTAs x 256 threads x 8 f32
per thread. Best of five runs; the run-to-run spread on a shared card is wide,
so the modal value is quoted.
|
time |
achieved bandwidth |
| current |
12.4 us |
1350 GB/s |
| with the store path implemented |
6.8 us |
2480 GB/s |
About 1.8x on this shape. The narrower the per-thread run, the smaller the
effect; a kernel that stores nothing is unaffected.
Suggested fix
Add the mirror case — destination global, register-side source statically
contiguous — packing into uint4 before the store, behind the same alignment
and unit-stride guards the load path already uses, and pass dst_gmem into
copy_fragment alongside src_gmem.
Worth factoring the shared predicates out at the same time: the
"coalesced, static, unit-stride" test and the 16B-contiguity runtime check are
identical for both directions, and duplicating them invites the two paths to
drift.
Note for whoever implements it: the vector path copies whole uint4 groups and
then the remainder one element at a time, so the run lengths that are not a
multiple of four are where an off-by-one would show up. Aligned-length cases
cannot reveal it.
Related, not covered here
tilefoundry::ops::copy_n (ops/copy/copy_impl.h) is still a plain scalar loop
with no vector path at all. Reshard does not route through it — it goes
tilefoundry::copy -> copy_fragment — so this issue leaves it alone, but any
op that does use copy_n gets narrow accesses.
PR follows.
Vectorized copy is load-only: every reshard back to global memory stores scalar
Affects:
main@e55e4b0Summary
copy_fragment(include/tilefoundry/runtime/cuda/tensor_view/shard_copy.cuh)gates its 128-bit path on the source being global:
So a
gmem -> rmemload vectorizes, but anrmem -> gmemstore never can — themirror condition is not implemented. Every reshard back to global memory emits
one narrow access per element, even when the register-side run is a statically
contiguous, correctly aligned block. A kernel whose loads are already wide still
stores scalar.
Evidence
Generated PTX for a store-heavy
gmem -> rmem -> gmemreshard, compiled forsm_90:There is no
st.global.v4.*anywhere in the kernel, while the load side of thesame kernel does use wide accesses.
Cost
A bandwidth-bound
gmem -> rmem -> gmemsquare, 1024 CTAs x 256 threads x 8 f32per thread. Best of five runs; the run-to-run spread on a shared card is wide,
so the modal value is quoted.
About 1.8x on this shape. The narrower the per-thread run, the smaller the
effect; a kernel that stores nothing is unaffected.
Suggested fix
Add the mirror case — destination global, register-side source statically
contiguous — packing into
uint4before the store, behind the same alignmentand unit-stride guards the load path already uses, and pass
dst_gmemintocopy_fragmentalongsidesrc_gmem.Worth factoring the shared predicates out at the same time: the
"coalesced, static, unit-stride" test and the 16B-contiguity runtime check are
identical for both directions, and duplicating them invites the two paths to
drift.
Note for whoever implements it: the vector path copies whole
uint4groups andthen the remainder one element at a time, so the run lengths that are not a
multiple of four are where an off-by-one would show up. Aligned-length cases
cannot reveal it.
Related, not covered here
tilefoundry::ops::copy_n(ops/copy/copy_impl.h) is still a plain scalar loopwith no vector path at all. Reshard does not route through it — it goes
tilefoundry::copy->copy_fragment— so this issue leaves it alone, but anyop that does use
copy_ngets narrow accesses.PR follows.