Skip to content

perf(runtime): vectorize the store side of copy_fragment - #69

Open
bigSheep123 wants to merge 1 commit into
tile-ai:mainfrom
bigSheep123:fix/vectorized-store
Open

perf(runtime): vectorize the store side of copy_fragment#69
bigSheep123 wants to merge 1 commit into
tile-ai:mainfrom
bigSheep123:fix/vectorized-store

Conversation

@bigSheep123

Copy link
Copy Markdown
Contributor

perf(runtime): vectorize the store side of copy_fragment

Fixes #68.

What was wrong

copy_fragment gated its 128-bit path on SrcIsGmem, so a gmem -> rmem load
vectorized but the rmem -> gmem store never did. Every reshard back to global
memory emitted one narrow access per element, even when the register-side run
was a statically contiguous, aligned block — so a kernel whose loads were
already wide still stored scalar.

The change

Add the mirror case — destination global, register-side source a statically
contiguous run — packing into uint4 before the store, behind the same
alignment and unit-stride guards the load path uses. copy() now passes both
src_gmem and dst_gmem into copy_fragment.

The two shared predicates are factored out rather than duplicated:
StaticContigView (coalesced, static, unit-stride) and is_contiguous_16b
(the runtime 16B alignment and contiguity check). Both directions now use one
copy of each, so they cannot drift apart.

Evidence

Generated PTX for a store-heavy reshard at sm_90 — deterministic, and the
clearest evidence that the path is actually taken:

before:   4 st.global.f32
after:   64 st.global.v4.f32  +  4 st.global.f32

Wall clock on a bandwidth-bound gmem -> rmem -> gmem square, 1024 CTAs x 256
threads x 8 f32 per thread. Best of five; the run-to-run spread on a shared card
is wide, so the modal value is quoted rather than a single reading:

time achieved bandwidth
before 12.4 us 1350 GB/s
after 6.8 us 2480 GB/s

About 1.8x on this shape. The narrower the per-thread run the smaller the
effect, and a kernel that stores nothing is unaffected.

Tests

tests/integration/test_vectorized_store.py covers the tail. The vector path
copies whole uint4 groups and then the remainder one element at a time, so a
run length that is not a multiple of four f32 is what an off-by-one there
would corrupt — the four cases are one exact group (4), and tails of two (6),
one (13) and three (259) elements.

These pass before this change as well. They are regression guards for the
newly reachable path, not a reproduction of a defect: the old behaviour was slow
rather than wrong. The defect the PTX above shows is a missing optimization, and
the risk this PR introduces is exactly the tail handling, which is what the
tests pin.

Each size is a separate top-level kernel on purpose — the compiled artifact is
cached per function, so a factory emitting several same-named kernels hands
every size whichever one compiled first.

Verification

pytest tests/ -q     644 passed (main @ e55e4b0)  ->  648 passed (+4 new)

Every .pre-commit-config.yaml hook clean: ruff, spec-rules, spec-refs,
spec-entropy, forward-references, comment-hygiene, no-machine-paths,
english-only, and clang-format over all C++ files.

`copy_fragment` gated its 128-bit path on `SrcIsGmem`, so a gmem -> rmem load
vectorized but the rmem -> gmem store never did: every reshard back to global
emitted one narrow access per element. A kernel whose loads were already wide
still stored scalar.

Add the mirror case -- destination global, register-side source a statically
contiguous run -- packing into `uint4` before the store, behind the same
alignment and unit-stride guards the load path uses. `copy()` now passes both
`src_gmem` and `dst_gmem`. The "coalesced, static, unit-stride" predicate and
the 16B-contiguity runtime check are factored into `StaticContigView` and
`is_contiguous_16b` so the two directions cannot drift apart.

Effect on the generated PTX for a store-heavy reshard, which is deterministic
and the clearest evidence here:

    before:  4 st.global.f32
    after:  64 st.global.v4.f32  +  4 st.global.f32

On a bandwidth-bound gmem -> rmem -> gmem square over 1024 CTAs x 256 threads
x 8 f32 (best of five, the run-to-run spread on a shared card is wide):

    before:  12.4 us   1350 GB/s
    after:    6.8 us   2480 GB/s

The added tests cover the tail: the vector path copies whole `uint4` groups and
then the remainder one element at a time, so a run length that is not a
multiple of four is what an off-by-one there would corrupt. They pass before
this change as well -- they guard the new path rather than reproduce a defect.

pytest tests/ -q: 644 passed before, 648 after (+4 new, no regressions).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant