GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 92.5% 99 / 0 / 107
Functions: 91.7% 11 / 0 / 12
Branches: 51.5% 34 / 0 / 66
src/cpu/noise.cpp
Line Branch Exec Source
1 // ─── CPU counter-based noise generation ─────────────────────────────────────
2 //
3 // Philox 4x32-10 stream (PyTorch / JAX compatible) implementing four host-side
4 // FP32 noise ops:
5 // randn — N(0, 1) via Box-Muller
6 // rand_uniform — U[0, 1)
7 // rand_bernoulli — 0/1 mask at probability p
8 // randn_truncated — rejection-sampled normal in [lo, hi]
9 //
10 // All four require Y FP32 and pre-sized; the op fills rows*cols elements in
11 // row-major linear order. Element i draws from substream (counter + i); see
12 // ops.h for the full ABI contract.
13 //
14 // The Philox construction here is byte-identical to src/cpu/sample_logits.cpp
15 // and the CUDA / Metal noise implementations, so a given (key, counter) yields
16 // the same draws on every backend.
17
18 #include <brotensor/tensor.h>
19
20 #include <cmath>
21 #include <cstddef>
22 #include <cstdint>
23 #include <stdexcept>
24 #include <string>
25
26 namespace brotensor::detail::cpu {
27
28 namespace {
29
30 ✗ [[noreturn]] void fail(const char* op, const std::string& reason) {
31 ✗ throw std::runtime_error(std::string("brotensor: ") + op + ": " + reason);
32 ✗ }
33
34 // ── Philox 4x32-10 ──────────────────────────────────────────────────────────
35 // Same construction as sample_logits.cpp; duplicated rather than shared so the
36 // translation units stay independent.
37
38 8060200 inline void mulhilo32(uint32_t a, uint32_t b, uint32_t& hi, uint32_t& lo) {
39 8060200 const uint64_t p = static_cast<uint64_t>(a) * static_cast<uint64_t>(b);
40 8060200 hi = static_cast<uint32_t>(p >> 32);
41 8060200 lo = static_cast<uint32_t>(p);
42 8060200 }
43
44 4030100 inline void philox_round(uint32_t ctr[4], const uint32_t key[2]) {
45 uint32_t hi0, lo0, hi1, lo1;
46 4030100 mulhilo32(0xD2511F53u, ctr[0], hi0, lo0);
47 4030100 mulhilo32(0xCD9E8D57u, ctr[2], hi1, lo1);
48 4030100 const uint32_t n0 = hi1 ^ ctr[1] ^ key[0];
49 4030100 const uint32_t n1 = lo1;
50 4030100 const uint32_t n2 = hi0 ^ ctr[3] ^ key[1];
51 4030100 const uint32_t n3 = lo0;
52 4030100 ctr[0] = n0; ctr[1] = n1; ctr[2] = n2; ctr[3] = n3;
53 4030100 }
54
55 // Run all ten rounds for substream `substream` under `key64`. Output is the
56 // four mixed counter words ctr[0..3] (each a uniformly distributed uint32).
57 403010 inline void philox4x32(uint64_t key64, uint64_t substream, uint32_t out[4]) {
58 806020 uint32_t key[2] = {
59 403010 static_cast<uint32_t>(key64 & 0xFFFFFFFFull),
60 403010 static_cast<uint32_t>(key64 >> 32),
61 };
62 806020 uint32_t ctr[4] = {
63 403010 static_cast<uint32_t>(substream & 0xFFFFFFFFull),
64 403010 static_cast<uint32_t>(substream >> 32),
65 0u, 0u,
66 };
67
2/2
✓ Branch 0 taken 4030100 times.
✓ Branch 1 taken 403010 times.
 4433110 for (int r = 0; r < 10; ++r) {
68 4030100 philox_round(ctr, key);
69
2/2
✓ Branch 0 taken 403010 times.
✓ Branch 1 taken 3627090 times.
 4030100 if (r < 9) {
70 3627090 key[0] += 0x9E3779B9u;
71 3627090 key[1] += 0xBB67AE85u;
72 3627090 }
73 4030100 }
74 403010 out[0] = ctr[0]; out[1] = ctr[1]; out[2] = ctr[2]; out[3] = ctr[3];
75 403010 }
76
77 // Top 24 bits / 2^24 -> uniform in [0, 1). Matches xavier_init / sample_logits.
78 555764 inline float u01_from(uint32_t w) {
79 555764 return static_cast<float>(w >> 8) * (1.0f / 16777216.0f);
80 }
81
82 // One uniform draw for substream `substream` (use ctr[0] only).
83 250256 inline float philox_uniform(uint64_t key64, uint64_t substream) {
84 uint32_t ctr[4];
85 250256 philox4x32(key64, substream, ctr);
86 250256 return u01_from(ctr[0]);
87 }
88
89 // One standard-normal draw via Box-Muller on (ctr[0], ctr[1]). We use
90 // 1 - u1 so log() stays finite for u1 == 0.
91 152754 inline float philox_normal(uint64_t key64, uint64_t substream) {
92 uint32_t ctr[4];
93 152754 philox4x32(key64, substream, ctr);
94 152754 const float u1 = 1.0f - u01_from(ctr[0]); // (0, 1]
95 152754 const float u2 = u01_from(ctr[1]); // [0, 1)
96 152754 constexpr double kTwoPi = 2.0 * 3.14159265358979323846;
97 152754 const float radius = std::sqrt(-2.0f * std::log(u1));
98 152754 const float theta =
99 152754 static_cast<float>(kTwoPi * static_cast<double>(u2));
100 152754 return radius * std::cos(theta);
101 }
102
103 // Common entry validation for the four ops.
104 16 inline std::size_t check_y(const char* op, const ::brotensor::Tensor& Y) {
105
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (Y.dtype != ::brotensor::Dtype::FP32) {
106 ✗ fail(op, "Y must be FP32");
107 }
108
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (Y.rows < 0 || Y.cols < 0) {
109 ✗ fail(op, "Y has negative dimension");
110 }
111 32 const std::size_t n = static_cast<std::size_t>(Y.rows) *
112 16 static_cast<std::size_t>(Y.cols);
113
3/4
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
 16 if (n != 0 && Y.data == nullptr) {
114 ✗ fail(op, "Y is uncommitted; pre-allocate before calling");
115 }
116 16 return n;
117 ✗ }
118
119 } // namespace
120
121 // ─── randn ──────────────────────────────────────────────────────────────────
122 6 void randn(uint64_t key, uint64_t counter, ::brotensor::Tensor& Y) {
123 6 const std::size_t n = check_y("randn", Y);
124
2/2
✓ Branch 0 taken 5 times.
✓ Branch 1 taken 1 time.
 6 if (n == 0) return;
125 5 float* y = Y.host_f32_mut();
126
2/2
✓ Branch 0 taken 100352 times.
✓ Branch 1 taken 5 times.
 100357 for (std::size_t i = 0; i < n; ++i) {
127 100352 y[i] = philox_normal(key, counter + static_cast<uint64_t>(i));
128 100352 }
129 6 }
130
131 // ─── rand_uniform ───────────────────────────────────────────────────────────
132 4 void rand_uniform(uint64_t key, uint64_t counter, ::brotensor::Tensor& Y) {
133 4 const std::size_t n = check_y("rand_uniform", Y);
134
2/2
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 1 time.
 4 if (n == 0) return;
135 3 float* y = Y.host_f32_mut();
136
2/2
✓ Branch 0 taken 100256 times.
✓ Branch 1 taken 3 times.
 100259 for (std::size_t i = 0; i < n; ++i) {
137 100256 y[i] = philox_uniform(key, counter + static_cast<uint64_t>(i));
138 100256 }
139 4 }
140
141 // ─── rand_bernoulli ─────────────────────────────────────────────────────────
142 4 void rand_bernoulli(float p, uint64_t key, uint64_t counter,
143 ::brotensor::Tensor& Y) {
144 4 const char* op = "rand_bernoulli";
145
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!(p >= 0.0f && p <= 1.0f)) {
146 ✗ fail(op, "p must be in [0, 1]");
147 }
148 4 const std::size_t n = check_y(op, Y);
149
2/2
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 1 time.
 4 if (n == 0) return;
150 3 float* y = Y.host_f32_mut();
151
2/2
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 150000 times.
 150003 for (std::size_t i = 0; i < n; ++i) {
152 150000 const float u = philox_uniform(key, counter + static_cast<uint64_t>(i));
153 150000 y[i] = (u < p) ? 1.0f : 0.0f;
154 150000 }
155 4 }
156
157 // ─── randn_truncated ────────────────────────────────────────────────────────
158 2 void randn_truncated(float lo, float hi, uint64_t key, uint64_t counter,
159 ::brotensor::Tensor& Y) {
160 2 const char* op = "randn_truncated";
161
1/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 2 if (!(lo < hi)) fail(op, "lo must be < hi");
162 2 const std::size_t n = check_y(op, Y);
163
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 1 time.
 2 if (n == 0) return;
164 1 float* y = Y.host_f32_mut();
165 1 constexpr int kMaxRetries = 64;
166
2/2
✓ Branch 0 taken 50000 times.
✓ Branch 1 taken 1 time.
 50001 for (std::size_t i = 0; i < n; ++i) {
167 50000 float z = 0.0f;
168
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 52402 times.
 52402 for (int r = 0; r < kMaxRetries; ++r) {
169 104804 const uint64_t sub = counter + static_cast<uint64_t>(i) +
170 104804 static_cast<uint64_t>(r) *
171 52402 static_cast<uint64_t>(n);
172 52402 z = philox_normal(key, sub);
173
4/4
✓ Branch 0 taken 51199 times.
✓ Branch 1 taken 1203 times.
✓ Branch 2 taken 50000 times.
✓ Branch 3 taken 1199 times.
 52402 if (z >= lo && z <= hi) break;
174 2402 }
175
1/2
✓ Branch 0 taken 50000 times.
✗ Branch 1 not taken.
 50000 if (z < lo) z = lo;
176
1/2
✓ Branch 0 taken 50000 times.
✗ Branch 1 not taken.
 50000 if (z > hi) z = hi;
177 50000 y[i] = z;
178 50000 }
179 2 }
180
181 } // namespace brotensor::detail::cpu
182