CUDA/HIP: VALU cross-lane ops for partial-top-k shuffles (~1.06-1.10x, follow-up to #15)#16
Merged
Merged
Conversation
…~1.06-1.10x)
Follow-up tuning of the partial-bitonic top-k, correctness-preserving (comparator/
merge networks byte-identical -> bit-exact TOP_K set), validated with interleaved
A/B rebuilds on gfx1201 and gfx1151.
The intra-wave xor shuffles were issuing through the LDS permute crossbar
(ds_swizzle_b32 / ds_bpermute_b32). Move the ones expressible as pure VALU
cross-lane gathers off the crossbar onto the ALU pipe, shortening the dependent
shuffle chains these latency-bound kernels sit on:
* xor ^1 / ^2 (intra-quad) -> DPP mov_dpp quad_perm (0xB1 / 0x4E).
* xor ^16 (widest intra-wave stride, entry stage of every 32-lane chain) ->
v_permlanex16_b32 in its identity-cross form.
* ^4 / ^8 stay on ds_swizzle_b32; the wave64 j==32 tail falls back to __shfl_xor.
* plus a middle-cross-wave-stride group-merge fuse in the shared/2wave Phase-A/B.
All cross-lane encodings verified bit-identical to __shfl_xor on gfx1201 and
gfx1151. test-backend-ops -o TOP_K: 0 fail on both archs. Perf (ne=[1000,16],
geomean over k=1/10/40/400 vs the prior partial-bitonic): gfx1201 ~1.06x,
gfx1151 ~1.10x (k=1 and k=400 gain most: gfx1151 1.13x / 1.15x).
Collaborator
Author
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
1 participant
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.
Draft — follow-up to #15, stacked on its branch (
feat/hip-topk-argsort), so the diff below is only this change. Found by our kernel optimizer on top of #15's partial-bitonic; kept separate so #15 stays a clean, self-contained win.What
The partial-bitonic top-k's intra-wave xor shuffles were all issuing through the LDS permute crossbar (
ds_swizzle_b32/ds_bpermute_b32). Move the ones expressible as pure VALU cross-lane gathers off the crossbar onto the ALU pipe, shortening the dependent shuffle chains these latency-bound kernels sit on:^1/^2(intra-quad) → DPPmov_dppquad_perm (0xB1/0x4E)^16(widest intra-wave stride, entry stage of every 32-lane chain) →v_permlanex16_b32(identity-cross form)^4/^8stay onds_swizzle_b32; the wave64j==32tail falls back to__shfl_xorComparator/merge networks are byte-identical → the produced top-k set is bit-identical.
Impact (
test-backend-ops perf -o TOP_K, ne=[1000,16], vs #15's partial-bitonic)Geomean ~1.06× (gfx1201) / ~1.10× (gfx1151) on top of #15.
Validation
quad_perm/permlanex16/ds_swizzleencodings verified bit-identical to__shfl_xoron gfx1201 and gfx1151 (on-device probe + the test gate).test-backend-ops test -o TOP_K:2/2 backends passed, 0 fail, on both archs.libggml-hip.soswap) to cancel warm-drift.Notes
Draft because it's incremental polish on already-diminishing returns (the optimizer plateaued for 15 iterations after this) — merge #15 first; this can follow. Same-arch dependency: nothing new (pure ISA-level shuffle lowering).