Add stochastic rounding and philox RNG under rocdl - #917
Closed
kashif wants to merge 1 commit into
Closed
Conversation
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
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.
Stacked on #916. Adds stochastic rounding and a Philox counter RNG, placed under the rocdl target package rather than as top-level helpers, following the discussion on #907.
Stochastic rounding does not fit the target-neutral IEEE rounding modes that #916 adds to
.to()(lowered througharith.truncf): it needs entropy and the hardware form is a target-specific instruction (CDNAv_cvt_*_sr, PTXcvt.rs). So it lives underrocdl/as a target-specific function. The f32 to bf16 path here is a portable software emulation; a CDNA hardwaresrvariant can slot in alongside it later.New API:
fx.rocdl.philox_4x32(counter, key)is the Random123 Philox 4x32 counter RNG (same generator asnumpy.random.Philox, cuRANDPhilox4_32_10, Tritontl.rand), returning four reproducible random words from a global index and a seed.fx.rocdl.stochastic_round_bf16(x, rand)perturbs the 16 mantissa bits bf16 discards with the random bits before truncating, so the probability of rounding up equals the fractional distance to the next bf16 value. Inf and NaN pass through unchanged.Consumer:
The RDNA WMMA GEMM (gfx11 and gfx12) gains a
roundingoption, default"rn"so nothing changes. Withrounding="rs"the f32 accumulator converts to bf16 through stochastic rounding, keyed per element from a runtimesr_seedlaunch argument.Tests:
The rounding is proven exactly unbiased in pure Python and covers Inf and NaN. Device tests confirm outputs land on the two nearest bf16 values with the right up probability and a matching mean, and that the RNG is uniform and collision free. The GEMM test checks the stochastic path stays within tolerance of round-to-nearest, varies with the runtime seed, and is reproducible for a fixed seed.
Note for review: this targets the #916 branch so the diff shows only the stochastic-rounding work. It should be merged after #916, and GitHub will retarget it to main automatically.