Skip to content

Vectorized copy is load-only: every reshard back to global memory stores scalar #68

Description

@bigSheep123

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions