perf(runtime): vectorize the store side of copy_fragment - #69
Open
bigSheep123 wants to merge 1 commit into
Open
Conversation
`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).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
perf(runtime): vectorize the store side of copy_fragment
Fixes #68.
What was wrong
copy_fragmentgated its 128-bit path onSrcIsGmem, so agmem -> rmemloadvectorized but the
rmem -> gmemstore never did. Every reshard back to globalmemory 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
uint4before the store, behind the samealignment and unit-stride guards the load path uses.
copy()now passes bothsrc_gmemanddst_gmemintocopy_fragment.The two shared predicates are factored out rather than duplicated:
StaticContigView(coalesced, static, unit-stride) andis_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 theclearest evidence that the path is actually taken:
Wall clock on a bandwidth-bound
gmem -> rmem -> gmemsquare, 1024 CTAs x 256threads 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:
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.pycovers the tail. The vector pathcopies whole
uint4groups and then the remainder one element at a time, so arun 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
Every
.pre-commit-config.yamlhook clean: ruff, spec-rules, spec-refs,spec-entropy, forward-references, comment-hygiene, no-machine-paths,
english-only, and clang-format over all C++ files.