GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 99.3% 151 / 0 / 152
Functions: 100.0% 8 / 0 / 8
Branches: 60.6% 109 / 0 / 180
src/cpu/masked_diffusion.cpp
Line Branch Exec Source
1 // ─── CPU masked-diffusion token selection ──────────────────────────────────
2 //
3 // FP32 scalar host implementation of `masked_diffusion_scores` and
4 // `masked_diffusion_commit` — the per-step token selection of a masked-
5 // diffusion language model over a (C codebooks, T frames) grid (the OmniVoice
6 // port in brosoundml). See ops/sampling.h for the full contract and the
7 // pseudocode this follows.
8 //
9 // This is the parity reference for the CUDA kernel in src/cuda/
10 // masked_diffusion.cu, so the arithmetic is spelled out in the exact order the
11 // kernel uses:
12 // * a row's log-sum-exp is max + (float)log(sum), the sum accumulated in
13 // FP64 from FP32 expf terms (order-independent to the last FP32 ulp);
14 // * the guided logit is a + gs * (a - b) with a = c - lse_c, b = u - lse_u,
15 // each operation rounded separately (the kernel uses __fadd_rn /
16 // __fmul_rn / __fsub_rn so nvcc cannot contract it into an FMA);
17 // * the noise comes from detail/hash_rng.h, identical integer arithmetic on
18 // both backends.
19 // The top-k class filter keeps the first k entries in (value desc, index asc)
20 // order: everything above the k-th value, then the lowest-indexed entries
21 // equal to it until k are kept.
22 //
23 // ── ACCUMULATION ────────────────────────────────────────────────────────────
24 // masked_diffusion_scores — pred, scores and confidence OVERWRITTEN.
25 // masked_diffusion_commit — tokens / unmask_step written only at idx[0..k).
26
27 #include <brotensor/tensor.h>
28 #include <brotensor/detail/hash_rng.h>
29 #include <brotensor/detail/cpu/thread_pool.h>
30
31 #include <algorithm>
32 #include <cmath>
33 #include <cstddef>
34 #include <cstdint>
35 #include <limits>
36 #include <stdexcept>
37 #include <string>
38 #include <vector>
39
40 namespace brotensor::detail::cpu {
41
42 namespace {
43
44 4 [[noreturn]] void fail(const char* op, const std::string& reason) {
45
7/14
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 4 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 4 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 4 times.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 4 times.
✓ Branch 12 taken 4 times.
✗ Branch 13 not taken.
 4 throw std::runtime_error(std::string("brotensor: ") + op + ": " + reason);
46 4 }
47
48 constexpr float kNegInf = -std::numeric_limits<float>::infinity();
49
50 313290 inline float gumbel_from_uniform(float u) {
51 313290 return -std::log(-std::log(u + 1e-10f) + 1e-10f);
52 }
53
54 // max + log(sum exp(x - max)) over a row of n floats; sum in FP64.
55 22442 inline float row_lse(const float* x, int n, float& max_out) {
56 22442 float m = kNegInf;
57
4/4
✓ Branch 0 taken 17937984 times.
✓ Branch 1 taken 22442 times.
✓ Branch 2 taken 165755 times.
✓ Branch 3 taken 17772229 times.
 17960426 for (int i = 0; i < n; ++i) m = (x[i] > m) ? x[i] : m;
58 22442 double s = 0.0;
59
2/2
✓ Branch 0 taken 16273192 times.
✓ Branch 1 taken 22442 times.
 16295634 for (int i = 0; i < n; ++i) s += static_cast<double>(std::exp(x[i] - m));
60 22442 max_out = m;
61 22442 return m + static_cast<float>(std::log(s));
62 }
63
64 // (a, ia) precedes (b, ib) in descending-value / ascending-index order.
65 7944170 inline bool prefers(float a, int ia, float b, int ib) {
66
2/2
✓ Branch 0 taken 7916002 times.
✓ Branch 1 taken 28168 times.
 7944170 if (a != b) return a > b;
67 28168 return ia < ib;
68 7944170 }
69
70 } // namespace
71
72 // ─── masked_diffusion_scores ────────────────────────────────────────────────
73
74 15 void masked_diffusion_scores(const ::brotensor::Tensor& logits,
75 const ::brotensor::Tensor& tokens,
76 int T, int C, int V, int mask_id,
77 float guidance_scale, float layer_penalty,
78 float position_temperature, float class_temperature,
79 float class_top_frac, uint64_t seed,
80 ::brotensor::Tensor& pred, ::brotensor::Tensor& scores,
81 ::brotensor::Tensor& confidence) {
82 15 const char* op = "masked_diffusion_scores";
83
1/4
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 18 if (logits.dtype != Dtype::FP32) fail(op, "logits must be FP32 (CPU backend is FP32-only)");
84
1/4
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 15 if (tokens.dtype != Dtype::INT32) fail(op, "tokens must be INT32");
85
1/4
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 15 if (T < 1 || C < 1 || V < 1) fail(op, "T, C and V must be >= 1");
86
3/4
✓ Branch 0 taken 14 times.
✓ Branch 1 taken 1 time.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
 15 if (mask_id < 0 || mask_id >= V) fail(op, "mask_id must be in [0, V)");
87 14 const bool guided = guidance_scale != 0.0f;
88
2/2
✓ Branch 0 taken 9 times.
✓ Branch 1 taken 5 times.
 14 const int R = guided ? 2 * T : T;
89
2/2
✓ Branch 0 taken 13 times.
✓ Branch 1 taken 1 time.
 14 if (logits.rows != R || logits.cols != C * V) {
90
4/8
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
✓ Branch 4 taken 1 time.
✗ Branch 5 not taken.
✓ Branch 6 taken 1 time.
✗ Branch 7 not taken.
 3 fail(op, "logits must be (" + std::to_string(R) + ", C*V=" +
91
5/10
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
✓ Branch 4 taken 1 time.
✗ Branch 5 not taken.
✓ Branch 6 taken 1 time.
✗ Branch 7 not taken.
✓ Branch 8 taken 1 time.
✗ Branch 9 not taken.
 2 std::to_string(C * V) + ") for T=" + std::to_string(T) +
92
2/4
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
 2 (guided ? " with" : " without") + " guidance, got (" +
93
5/10
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
✓ Branch 4 taken 1 time.
✗ Branch 5 not taken.
✓ Branch 6 taken 1 time.
✗ Branch 7 not taken.
✓ Branch 8 taken 1 time.
✗ Branch 9 not taken.
 2 std::to_string(logits.rows) + ", " + std::to_string(logits.cols) + ")");
94 }
95
3/4
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 1 time.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
 13 if (tokens.rows != C || tokens.cols != T) fail(op, "tokens must be (C, T)");
96
97 12 const int rows = C * T;
98
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 12 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 12 if (pred.rows != C || pred.cols != T || pred.dtype != Dtype::INT32) {
99 12 pred.resize(C, T, Dtype::INT32);
100 12 }
101
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 12 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 12 if (scores.rows != C || scores.cols != T || scores.dtype != Dtype::FP32) {
102 12 scores.resize(C, T, Dtype::FP32);
103 12 }
104
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 12 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 12 if (confidence.rows != C || confidence.cols != T || confidence.dtype != Dtype::FP32) {
105 12 confidence.resize(C, T, Dtype::FP32);
106 12 }
107
108 // Top-k class filter width, fixed per call (ceil(frac * V), clamped).
109 12 int k_keep = V;
110
2/2
✓ Branch 0 taken 9 times.
✓ Branch 1 taken 3 times.
 12 if (class_temperature > 0.0f) {
111 3 const double kd = std::ceil(static_cast<double>(class_top_frac) * V);
112
2/4
✗ Branch 0 not taken.
✓ Branch 1 taken 3 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 3 times.
 3 k_keep = (kd < 1.0) ? 1 : (kd > V ? V : static_cast<int>(kd));
113 3 }
114
115 12 const float* L = logits.host_f32();
116 12 const int32_t* tok = static_cast<const int32_t*>(tokens.host_raw());
117 12 int32_t* pr = static_cast<int32_t*>(pred.host_raw_mut());
118 12 float* sc = scores.host_f32_mut();
119 12 float* cf = confidence.host_f32_mut();
120 12 const std::size_t row_stride = static_cast<std::size_t>(C) * V;
121
122
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 9635 parallel_for(static_cast<std::size_t>(rows), [&](std::size_t pi) {
123 9623 const int p = static_cast<int>(pi);
124 9623 const int c = p / T;
125 9623 const int t = p - c * T;
126 19246 const float* cond = L + static_cast<std::size_t>(t) * row_stride +
127 9623 static_cast<std::size_t>(c) * V;
128
129
4/4
✓ Branch 0 taken 9620 times.
✓ Branch 1 taken 3 times.
✓ Branch 2 taken 9620 times.
✓ Branch 3 taken 3 times.
 9623 thread_local std::vector<float> lp_buf, sort_buf;
130
2/2
✓ Branch 0 taken 9617 times.
✓ Branch 1 taken 6 times.
 9623 if (static_cast<int>(lp_buf.size()) < V) {
131 6 lp_buf.resize(V);
132 6 sort_buf.resize(V);
133 6 }
134 9623 float* lp = lp_buf.data();
135
136 // ── log_probs ──
137 9623 float dummy = 0.0f;
138
2/2
✓ Branch 0 taken 6423 times.
✓ Branch 1 taken 3200 times.
 9623 if (guided) {
139 12846 const float* unc = L + static_cast<std::size_t>(T + t) * row_stride +
140 6423 static_cast<std::size_t>(c) * V;
141 6423 const float lse_c = row_lse(cond, V, dummy);
142 6423 const float lse_u = row_lse(unc, V, dummy);
143
2/2
✓ Branch 0 taken 6133500 times.
✓ Branch 1 taken 6423 times.
 6139923 for (int v = 0; v < V; ++v) {
144 6133500 const float a = cond[v] - lse_c;
145 6133500 const float b = unc[v] - lse_u;
146 6133500 lp[v] = a + guidance_scale * (a - b);
147 6133500 }
148 6423 const float lse_g = row_lse(lp, V, dummy);
149
2/2
✓ Branch 0 taken 6193114 times.
✓ Branch 1 taken 6423 times.
 6199537 for (int v = 0; v < V; ++v) lp[v] = lp[v] - lse_g;
150 6423 } else {
151 3200 const float lse = row_lse(cond, V, dummy);
152
2/2
✓ Branch 0 taken 2439474 times.
✓ Branch 1 taken 3200 times.
 2442674 for (int v = 0; v < V; ++v) lp[v] = cond[v] - lse;
153 }
154 9623 lp[mask_id] = kNegInf;
155
156 // ── confidence: max of the UNfiltered log_probs ──
157 9623 float conf = kNegInf;
158
4/4
✓ Branch 0 taken 8279479 times.
✓ Branch 1 taken 9623 times.
✓ Branch 2 taken 71432 times.
✓ Branch 3 taken 8208047 times.
 8289102 for (int v = 0; v < V; ++v) conf = (lp[v] > conf) ? lp[v] : conf;
159 9623 cf[p] = conf; // raw, for every cell — masked or already decided
160
161 // ── prediction ──
162 9623 int best = 0;
163
2/2
✓ Branch 0 taken 3208 times.
✓ Branch 1 taken 6415 times.
 9623 if (class_temperature > 0.0f) {
164 // k-th largest value of the row (ties resolved by index below).
165 3208 float* srt = sort_buf.data();
166 3208 std::copy(lp, lp + V, srt);
167 3208 std::nth_element(srt, srt + (k_keep - 1), srt + V,
168 6530793 [](float a, float b) { return a > b; });
169 3208 const float thr = srt[k_keep - 1];
170 3208 int count_gt = 0;
171
2/2
✓ Branch 0 taken 3173483 times.
✓ Branch 1 taken 3208 times.
 3176691 for (int v = 0; v < V; ++v) count_gt += (lp[v] > thr) ? 1 : 0;
172 3208 const int m = k_keep - count_gt; // equal-to-threshold slots
173
174 3208 const uint64_t base = static_cast<uint64_t>(p) * static_cast<uint64_t>(V);
175 3208 float best_s = kNegInf;
176 3208 int seen_eq = 0;
177
2/2
✓ Branch 0 taken 2847023 times.
✓ Branch 1 taken 3208 times.
 2850231 for (int v = 0; v < V; ++v) {
178 2847023 bool kept = lp[v] > thr;
179
4/4
✓ Branch 0 taken 2571230 times.
✓ Branch 1 taken 275793 times.
✓ Branch 2 taken 2568017 times.
✓ Branch 3 taken 3213 times.
 2847023 if (!kept && lp[v] == thr) {
180 3213 kept = seen_eq < m;
181 3213 ++seen_eq;
182 3213 }
183 2847023 float s = kNegInf;
184
2/2
✓ Branch 0 taken 2538850 times.
✓ Branch 1 taken 308173 times.
 2847023 if (kept) {
185 616346 const float u = hash_uniform(seed, kHashDomainMaskedDiffusionClass,
186 308173 base + static_cast<uint64_t>(v));
187 308173 s = lp[v] / class_temperature + gumbel_from_uniform(u);
188 308173 }
189
2/2
✓ Branch 0 taken 2830240 times.
✓ Branch 1 taken 16783 times.
 2847023 if (prefers(s, v, best_s, best)) { best_s = s; best = v; }
190 2847023 }
191 3208 } else {
192 6415 float best_v = kNegInf;
193
2/2
✓ Branch 0 taken 5095003 times.
✓ Branch 1 taken 6415 times.
 5101418 for (int v = 0; v < V; ++v) {
194
2/2
✓ Branch 0 taken 5047510 times.
✓ Branch 1 taken 47493 times.
 5095003 if (prefers(lp[v], v, best_v, best)) { best_v = lp[v]; best = v; }
195 5095003 }
196 }
197 9623 pr[p] = best;
198
199 // ── score ──
200 9623 float s = kNegInf;
201
2/2
✓ Branch 0 taken 3825 times.
✓ Branch 1 taken 5798 times.
 9623 if (tok[p] == mask_id) {
202 5798 s = conf - static_cast<float>(c) * layer_penalty;
203
2/2
✓ Branch 0 taken 1915 times.
✓ Branch 1 taken 3883 times.
 5798 if (position_temperature > 0.0f) {
204 7766 const float u = hash_uniform(seed, kHashDomainMaskedDiffusionPosition,
205 3883 static_cast<uint64_t>(p));
206 3883 s = s / position_temperature + gumbel_from_uniform(u);
207 3883 }
208 5798 }
209 9623 sc[p] = s;
210 9623 });
211 15 }
212
213 // ─── masked_diffusion_commit ────────────────────────────────────────────────
214
215 7 void masked_diffusion_commit(const ::brotensor::Tensor& pred,
216 const ::brotensor::Tensor& idx, int k, int step,
217 ::brotensor::Tensor& tokens,
218 ::brotensor::Tensor& unmask_step) {
219 7 const char* op = "masked_diffusion_commit";
220
1/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 8 if (pred.dtype != Dtype::INT32) fail(op, "pred must be INT32");
221
1/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 7 if (idx.dtype != Dtype::INT32) fail(op, "idx must be INT32");
222
1/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 7 if (tokens.dtype != Dtype::INT32) fail(op, "tokens must be INT32");
223
1/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 7 if (unmask_step.dtype != Dtype::INT32) fail(op, "unmask_step must be INT32");
224
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 14 if (tokens.rows != pred.rows || tokens.cols != pred.cols ||
225 7 unmask_step.rows != pred.rows || unmask_step.cols != pred.cols) {
226 ✗ fail(op, "pred, tokens and unmask_step must share one (C, T) shape");
227 }
228
1/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 7 if (k < 0) fail(op, "k must be >= 0");
229
3/4
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
 7 if (k > idx.size()) fail(op, "idx must hold at least k entries");
230
2/2
✓ Branch 0 taken 5 times.
✓ Branch 1 taken 1 time.
 6 if (k == 0) return;
231
232 5 const int n = pred.size();
233 5 const int32_t* pr = static_cast<const int32_t*>(pred.host_raw());
234 5 const int32_t* ix = static_cast<const int32_t*>(idx.host_raw());
235 5 int32_t* tok = static_cast<int32_t*>(tokens.host_raw_mut());
236 5 int32_t* us = static_cast<int32_t*>(unmask_step.host_raw_mut());
237
2/2
✓ Branch 0 taken 5 times.
✓ Branch 1 taken 456 times.
 461 for (int i = 0; i < k; ++i) {
238 456 const int32_t p = ix[i];
239
4/4
✓ Branch 0 taken 455 times.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 1 time.
✓ Branch 3 taken 454 times.
 456 if (p < 0 || p >= n) continue;
240 454 tok[p] = pr[p];
241 454 us[p] = step;
242 454 }
243 7 }
244
245 } // namespace brotensor::detail::cpu
246