include/brotensor/detail/hash_rng.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | |||
| 3 | // brotensor detail/hash_rng.h — counter-based hash RNG shared by the CPU and | ||
| 4 | // CUDA backends (and by tests that need to reproduce an op's noise). | ||
| 5 | // | ||
| 6 | // A stateless, order-independent uniform generator: the uniform for element | ||
| 7 | // `index` of stream (`seed`, `domain`) is a pure function of those three | ||
| 8 | // integers, so a CPU loop and a CUDA grid produce bit-identical noise for | ||
| 9 | // bit-identical inputs. The mixer is splitmix64 (Steele/Lea/Flood) — a | ||
| 10 | // bijective 64-bit finaliser with full avalanche, which is exactly what a | ||
| 11 | // counter-based generator needs; the input is `seed ^ (domain + index)` so a | ||
| 12 | // caller varies `seed` per step and each op keeps a distinct `domain` for | ||
| 13 | // each of its noise streams (a position stream and a vocabulary stream must | ||
| 14 | // not collide at index 0). | ||
| 15 | // | ||
| 16 | // The uniform takes the high 24 bits of the hash to a float in [0, 1) — the | ||
| 17 | // same mapping the Philox ops use (rand_uniform, sample_logits), so the value | ||
| 18 | // set is identical: multiples of 2^-24, never 1.0. | ||
| 19 | // | ||
| 20 | // Every function is `inline` and, under nvcc, `__host__ __device__`, so the | ||
| 21 | // header compiles into both a .cpp and a .cu translation unit with identical | ||
| 22 | // integer arithmetic. Metal cannot include this header (MSL is compiled from | ||
| 23 | // source strings); src/metal/masked_diffusion.mm carries a verbatim copy. | ||
| 24 | |||
| 25 | #include <cstdint> | ||
| 26 | |||
| 27 | #if defined(__CUDACC__) | ||
| 28 | #define BROTENSOR_HASH_RNG_HD __host__ __device__ | ||
| 29 | #else | ||
| 30 | #define BROTENSOR_HASH_RNG_HD | ||
| 31 | #endif | ||
| 32 | |||
| 33 | namespace brotensor::detail { | ||
| 34 | |||
| 35 | // Domain constants for the ops that draw from this generator. Arbitrary | ||
| 36 | // large odd constants, far apart, so `domain + index` never overlaps across | ||
| 37 | // streams for any index a tensor can hold. | ||
| 38 | inline constexpr std::uint64_t kHashDomainMaskedDiffusionPosition = 0x5851F42D4C957F2DULL; | ||
| 39 | inline constexpr std::uint64_t kHashDomainMaskedDiffusionClass = 0x14057B7EF767814FULL; | ||
| 40 | |||
| 41 | // splitmix64 finaliser: one round of the SplitMix64 output function applied | ||
| 42 | // to `x` (including the golden-ratio increment, so x == 0 does not map to 0). | ||
| 43 | 649404 | BROTENSOR_HASH_RNG_HD inline std::uint64_t splitmix64(std::uint64_t x) { | |
| 44 | 649404 | x += 0x9E3779B97F4A7C15ULL; | |
| 45 | 649404 | x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ULL; | |
| 46 | 649404 | x = (x ^ (x >> 27)) * 0x94D049BB133111EBULL; | |
| 47 | 649404 | return x ^ (x >> 31); | |
| 48 | } | ||
| 49 | |||
| 50 | // Raw 64-bit hash for element `index` of stream (seed, domain). | ||
| 51 | 649685 | BROTENSOR_HASH_RNG_HD inline std::uint64_t hash_u64(std::uint64_t seed, | |
| 52 | std::uint64_t domain, | ||
| 53 | std::uint64_t index) { | ||
| 54 | 649685 | return splitmix64(seed ^ (domain + index)); | |
| 55 | } | ||
| 56 | |||
| 57 | // Uniform in [0, 1): top 24 bits of the hash / 2^24. Exact in FP32. | ||
| 58 | 649761 | BROTENSOR_HASH_RNG_HD inline float hash_uniform(std::uint64_t seed, | |
| 59 | std::uint64_t domain, | ||
| 60 | std::uint64_t index) { | ||
| 61 | 649761 | return static_cast<float>(hash_u64(seed, domain, index) >> 40) * | |
| 62 | (1.0f / 16777216.0f); | ||
| 63 | } | ||
| 64 | |||
| 65 | } // namespace brotensor::detail | ||
| 66 | |||
| 67 | #undef BROTENSOR_HASH_RNG_HD | ||
| 68 |