GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 94.7% 198 / 0 / 209
Functions: 100.0% 9 / 0 / 9
Branches: 64.7% 123 / 0 / 190
src/cpu/sample_logits.cpp
Line Branch Exec Source
1 // ─── CPU autoregressive logit sampling (CHUNK 7, family F) ──────────────────
2 //
3 // FP32 scalar host implementation of `sample_logits` — the next-token sampler
4 // shared by autoregressive generation loops (brosoundml codec-LM decoding and
5 // the brolm language-model project). CPU is FP32-only; the GPU vtable slot for
6 // this op stays null.
7 //
8 // Per-row algorithm (one row of an (N, V) logit matrix at a time):
9 // temperature scale -> softmax -> optional top-k filter -> optional top-p
10 // (nucleus) filter -> renormalize over the kept set -> inverse-CDF draw with
11 // a Philox-generated uniform in [0, 1). temperature == 0 short-circuits to a
12 // deterministic argmax (no RNG consumed).
13 //
14 // ── INT32 output ────────────────────────────────────────────────────────────
15 // indices — (N, 1) INT32 sampled token ids. Resized AND dtype-set to INT32,
16 // accessed via host_raw / host_raw_mut cast to int32_t* (host_f32 throws on a
17 // non-FP32 dtype).
18 //
19 // ── Philox (key, counter) ABI ───────────────────────────────────────────────
20 // Standard Philox 4x32-10 counter-based generator (PyTorch / JAX compatible).
21 // Row n is keyed by `key` and counter block {lo(counter+n), hi(counter+n),
22 // 0, 0}; the first of the four uint32 outputs becomes a uniform in [0, 1).
23 // See ops.h for the full ABI contract.
24
25 #include <brotensor/tensor.h>
26
27 #include <algorithm>
28 #include <cmath>
29 #include <cstddef>
30 #include <cstdint>
31 #include <stdexcept>
32 #include <string>
33 #include <vector>
34
35 namespace brotensor::detail::cpu {
36
37 namespace {
38
39 6 [[noreturn]] void fail(const char* op, const std::string& reason) {
40
7/14
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 6 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 6 times.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 6 times.
✓ Branch 12 taken 6 times.
✗ Branch 13 not taken.
 6 throw std::runtime_error(std::string("brotensor: ") + op + ": " + reason);
41 6 }
42
43 // ─── Philox 4x32-10 counter-based RNG ───────────────────────────────────────
44 //
45 // The reference 4x32-10 generator from Salmon et al. "Parallel Random Numbers:
46 // As Easy as 1, 2, 3" — the same construction PyTorch's CUDAGeneratorImpl and
47 // JAX's threefry/philox use. State is a 128-bit counter (4x uint32) plus a
48 // 64-bit key (2x uint32). One call to philox4x32() produces four uint32s; we
49 // consume the first.
50
51 constexpr uint32_t kPhiloxM0 = 0xD2511F53u; // multiplier, lane 0
52 constexpr uint32_t kPhiloxM1 = 0xCD9E8D57u; // multiplier, lane 1
53 constexpr uint32_t kPhiloxW0 = 0x9E3779B9u; // key bump, word 0 (golden ratio)
54 constexpr uint32_t kPhiloxW1 = 0xBB67AE85u; // key bump, word 1 (sqrt(3)-1)
55
56 // 32x32 -> 64 multiply, returning the high and low 32-bit words.
57 1006960 inline void mulhilo32(uint32_t a, uint32_t b, uint32_t& hi, uint32_t& lo) {
58 1006960 const uint64_t product = static_cast<uint64_t>(a) * static_cast<uint64_t>(b);
59 1006960 hi = static_cast<uint32_t>(product >> 32);
60 1006960 lo = static_cast<uint32_t>(product);
61 1006960 }
62
63 // One Philox round mixing the 128-bit counter under the current key words.
64 503480 inline void philox_round(uint32_t ctr[4], const uint32_t key[2]) {
65 uint32_t hi0, lo0, hi1, lo1;
66 503480 mulhilo32(kPhiloxM0, ctr[0], hi0, lo0);
67 503480 mulhilo32(kPhiloxM1, ctr[2], hi1, lo1);
68 503480 const uint32_t new0 = hi1 ^ ctr[1] ^ key[0];
69 503480 const uint32_t new1 = lo1;
70 503480 const uint32_t new2 = hi0 ^ ctr[3] ^ key[1];
71 503480 const uint32_t new3 = lo0;
72 503480 ctr[0] = new0;
73 503480 ctr[1] = new1;
74 503480 ctr[2] = new2;
75 503480 ctr[3] = new3;
76 503480 }
77
78 // Full Philox 4x32-10: ten rounds, bumping the key between rounds.
79 50348 inline void philox4x32(uint32_t ctr[4], uint32_t key[2]) {
80
2/2
✓ Branch 0 taken 503480 times.
✓ Branch 1 taken 50348 times.
 553828 for (int round = 0; round < 10; ++round) {
81 503480 philox_round(ctr, key);
82
2/2
✓ Branch 0 taken 50348 times.
✓ Branch 1 taken 453132 times.
 503480 if (round < 9) {
83 453132 key[0] += kPhiloxW0;
84 453132 key[1] += kPhiloxW1;
85 453132 }
86 503480 }
87 50348 }
88
89 // Draw one uniform in [0, 1) for substream `substream`, seeded by `key64`.
90 // The 64-bit key is split low-word-first into the two Philox key words; the
91 // substream index is the low 64 bits of the 128-bit counter (low-word-first),
92 // the upper 64 counter bits are zero. The first of the four uint32 outputs is
93 // mapped to [0, 1) via its top 24 bits / 2^24 — matching xavier_init's u01.
94 50348 inline float philox_uniform(uint64_t key64, uint64_t substream) {
95 100696 uint32_t key[2] = {
96 50348 static_cast<uint32_t>(key64 & 0xFFFFFFFFull),
97 50348 static_cast<uint32_t>(key64 >> 32),
98 };
99 100696 uint32_t ctr[4] = {
100 50348 static_cast<uint32_t>(substream & 0xFFFFFFFFull),
101 50348 static_cast<uint32_t>(substream >> 32),
102 0u,
103 0u,
104 };
105 50348 philox4x32(ctr, key);
106 50348 return static_cast<float>(ctr[0] >> 8) / 16777216.0f; // top 24 bits
107 }
108
109 } // namespace
110
111 // ─── sample_logits ──────────────────────────────────────────────────────────
112
113 50256 void sample_logits(const ::brotensor::Tensor& logits, float temperature,
114 int top_k, float top_p, uint64_t key, uint64_t counter,
115 ::brotensor::Tensor& indices) {
116
1/2
✓ Branch 0 taken 50256 times.
✗ Branch 1 not taken.
 50256 if (logits.dtype != ::brotensor::Dtype::FP32) {
117
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 3 fail("sample_logits", "logits must be FP32 (CPU backend is FP32-only)");
118 }
119
2/2
✓ Branch 0 taken 50255 times.
✓ Branch 1 taken 1 time.
 50256 if (temperature < 0.0f) {
120
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 fail("sample_logits", "temperature must be >= 0");
121 }
122
2/2
✓ Branch 0 taken 50254 times.
✓ Branch 1 taken 1 time.
 50255 if (top_k < 0) {
123
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 fail("sample_logits", "top_k must be >= 0");
124 }
125
2/2
✓ Branch 0 taken 50253 times.
✓ Branch 1 taken 1 time.
 50254 if (top_p < 0.0f) {
126
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 fail("sample_logits", "top_p must be >= 0");
127 }
128
129 50253 const int N = logits.rows;
130 50253 const int V = logits.cols;
131
2/4
✓ Branch 0 taken 50253 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 50253 times.
✗ Branch 3 not taken.
 50253 if (N > 0 && V == 0) {
132 ✗ fail("sample_logits", "vocabulary size (logits.cols) must be > 0");
133 }
134
135 // indices: (N, 1) INT32 — resize AND dtype-set.
136
4/6
✓ Branch 0 taken 63 times.
✓ Branch 1 taken 50190 times.
✓ Branch 2 taken 63 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 63 times.
 50253 if (indices.rows != N || indices.cols != 1 ||
137 63 indices.dtype != ::brotensor::Dtype::INT32) {
138 50190 indices.resize(N, 1, ::brotensor::Dtype::INT32);
139 50190 }
140
1/2
✓ Branch 0 taken 50253 times.
✗ Branch 1 not taken.
 50253 if (N == 0) return;
141
142 50253 const float* lp = logits.host_f32();
143 50253 int32_t* ip = static_cast<int32_t*>(indices.host_raw_mut());
144
145 // Scratch reused across rows.
146 50253 std::vector<float> prob(static_cast<std::size_t>(V));
147
1/2
✓ Branch 0 taken 50253 times.
✗ Branch 1 not taken.
 50253 std::vector<int> order(static_cast<std::size_t>(V));
148
149
2/2
✓ Branch 0 taken 50363 times.
✓ Branch 1 taken 50253 times.
 100616 for (int n = 0; n < N; ++n) {
150 50363 const float* row = lp + static_cast<std::size_t>(n) * V;
151
152 // ── Greedy: temperature == 0 -> deterministic argmax, no RNG. ──
153
2/2
✓ Branch 0 taken 21 times.
✓ Branch 1 taken 50342 times.
 50363 if (temperature == 0.0f) {
154 21 float best_v = -3.4028235e38f; // -FLT_MAX
155 21 int best_i = 0;
156
2/2
✓ Branch 0 taken 2333 times.
✓ Branch 1 taken 21 times.
 2354 for (int v = 0; v < V; ++v) {
157
2/2
✓ Branch 0 taken 2237 times.
✓ Branch 1 taken 96 times.
 2333 if (row[v] > best_v) { best_v = row[v]; best_i = v; }
158 2333 }
159 21 ip[n] = static_cast<int32_t>(best_i);
160 21 continue;
161 }
162
163 // ── 1. temperature scale + 2. softmax (numerically stable). ──
164 50342 float max_logit = -3.4028235e38f;
165
2/2
✓ Branch 0 taken 136240 times.
✓ Branch 1 taken 50342 times.
 186582 for (int v = 0; v < V; ++v) {
166 136240 const float s = row[v] / temperature;
167
2/2
✓ Branch 0 taken 25234 times.
✓ Branch 1 taken 111006 times.
 136240 if (s > max_logit) max_logit = s;
168 136240 }
169 50342 double sum = 0.0;
170
2/2
✓ Branch 0 taken 136240 times.
✓ Branch 1 taken 50342 times.
 186582 for (int v = 0; v < V; ++v) {
171 136240 const float s = row[v] / temperature;
172 136240 const float e = std::exp(s - max_logit);
173 136240 prob[v] = e;
174 136240 sum += e;
175 136240 }
176
1/2
✓ Branch 0 taken 50342 times.
✗ Branch 1 not taken.
 50342 const float inv_sum = (sum > 0.0) ? static_cast<float>(1.0 / sum) : 0.0f;
177
2/2
✓ Branch 0 taken 136240 times.
✓ Branch 1 taken 50342 times.
 186582 for (int v = 0; v < V; ++v) prob[v] *= inv_sum;
178
179 // ── 3. top-k filter: keep the top_k highest-probability tokens. ──
180 // Only the top `bound` entries are ever consulted below (top-p can
181 // only shrink the kept set further, never grow it past top_k), so a
182 // partial_sort bounded to `bound` reproduces the same output as a
183 // full sort of all V entries while avoiding O(V log V) work. The
184 // comparator breaks ties by ascending index to match std::stable_sort
185 // applied to the identity-ordered 0..V-1 sequence.
186 50342 int bound = V;
187
3/4
✓ Branch 0 taken 92 times.
✓ Branch 1 taken 50250 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 92 times.
 50342 if (top_k > 0 && top_k < bound) bound = top_k;
188
189
2/2
✓ Branch 0 taken 136240 times.
✓ Branch 1 taken 50342 times.
 186582 for (int v = 0; v < V; ++v) order[v] = v;
190
2/4
✓ Branch 0 taken 50342 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 50342 times.
✗ Branch 3 not taken.
 100684 std::partial_sort(order.begin(), order.begin() + bound, order.end(),
191 212240 [&](int a, int b) {
192
2/2
✓ Branch 0 taken 36747 times.
✓ Branch 1 taken 125151 times.
 287049 return prob[a] > prob[b] ||
193
2/2
✓ Branch 0 taken 124134 times.
✓ Branch 1 taken 1017 times.
 125151 (prob[a] == prob[b] && a < b);
194 });
195
196 50342 int keep = bound;
197
198 // ── 4. top-p (nucleus): keep the smallest high-prob set whose
199 // cumulative probability >= top_p, applied to the top-k
200 // survivors. top_p >= 1.0 disables it. ──
201
2/2
✓ Branch 0 taken 50176 times.
✓ Branch 1 taken 166 times.
 50342 if (top_p < 1.0f) {
202 166 double cum = 0.0;
203 166 int nucleus = 0;
204
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 294 times.
 294 for (int r = 0; r < keep; ++r) {
205 294 cum += prob[order[r]];
206 294 ++nucleus;
207
2/2
✓ Branch 0 taken 166 times.
✓ Branch 1 taken 128 times.
 294 if (cum >= static_cast<double>(top_p)) break;
208 128 }
209
1/2
✓ Branch 0 taken 166 times.
✗ Branch 1 not taken.
 166 if (nucleus < 1) nucleus = 1; // always keep at least one token
210 166 keep = nucleus;
211 166 }
212
213 // ── 5. renormalize over the kept set. ──
214 50342 double kept_sum = 0.0;
215
2/2
✓ Branch 0 taken 131058 times.
✓ Branch 1 taken 50342 times.
 181400 for (int r = 0; r < keep; ++r) kept_sum += prob[order[r]];
216
217 // ── 6. inverse-CDF draw with a Philox uniform for substream
218 // (counter + n). ──
219
1/2
✓ Branch 0 taken 50342 times.
✗ Branch 1 not taken.
 50342 const float u = philox_uniform(key, counter + static_cast<uint64_t>(n));
220 50342 int chosen = order[0];
221
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 50342 times.
 50342 if (kept_sum > 0.0) {
222 50342 const double target = static_cast<double>(u) * kept_sum;
223 50342 double acc = 0.0;
224 50342 chosen = order[keep - 1]; // fallback: last kept (covers u≈1).
225
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 67524 times.
 67524 for (int r = 0; r < keep; ++r) {
226 67524 acc += prob[order[r]];
227
2/2
✓ Branch 0 taken 50342 times.
✓ Branch 1 taken 17182 times.
 67524 if (target < acc) { chosen = order[r]; break; }
228 17182 }
229 50342 }
230 50342 ip[n] = static_cast<int32_t>(chosen);
231 50342 }
232 50256 }
233
234 // ─── sample_logits_into ─────────────────────────────────────────────────────
235 //
236 // Graph-capturable variant: the Philox base counter lives in a device tensor
237 // (counter[0]) and is advanced by N on completion; scratch is caller-owned.
238 // On CPU there is no graph capture, so this is the same per-row pipeline run
239 // directly — its purpose here is to keep the op surface and the (key, counter)
240 // draw byte-identical to the CUDA/Metal capture path. Draw for row n uses
241 // substream (counter[0] + n), matching sample_logits with that base counter.
242
243 7 void sample_logits_into(const ::brotensor::Tensor& logits, float temperature,
244 int top_k, float top_p, uint64_t key,
245 ::brotensor::Tensor& counter,
246 ::brotensor::Tensor& scratch,
247 ::brotensor::Tensor& indices) {
248 7 const char* op = "sample_logits_into";
249
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (logits.dtype != ::brotensor::Dtype::FP32) {
250
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 3 fail(op, "logits must be FP32 (CPU backend is FP32-only)");
251 }
252
1/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 7 if (temperature < 0.0f) fail(op, "temperature must be >= 0");
253
1/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 7 if (top_k < 0) fail(op, "top_k must be >= 0");
254
1/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 7 if (top_p < 0.0f) fail(op, "top_p must be >= 0");
255
256 7 const int N = logits.rows;
257 7 const int V = logits.cols;
258
2/6
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 7 if (N > 0 && V == 0) fail(op, "vocabulary size (logits.cols) must be > 0");
259
260
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 1 time.
 7 if (counter.dtype != ::brotensor::Dtype::INT32 ||
261 6 static_cast<std::size_t>(counter.rows) * counter.cols < 1) {
262
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 fail(op, "counter must be an INT32 tensor with >= 1 element");
263 }
264 6 const std::size_t nv = static_cast<std::size_t>(N) * V;
265
2/2
✓ Branch 0 taken 5 times.
✓ Branch 1 taken 1 time.
 6 if (scratch.dtype != ::brotensor::Dtype::FP32 ||
266 5 static_cast<std::size_t>(scratch.rows) * scratch.cols < 3 * nv) {
267
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 fail(op, "scratch must be FP32 with at least 3*N*V elements");
268 }
269 // indices must be a pre-sized (N,1) INT32 — never resized here (a resize
270 // would allocate, which the CUDA path forbids mid-capture).
271
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 1 time.
 5 if (indices.rows != N || indices.cols != 1 ||
272 4 indices.dtype != ::brotensor::Dtype::INT32) {
273
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 fail(op, "indices must be a pre-sized (N,1) INT32 tensor");
274 }
275
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (N == 0) return;
276
277 4 int32_t* cp = static_cast<int32_t*>(counter.host_raw_mut());
278 4 const uint64_t base =
279 4 static_cast<uint64_t>(static_cast<uint32_t>(cp[0]));
280
281 4 const float* lp = logits.host_f32();
282 4 int32_t* ip = static_cast<int32_t*>(indices.host_raw_mut());
283
284 4 std::vector<float> prob(static_cast<std::size_t>(V));
285
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 std::vector<int> order(static_cast<std::size_t>(V));
286
287
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 4 times.
 12 for (int n = 0; n < N; ++n) {
288 8 const float* row = lp + static_cast<std::size_t>(n) * V;
289
290
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 6 times.
 8 if (temperature == 0.0f) {
291 2 float best_v = -3.4028235e38f;
292 2 int best_i = 0;
293
2/2
✓ Branch 0 taken 10 times.
✓ Branch 1 taken 2 times.
 12 for (int v = 0; v < V; ++v) {
294
2/2
✓ Branch 0 taken 7 times.
✓ Branch 1 taken 3 times.
 10 if (row[v] > best_v) { best_v = row[v]; best_i = v; }
295 10 }
296 2 ip[n] = static_cast<int32_t>(best_i);
297 2 continue;
298 }
299
300 6 float max_logit = -3.4028235e38f;
301
2/2
✓ Branch 0 taken 48 times.
✓ Branch 1 taken 6 times.
 54 for (int v = 0; v < V; ++v) {
302 48 const float s = row[v] / temperature;
303
2/2
✓ Branch 0 taken 32 times.
✓ Branch 1 taken 16 times.
 48 if (s > max_logit) max_logit = s;
304 48 }
305 6 double sum = 0.0;
306
2/2
✓ Branch 0 taken 48 times.
✓ Branch 1 taken 6 times.
 54 for (int v = 0; v < V; ++v) {
307 48 const float s = row[v] / temperature;
308 48 const float e = std::exp(s - max_logit);
309 48 prob[v] = e;
310 48 sum += e;
311 48 }
312
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 const float inv_sum = (sum > 0.0) ? static_cast<float>(1.0 / sum) : 0.0f;
313
2/2
✓ Branch 0 taken 48 times.
✓ Branch 1 taken 6 times.
 54 for (int v = 0; v < V; ++v) prob[v] *= inv_sum;
314
315 // See the matching comment in sample_logits() above: bound to the
316 // number of entries top-p/top-k can ever consult, and break ties by
317 // ascending index to match std::stable_sort's behavior exactly.
318 6 int bound = V;
319
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 6 if (top_k > 0 && top_k < bound) bound = top_k;
320
321
2/2
✓ Branch 0 taken 48 times.
✓ Branch 1 taken 6 times.
 54 for (int v = 0; v < V; ++v) order[v] = v;
322
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
 12 std::partial_sort(order.begin(), order.begin() + bound, order.end(),
323 138 [&](int a, int b) {
324
2/2
✓ Branch 0 taken 38 times.
✓ Branch 1 taken 94 times.
 226 return prob[a] > prob[b] ||
325
2/2
✓ Branch 0 taken 84 times.
✓ Branch 1 taken 10 times.
 94 (prob[a] == prob[b] && a < b);
326 });
327
328 6 int keep = bound;
329
330
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (top_p < 1.0f) {
331 ✗ double cum = 0.0;
332 ✗ int nucleus = 0;
333 ✗ for (int r = 0; r < keep; ++r) {
334 ✗ cum += prob[order[r]];
335 ✗ ++nucleus;
336 ✗ if (cum >= static_cast<double>(top_p)) break;
337 ✗ }
338 ✗ if (nucleus < 1) nucleus = 1;
339 ✗ keep = nucleus;
340 ✗ }
341
342 6 double kept_sum = 0.0;
343
2/2
✓ Branch 0 taken 48 times.
✓ Branch 1 taken 6 times.
 54 for (int r = 0; r < keep; ++r) kept_sum += prob[order[r]];
344
345
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 const float u = philox_uniform(key, base + static_cast<uint64_t>(n));
346 6 int chosen = order[0];
347
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (kept_sum > 0.0) {
348 6 const double target = static_cast<double>(u) * kept_sum;
349 6 double acc = 0.0;
350 6 chosen = order[keep - 1];
351
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 21 times.
 21 for (int r = 0; r < keep; ++r) {
352 21 acc += prob[order[r]];
353
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 15 times.
 21 if (target < acc) { chosen = order[r]; break; }
354 15 }
355 6 }
356 6 ip[n] = static_cast<int32_t>(chosen);
357 6 }
358
359 // Advance the base counter by the rows drawn (matches the device path so a
360 // repeated call continues the same Philox stream). Greedy consumes no RNG,
361 // so its counter is left untouched — mirroring the CUDA/Metal path.
362
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 3 times.
 4 if (temperature != 0.0f)
363 3 cp[0] = static_cast<int32_t>(
364 3 static_cast<uint32_t>(base + static_cast<uint64_t>(N)));
365 7 }
366
367 } // namespace brotensor::detail::cpu
368