GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 92.0% 1109 / 0 / 1206
Functions: 92.5% 62 / 0 / 67
Branches: 73.4% 683 / 0 / 930
src/cpu/ops_impl.cpp
Line Branch Exec Source
1 // ─── CPU op implementations ────────────────────────────────────────────────
2 //
3 // Plain FP32 host loops. Ports the 16 ops formerly in src/cpu/ops.cpp to
4 // the unified `brotensor::Tensor` type. Math is unchanged from the old
5 // host-only impls — only buffer access (host_f32 / host_f32_mut / at)
6 // differs from the legacy `.ptr()` / `.data.data()` accessors.
7
8 #include <brotensor/tensor.h>
9 #include <brotensor/detail/cpu/thread_pool.h>
10
11 #include <algorithm>
12 #include <cassert>
13 #include <cmath>
14 #include <cstddef>
15 #include <cstdint>
16 #include <cstring>
17 #include <stdexcept>
18 #include <vector>
19
20 namespace brotensor::detail::cpu {
21
22 // ─── Dense layer ───────────────────────────────────────────────────────────
23
24 881 void linear_forward(const ::brotensor::Tensor& W, const ::brotensor::Tensor& b,
25 const ::brotensor::Tensor& x, ::brotensor::Tensor& y) {
26 881 const int out = W.rows;
27 881 const int in = W.cols;
28
1/2
✓ Branch 0 taken 881 times.
✗ Branch 1 not taken.
 881 assert(x.size() == in);
29
1/2
✓ Branch 0 taken 881 times.
✗ Branch 1 not taken.
 881 assert(b.size() == out);
30
31
4/6
✓ Branch 0 taken 77 times.
✓ Branch 1 taken 804 times.
✓ Branch 2 taken 77 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 77 times.
 881 if (y.rows != out || y.cols != 1 || y.dtype != ::brotensor::Dtype::FP32) {
32 804 y.resize(out, 1, ::brotensor::Dtype::FP32);
33 804 }
34
35 881 const float* Wp = W.host_f32();
36 881 const float* xp = x.host_f32();
37 881 const float* bp = b.host_f32();
38 881 float* yp = y.host_f32_mut();
39
40
2/2
✓ Branch 0 taken 4722 times.
✓ Branch 1 taken 881 times.
 5603 for (int i = 0; i < out; ++i) {
41 4722 float acc = bp[i];
42 4722 const float* row = Wp + static_cast<std::size_t>(i) * in;
43
2/2
✓ Branch 0 taken 47268 times.
✓ Branch 1 taken 4722 times.
 51990 for (int j = 0; j < in; ++j) acc += row[j] * xp[j];
44 4722 yp[i] = acc;
45 4722 }
46 881 }
47
48 20 void linear_backward(const ::brotensor::Tensor& W, const ::brotensor::Tensor& x,
49 const ::brotensor::Tensor& dY,
50 ::brotensor::Tensor& dX, ::brotensor::Tensor& dW,
51 ::brotensor::Tensor& dB) {
52 20 const int out = W.rows;
53 20 const int in = W.cols;
54
1/2
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
 20 assert(x.size() == in);
55
1/2
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
 20 assert(dY.size() == out);
56
2/4
✗ Branch 0 not taken.
✓ Branch 1 taken 20 times.
✓ Branch 2 taken 20 times.
✗ Branch 3 not taken.
 20 assert(dW.rows == out && dW.cols == in);
57
1/2
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
 20 assert(dB.size() == out);
58
59
3/4
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 12 times.
 20 if (dX.size() != in || dX.dtype != ::brotensor::Dtype::FP32) {
60 8 dX.resize(in, 1, ::brotensor::Dtype::FP32);
61 8 }
62
63 20 const float* Wp = W.host_f32();
64 20 const float* xp = x.host_f32();
65 20 const float* dYp = dY.host_f32();
66 20 float* dXp = dX.host_f32_mut();
67 20 float* dWp = dW.host_f32_mut();
68 20 float* dBp = dB.host_f32_mut();
69
70 // dX = W^T * dY (overwrite)
71
2/2
✓ Branch 0 taken 292 times.
✓ Branch 1 taken 20 times.
 312 for (int j = 0; j < in; ++j) dXp[j] = 0.0f;
72
2/2
✓ Branch 0 taken 244 times.
✓ Branch 1 taken 20 times.
 264 for (int i = 0; i < out; ++i) {
73 244 const float g = dYp[i];
74 244 const float* row = Wp + static_cast<std::size_t>(i) * in;
75
2/2
✓ Branch 0 taken 18928 times.
✓ Branch 1 taken 244 times.
 19172 for (int j = 0; j < in; ++j) dXp[j] += row[j] * g;
76 244 }
77 // dW += dY * x^T (outer product, accumulated)
78
2/2
✓ Branch 0 taken 244 times.
✓ Branch 1 taken 20 times.
 264 for (int i = 0; i < out; ++i) {
79 244 const float g = dYp[i];
80 244 float* row = dWp + static_cast<std::size_t>(i) * in;
81
2/2
✓ Branch 0 taken 18928 times.
✓ Branch 1 taken 244 times.
 19172 for (int j = 0; j < in; ++j) row[j] += g * xp[j];
82 244 }
83 // dB += dY
84
2/2
✓ Branch 0 taken 244 times.
✓ Branch 1 taken 20 times.
 264 for (int i = 0; i < out; ++i) dBp[i] += dYp[i];
85 20 }
86
87 // ─── Elementwise activations ───────────────────────────────────────────────
88
89 127 static void ensure_match(const ::brotensor::Tensor& src, ::brotensor::Tensor& dst) {
90
4/6
✓ Branch 0 taken 126 times.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 126 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 126 times.
 127 if (dst.rows != src.rows || dst.cols != src.cols ||
91 126 dst.dtype != ::brotensor::Dtype::FP32) {
92 1 dst.resize(src.rows, src.cols, ::brotensor::Dtype::FP32);
93 1 }
94 127 }
95
96 18 void relu_forward(const ::brotensor::Tensor& x, ::brotensor::Tensor& y) {
97 18 ensure_match(x, y);
98 18 const int n = x.size();
99 18 const float* xp = x.host_f32();
100 18 float* yp = y.host_f32_mut();
101
4/4
✓ Branch 0 taken 1860 times.
✓ Branch 1 taken 18 times.
✓ Branch 2 taken 901 times.
✓ Branch 3 taken 959 times.
 1878 for (int i = 0; i < n; ++i) yp[i] = xp[i] > 0.0f ? xp[i] : 0.0f;
102 18 }
103
104 6 void relu_backward(const ::brotensor::Tensor& x, const ::brotensor::Tensor& dY,
105 ::brotensor::Tensor& dX) {
106 6 ensure_match(x, dX);
107 6 const int n = x.size();
108 6 const float* xp = x.host_f32();
109 6 const float* dYp = dY.host_f32();
110 6 float* dXp = dX.host_f32_mut();
111
4/4
✓ Branch 0 taken 1550 times.
✓ Branch 1 taken 6 times.
✓ Branch 2 taken 752 times.
✓ Branch 3 taken 798 times.
 1556 for (int i = 0; i < n; ++i) dXp[i] = xp[i] > 0.0f ? dYp[i] : 0.0f;
112 6 }
113
114 21 void tanh_forward(const ::brotensor::Tensor& x, ::brotensor::Tensor& y) {
115 21 ensure_match(x, y);
116 21 const int n = x.size();
117 21 const float* xp = x.host_f32();
118 21 float* yp = y.host_f32_mut();
119
2/2
✓ Branch 0 taken 2395 times.
✓ Branch 1 taken 21 times.
 2416 for (int i = 0; i < n; ++i) yp[i] = std::tanh(xp[i]);
120 21 }
121
122 6 void tanh_backward(const ::brotensor::Tensor& y, const ::brotensor::Tensor& dY,
123 ::brotensor::Tensor& dX) {
124 6 ensure_match(y, dX);
125 6 const int n = y.size();
126 6 const float* yp = y.host_f32();
127 6 const float* dYp = dY.host_f32();
128 6 float* dXp = dX.host_f32_mut();
129
2/2
✓ Branch 0 taken 1551 times.
✓ Branch 1 taken 6 times.
 1557 for (int i = 0; i < n; ++i) dXp[i] = dYp[i] * (1.0f - yp[i] * yp[i]);
130 6 }
131
132 20 void sigmoid_forward(const ::brotensor::Tensor& x, ::brotensor::Tensor& y) {
133 20 ensure_match(x, y);
134 20 const int n = x.size();
135 20 const float* xp = x.host_f32();
136 20 float* yp = y.host_f32_mut();
137
2/2
✓ Branch 0 taken 1883 times.
✓ Branch 1 taken 20 times.
 1903 for (int i = 0; i < n; ++i) yp[i] = 1.0f / (1.0f + std::exp(-xp[i]));
138 20 }
139
140 6 void sigmoid_backward(const ::brotensor::Tensor& y, const ::brotensor::Tensor& dY,
141 ::brotensor::Tensor& dX) {
142 6 ensure_match(y, dX);
143 6 const int n = y.size();
144 6 const float* yp = y.host_f32();
145 6 const float* dYp = dY.host_f32();
146 6 float* dXp = dX.host_f32_mut();
147
2/2
✓ Branch 0 taken 1551 times.
✓ Branch 1 taken 6 times.
 1557 for (int i = 0; i < n; ++i) dXp[i] = dYp[i] * yp[i] * (1.0f - yp[i]);
148 6 }
149
150 // ─── Softmax + xent ────────────────────────────────────────────────────────
151
152 24 void softmax_forward(const ::brotensor::Tensor& logits, ::brotensor::Tensor& probs,
153 const float* mask) {
154 24 ensure_match(logits, probs);
155 24 const int n = logits.size();
156 24 const float* lp = logits.host_f32();
157 24 float* pp = probs.host_f32_mut();
158
159 24 float m = -1e30f;
160
2/2
✓ Branch 0 taken 825 times.
✓ Branch 1 taken 24 times.
 849 for (int i = 0; i < n; ++i) {
161
4/4
✓ Branch 0 taken 469 times.
✓ Branch 1 taken 356 times.
✓ Branch 2 taken 212 times.
✓ Branch 3 taken 257 times.
 825 if (mask && mask[i] < 0.5f) continue;
162
2/2
✓ Branch 0 taken 465 times.
✓ Branch 1 taken 103 times.
 568 if (lp[i] > m) m = lp[i];
163 568 }
164 24 float s = 0.0f;
165
2/2
✓ Branch 0 taken 825 times.
✓ Branch 1 taken 24 times.
 849 for (int i = 0; i < n; ++i) {
166
4/4
✓ Branch 0 taken 469 times.
✓ Branch 1 taken 356 times.
✓ Branch 2 taken 212 times.
✓ Branch 3 taken 257 times.
 825 if (mask && mask[i] < 0.5f) { pp[i] = 0.0f; continue; }
167 568 pp[i] = std::exp(lp[i] - m);
168 568 s += pp[i];
169 568 }
170
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 const float inv = s > 0.0f ? 1.0f / s : 0.0f;
171
2/2
✓ Branch 0 taken 825 times.
✓ Branch 1 taken 24 times.
 849 for (int i = 0; i < n; ++i) pp[i] *= inv;
172 24 }
173
174 27 void softmax_rows_forward(const ::brotensor::Tensor& X, ::brotensor::Tensor& Y,
175 int rows, int cols) {
176
4/6
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 26 times.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 1 time.
 27 if (Y.rows != X.rows || Y.cols != X.cols || Y.dtype != X.dtype)
177 26 Y.resize(X.rows, X.cols, X.dtype);
178 27 const float* xp = X.host_f32();
179 27 float* yp = Y.host_f32_mut();
180
2/2
✓ Branch 0 taken 27 times.
✓ Branch 1 taken 263 times.
 290 for (int r = 0; r < rows; ++r) {
181 263 const float* lp = xp + static_cast<std::size_t>(r) * cols;
182 263 float* pp = yp + static_cast<std::size_t>(r) * cols;
183 263 float m = -1e30f;
184
4/4
✓ Branch 0 taken 21662 times.
✓ Branch 1 taken 263 times.
✓ Branch 2 taken 20564 times.
✓ Branch 3 taken 1098 times.
 21925 for (int i = 0; i < cols; ++i) if (lp[i] > m) m = lp[i];
185 263 float s = 0.0f;
186
2/2
✓ Branch 0 taken 21662 times.
✓ Branch 1 taken 263 times.
 21925 for (int i = 0; i < cols; ++i) { pp[i] = std::exp(lp[i] - m); s += pp[i]; }
187
1/2
✓ Branch 0 taken 263 times.
✗ Branch 1 not taken.
 263 const float inv = s > 0.0f ? 1.0f / s : 0.0f;
188
2/2
✓ Branch 0 taken 21662 times.
✓ Branch 1 taken 263 times.
 21925 for (int i = 0; i < cols; ++i) pp[i] *= inv;
189 263 }
190 27 }
191
192 11 void softmax_backward(const ::brotensor::Tensor& probs,
193 const ::brotensor::Tensor& dProbs,
194 ::brotensor::Tensor& dLogits) {
195 11 ensure_match(probs, dLogits);
196 11 const int n = probs.size();
197 11 const float* pp = probs.host_f32();
198 11 const float* dp = dProbs.host_f32();
199 11 float* dz = dLogits.host_f32_mut();
200 11 float dot = 0.0f;
201
2/2
✓ Branch 0 taken 638 times.
✓ Branch 1 taken 11 times.
 649 for (int i = 0; i < n; ++i) dot += dp[i] * pp[i];
202
2/2
✓ Branch 0 taken 638 times.
✓ Branch 1 taken 11 times.
 649 for (int i = 0; i < n; ++i) dz[i] = pp[i] * (dp[i] - dot);
203 11 }
204
205 10 float softmax_xent_segment(const float* lp, const float* tp,
206 float* pp, float* dz,
207 int n, const float* mask) {
208 // Stable softmax over the segment.
209 10 float m = -1e30f;
210
2/2
✓ Branch 0 taken 535 times.
✓ Branch 1 taken 10 times.
 545 for (int i = 0; i < n; ++i) {
211
4/4
✓ Branch 0 taken 196 times.
✓ Branch 1 taken 339 times.
✓ Branch 2 taken 115 times.
✓ Branch 3 taken 81 times.
 535 if (mask && mask[i] < 0.5f) continue;
212
2/2
✓ Branch 0 taken 418 times.
✓ Branch 1 taken 36 times.
 454 if (lp[i] > m) m = lp[i];
213 454 }
214 10 float s = 0.0f;
215
2/2
✓ Branch 0 taken 535 times.
✓ Branch 1 taken 10 times.
 545 for (int i = 0; i < n; ++i) {
216
4/4
✓ Branch 0 taken 196 times.
✓ Branch 1 taken 339 times.
✓ Branch 2 taken 115 times.
✓ Branch 3 taken 81 times.
 535 if (mask && mask[i] < 0.5f) { pp[i] = 0.0f; continue; }
217 454 pp[i] = std::exp(lp[i] - m);
218 454 s += pp[i];
219 454 }
220
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 const float inv = s > 0.0f ? 1.0f / s : 0.0f;
221
2/2
✓ Branch 0 taken 535 times.
✓ Branch 1 taken 10 times.
 545 for (int i = 0; i < n; ++i) pp[i] *= inv;
222
223 // xent + dLogits = (p - t).
224 10 float loss = 0.0f;
225
2/2
✓ Branch 0 taken 535 times.
✓ Branch 1 taken 10 times.
 545 for (int i = 0; i < n; ++i) {
226
4/4
✓ Branch 0 taken 196 times.
✓ Branch 1 taken 339 times.
✓ Branch 2 taken 115 times.
✓ Branch 3 taken 81 times.
 535 if (mask && mask[i] < 0.5f) { dz[i] = 0.0f; continue; }
227
2/2
✓ Branch 0 taken 9 times.
✓ Branch 1 taken 445 times.
 454 if (tp[i] > 0.0f) {
228
1/2
✓ Branch 0 taken 445 times.
✗ Branch 1 not taken.
 445 const float p = pp[i] > 1e-12f ? pp[i] : 1e-12f;
229 445 loss -= tp[i] * std::log(p);
230 445 }
231 454 dz[i] = pp[i] - tp[i];
232 454 }
233 10 return loss;
234 }
235
236 3 float softmax_xent(const ::brotensor::Tensor& logits,
237 const ::brotensor::Tensor& target,
238 ::brotensor::Tensor& probs, ::brotensor::Tensor& dLogits,
239 const float* mask) {
240 3 ensure_match(logits, probs);
241 3 ensure_match(logits, dLogits);
242 6 return softmax_xent_segment(logits.host_f32(), target.host_f32(),
243 3 probs.host_f32_mut(), dLogits.host_f32_mut(),
244 3 logits.size(), mask);
245 }
246
247 // ─── Scalar MSE ────────────────────────────────────────────────────────────
248
249 3 float mse_scalar(float pred, float target, float& dPred) {
250 3 const float d = pred - target;
251 3 dPred = d;
252 3 return 0.5f * d * d;
253 }
254
255 // ─── Misc elementwise ──────────────────────────────────────────────────────
256
257 298 void add_inplace(::brotensor::Tensor& y, const ::brotensor::Tensor& x) {
258 298 const int n = y.size();
259
1/2
✓ Branch 0 taken 298 times.
✗ Branch 1 not taken.
 298 assert(x.size() == n);
260 298 const float* xp = x.host_f32();
261 298 float* yp = y.host_f32_mut();
262
2/2
✓ Branch 0 taken 29821 times.
✓ Branch 1 taken 298 times.
 30119 for (int i = 0; i < n; ++i) yp[i] += xp[i];
263 298 }
264
265 8 void add_scalar_inplace(::brotensor::Tensor& y, float s) {
266 8 const int n = y.size();
267 8 float* yp = y.host_f32_mut();
268
2/2
✓ Branch 0 taken 2829 times.
✓ Branch 1 taken 8 times.
 2837 for (int i = 0; i < n; ++i) yp[i] += s;
269 8 }
270
271 22 void add_channel_bias_inplace(::brotensor::Tensor& y, const ::brotensor::Tensor& bias,
272 int C, int L) {
273 22 float* yp = y.host_f32_mut();
274 22 const float* bp = bias.host_f32();
275
2/2
✓ Branch 0 taken 1051 times.
✓ Branch 1 taken 22 times.
 1073 for (int c = 0; c < C; ++c) {
276 1051 const float b = bp[c];
277 1051 float* row = yp + static_cast<std::size_t>(c) * L;
278
2/2
✓ Branch 0 taken 56908 times.
✓ Branch 1 taken 1051 times.
 57959 for (int i = 0; i < L; ++i) row[i] += b;
279 1051 }
280 22 }
281
282 // ─── Xavier-uniform init ───────────────────────────────────────────────────
283
284 // splitmix64 advanced by reference; deterministic, no external dep.
285 105 static inline uint64_t splitmix(uint64_t& s) {
286 105 s += 0x9E3779B97F4A7C15ULL;
287 105 uint64_t z = s;
288 105 z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL;
289 105 z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL;
290 105 return z ^ (z >> 31);
291 }
292
293 105 static inline float u01(uint64_t& s) {
294 // top 24 bits → [0, 1)
295 105 return static_cast<float>((splitmix(s) >> 40)) / 16777216.0f;
296 }
297
298 3 void xavier_init(::brotensor::Tensor& W, uint64_t& rng_state) {
299 // fan-in = cols, fan-out = rows ; limit = sqrt(6 / (in + out))
300 3 const float limit = std::sqrt(6.0f / static_cast<float>(W.rows + W.cols));
301 3 const int n = W.size();
302 3 float* wp = W.host_f32_mut();
303
2/2
✓ Branch 0 taken 105 times.
✓ Branch 1 taken 3 times.
 108 for (int i = 0; i < n; ++i) {
304 105 const float u = u01(rng_state); // [0,1)
305 105 wp[i] = (u * 2.0f - 1.0f) * limit; // [-limit, +limit]
306 105 }
307 3 }
308
309 // ─── Optimisers ────────────────────────────────────────────────────────────
310
311 ✗ void sgd_step(::brotensor::Tensor& param, ::brotensor::Tensor& grad,
312 ::brotensor::Tensor& velocity, float lr, float momentum) {
313 ✗ const int n = param.size();
314 ✗ assert(grad.size() == n && velocity.size() == n);
315 ✗ float* p = param.host_f32_mut();
316 ✗ const float* g = grad.host_f32();
317 ✗ float* v = velocity.host_f32_mut();
318 ✗ for (int i = 0; i < n; ++i) {
319 ✗ const float vi = momentum * v[i] + g[i];
320 ✗ v[i] = vi;
321 ✗ p[i] -= lr * vi;
322 ✗ }
323 ✗ }
324
325 25 void adam_step(::brotensor::Tensor& param, const ::brotensor::Tensor& grad,
326 ::brotensor::Tensor& m, ::brotensor::Tensor& v,
327 float lr, float beta1, float beta2, float eps, int step) {
328 25 const int n = param.size();
329
3/6
✓ Branch 0 taken 25 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 25 times.
✓ Branch 4 taken 25 times.
✗ Branch 5 not taken.
 25 assert(grad.size() == n && m.size() == n && v.size() == n);
330 25 const float inv_bc1 = 1.0f / (1.0f - std::pow(beta1, static_cast<float>(step)));
331 25 const float inv_bc2 = 1.0f / (1.0f - std::pow(beta2, static_cast<float>(step)));
332 25 float* p = param.host_f32_mut();
333 25 const float* g = grad.host_f32();
334 25 float* mp = m.host_f32_mut();
335 25 float* vp = v.host_f32_mut();
336
2/2
✓ Branch 0 taken 6337 times.
✓ Branch 1 taken 25 times.
 6362 for (int i = 0; i < n; ++i) {
337 6337 const float gi = g[i];
338 6337 const float mi = beta1 * mp[i] + (1.0f - beta1) * gi;
339 6337 const float vi = beta2 * vp[i] + (1.0f - beta2) * gi * gi;
340 6337 mp[i] = mi;
341 6337 vp[i] = vi;
342 6337 const float m_hat = mi * inv_bc1;
343 6337 const float v_hat = vi * inv_bc2;
344 6337 p[i] -= lr * m_hat / (std::sqrt(v_hat) + eps);
345 6337 }
346 25 }
347
348 // ─── Elementwise scale ─────────────────────────────────────────────────────
349
350 280 void scale_inplace(::brotensor::Tensor& y, float s) {
351 280 const int n = y.size();
352 280 float* yp = y.host_f32_mut();
353
2/2
✓ Branch 0 taken 2877 times.
✓ Branch 1 taken 280 times.
 3157 for (int i = 0; i < n; ++i) yp[i] *= s;
354 280 }
355
356 // ─── LayerNorm ─────────────────────────────────────────────────────────────
357
358 3 void layernorm_forward(const ::brotensor::Tensor& x,
359 const ::brotensor::Tensor& gamma,
360 const ::brotensor::Tensor& beta,
361 ::brotensor::Tensor& y, ::brotensor::Tensor& xhat,
362 float& mean_out, float& rstd_out, float eps) {
363 3 const int n = x.size();
364
3/6
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 3 times.
 3 if (y.rows != x.rows || y.cols != x.cols ||
365 3 y.dtype != ::brotensor::Dtype::FP32) {
366 ✗ y.resize(x.rows, x.cols, ::brotensor::Dtype::FP32);
367 ✗ }
368
3/6
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 3 times.
 3 if (xhat.rows != x.rows || xhat.cols != x.cols ||
369 3 xhat.dtype != ::brotensor::Dtype::FP32) {
370 ✗ xhat.resize(x.rows, x.cols, ::brotensor::Dtype::FP32);
371 ✗ }
372
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (n == 0) { mean_out = 0.0f; rstd_out = 0.0f; return; }
373
374 3 const float* xp = x.host_f32();
375 3 const float* gp = gamma.host_f32();
376 3 const float* bp = beta.host_f32();
377 3 float* yp = y.host_f32_mut();
378 3 float* xhp = xhat.host_f32_mut();
379
380 3 float mean = 0.0f;
381
2/2
✓ Branch 0 taken 336 times.
✓ Branch 1 taken 3 times.
 339 for (int i = 0; i < n; ++i) mean += xp[i];
382 3 mean /= static_cast<float>(n);
383
384 3 float var = 0.0f;
385
2/2
✓ Branch 0 taken 336 times.
✓ Branch 1 taken 3 times.
 339 for (int i = 0; i < n; ++i) {
386 336 const float d = xp[i] - mean;
387 336 var += d * d;
388 336 }
389 3 var /= static_cast<float>(n);
390 3 const float rstd = 1.0f / std::sqrt(var + eps);
391
392
2/2
✓ Branch 0 taken 336 times.
✓ Branch 1 taken 3 times.
 339 for (int i = 0; i < n; ++i) {
393 336 const float xh = (xp[i] - mean) * rstd;
394 336 xhp[i] = xh;
395 336 yp[i] = gp[i] * xh + bp[i];
396 336 }
397 3 mean_out = mean;
398 3 rstd_out = rstd;
399 3 }
400
401 6 void layernorm_backward(const ::brotensor::Tensor& dY,
402 const ::brotensor::Tensor& xhat,
403 const ::brotensor::Tensor& gamma, float rstd,
404 ::brotensor::Tensor& dX,
405 ::brotensor::Tensor& dGamma,
406 ::brotensor::Tensor& dBeta) {
407 6 const int n = dY.size();
408
3/6
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 6 times.
 6 if (dX.rows != dY.rows || dX.cols != dY.cols ||
409 6 dX.dtype != ::brotensor::Dtype::FP32) {
410 ✗ dX.resize(dY.rows, dY.cols, ::brotensor::Dtype::FP32);
411 ✗ }
412
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (n == 0) return;
413
414 6 const float* dYp = dY.host_f32();
415 6 const float* xhp = xhat.host_f32();
416 6 const float* gp = gamma.host_f32();
417 6 float* dXp = dX.host_f32_mut();
418 6 float* dGp = dGamma.host_f32_mut();
419 6 float* dBp = dBeta.host_f32_mut();
420
421 // dGamma/dBeta accumulate; caller zeros.
422
2/2
✓ Branch 0 taken 672 times.
✓ Branch 1 taken 6 times.
 678 for (int i = 0; i < n; ++i) {
423 672 dGp[i] += dYp[i] * xhp[i];
424 672 dBp[i] += dYp[i];
425 672 }
426 // sum_dxh, sum_dxh_xhat where dxh = dY * gamma.
427 6 float sum_dxh = 0.0f, sum_dxh_xhat = 0.0f;
428
2/2
✓ Branch 0 taken 672 times.
✓ Branch 1 taken 6 times.
 678 for (int i = 0; i < n; ++i) {
429 672 const float dxh = dYp[i] * gp[i];
430 672 sum_dxh += dxh;
431 672 sum_dxh_xhat += dxh * xhp[i];
432 672 }
433 6 const float nf = static_cast<float>(n);
434 6 const float scale = rstd / nf;
435
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 672 times.
 678 for (int i = 0; i < n; ++i) {
436 672 const float dxh = dYp[i] * gp[i];
437 672 dXp[i] = scale * (nf * dxh - sum_dxh - xhp[i] * sum_dxh_xhat);
438 672 }
439 6 }
440
441 // ─── Single-head scaled dot-product self-attention ─────────────────────────
442
443 7 void attention_forward(const ::brotensor::Tensor& X,
444 const ::brotensor::Tensor& Wq, const ::brotensor::Tensor& Wk,
445 const ::brotensor::Tensor& Wv, const ::brotensor::Tensor& Wo,
446 const float* d_mask,
447 ::brotensor::Tensor& Q, ::brotensor::Tensor& K,
448 ::brotensor::Tensor& V, ::brotensor::Tensor& Attn,
449 ::brotensor::Tensor& Y_pre_Wo, ::brotensor::Tensor& O) {
450 using ::brotensor::Dtype;
451 7 const int N = X.rows;
452 7 const int D = X.cols;
453
3/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 taken 7 times.
 7 if (Q.rows != N || Q.cols != D || Q.dtype != Dtype::FP32) Q.resize(N, D, Dtype::FP32);
454
3/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 taken 7 times.
 7 if (K.rows != N || K.cols != D || K.dtype != Dtype::FP32) K.resize(N, D, Dtype::FP32);
455
3/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 taken 7 times.
 7 if (V.rows != N || V.cols != D || V.dtype != Dtype::FP32) V.resize(N, D, Dtype::FP32);
456
3/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 taken 7 times.
 7 if (Attn.rows != N || Attn.cols != N || Attn.dtype != Dtype::FP32) Attn.resize(N, N, Dtype::FP32);
457
3/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 taken 7 times.
 7 if (Y_pre_Wo.rows != N || Y_pre_Wo.cols != D || Y_pre_Wo.dtype != Dtype::FP32) Y_pre_Wo.resize(N, D, Dtype::FP32);
458
3/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 taken 7 times.
 7 if (O.rows != N || O.cols != D || O.dtype != Dtype::FP32) O.resize(N, D, Dtype::FP32);
459
2/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 7 times.
 7 if (N == 0 || D == 0) return;
460
461 7 const float* Xp = X.host_f32();
462 7 const float* Wqp = Wq.host_f32();
463 7 const float* Wkp = Wk.host_f32();
464 7 const float* Wvp = Wv.host_f32();
465 7 const float* Wop = Wo.host_f32();
466 7 float* Qp = Q.host_f32_mut();
467 7 float* Kp = K.host_f32_mut();
468 7 float* Vp = V.host_f32_mut();
469 7 float* Ap = Attn.host_f32_mut();
470 7 float* Yp = Y_pre_Wo.host_f32_mut();
471 7 float* Op = O.host_f32_mut();
472
473 // Q/K/V projections: out(i,j) = sum_k X(i,k) * W(j,k) (W stored (D,D)).
474 // Each token i exclusively owns row i of Q/K/V (W* are read-only), so the
475 // token axis parallelizes with no cross-thread writes.
476
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 73 parallel_for(static_cast<std::size_t>(N), [&](std::size_t ii) {
477 66 const int i = static_cast<int>(ii);
478 66 const float* xr = Xp + static_cast<std::size_t>(i) * D;
479
2/2
✓ Branch 0 taken 2915 times.
✓ Branch 1 taken 66 times.
 2981 for (int j = 0; j < D; ++j) {
480 2915 const float* wq = Wqp + static_cast<std::size_t>(j) * D;
481 2915 const float* wk = Wkp + static_cast<std::size_t>(j) * D;
482 2915 const float* wv = Wvp + static_cast<std::size_t>(j) * D;
483 2915 float aq = 0.0f, ak = 0.0f, av = 0.0f;
484
2/2
✓ Branch 0 taken 144801 times.
✓ Branch 1 taken 2915 times.
 147716 for (int k = 0; k < D; ++k) {
485 144801 aq += xr[k] * wq[k];
486 144801 ak += xr[k] * wk[k];
487 144801 av += xr[k] * wv[k];
488 144801 }
489 2915 const std::size_t idx = static_cast<std::size_t>(i) * D + j;
490 2915 Qp[idx] = aq;
491 2915 Kp[idx] = ak;
492 2915 Vp[idx] = av;
493 2915 }
494 66 });
495
496 7 const float inv_sqrtd = 1.0f / std::sqrt(static_cast<float>(D));
497
498 // Scores → masked row softmax → Attn. Each i owns row i of Attn
499 // exclusively (Q/K are read-only here), so this parallelizes over i.
500
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 70 parallel_for(static_cast<std::size_t>(N), [&](std::size_t ii) {
501 63 const int i = static_cast<int>(ii);
502 63 float* arow = Ap + static_cast<std::size_t>(i) * N;
503
4/4
✓ Branch 0 taken 32 times.
✓ Branch 1 taken 31 times.
✓ Branch 2 taken 18 times.
✓ Branch 3 taken 14 times.
 63 if (d_mask && d_mask[i] < 0.5f) {
504
2/2
✓ Branch 0 taken 184 times.
✓ Branch 1 taken 14 times.
 198 for (int j = 0; j < N; ++j) arow[j] = 0.0f;
505 14 return;
506 }
507 49 const float* qr = Qp + static_cast<std::size_t>(i) * D;
508 49 float m = -1e30f;
509
2/2
✓ Branch 0 taken 545 times.
✓ Branch 1 taken 49 times.
 594 for (int j = 0; j < N; ++j) {
510
4/4
✓ Branch 0 taken 200 times.
✓ Branch 1 taken 345 times.
✓ Branch 2 taken 122 times.
✓ Branch 3 taken 78 times.
 545 if (d_mask && d_mask[j] < 0.5f) continue;
511 467 const float* kr = Kp + static_cast<std::size_t>(j) * D;
512 467 float s = 0.0f;
513
2/2
✓ Branch 0 taken 15322 times.
✓ Branch 1 taken 467 times.
 15789 for (int k = 0; k < D; ++k) s += qr[k] * kr[k];
514 467 s *= inv_sqrtd;
515 467 arow[j] = s;
516
2/2
✓ Branch 0 taken 320 times.
✓ Branch 1 taken 147 times.
 467 if (s > m) m = s;
517 467 }
518 49 float sum = 0.0f;
519
2/2
✓ Branch 0 taken 541 times.
✓ Branch 1 taken 49 times.
 590 for (int j = 0; j < N; ++j) {
520
4/4
✓ Branch 0 taken 200 times.
✓ Branch 1 taken 341 times.
✓ Branch 2 taken 122 times.
✓ Branch 3 taken 78 times.
 541 if (d_mask && d_mask[j] < 0.5f) { arow[j] = 0.0f; continue; }
521 463 const float e = std::exp(arow[j] - m);
522 463 arow[j] = e;
523 463 sum += e;
524 463 }
525
1/2
✓ Branch 0 taken 49 times.
✗ Branch 1 not taken.
 49 const float inv = sum > 0.0f ? 1.0f / sum : 0.0f;
526
2/2
✓ Branch 0 taken 49 times.
✓ Branch 1 taken 560 times.
 609 for (int j = 0; j < N; ++j) arow[j] *= inv;
527 63 });
528
529 // Y_pre_Wo = Attn @ V. Each i owns row i of Y_pre_Wo exclusively (Attn/V
530 // read-only), so this parallelizes over i.
531
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 73 parallel_for(static_cast<std::size_t>(N), [&](std::size_t ii) {
532 66 const int i = static_cast<int>(ii);
533 66 const float* arow = Ap + static_cast<std::size_t>(i) * N;
534 66 float* yr = Yp + static_cast<std::size_t>(i) * D;
535
2/2
✓ Branch 0 taken 2313 times.
✓ Branch 1 taken 66 times.
 2379 for (int k = 0; k < D; ++k) {
536 2313 float acc = 0.0f;
537
2/2
✓ Branch 0 taken 24427 times.
✓ Branch 1 taken 2313 times.
 26740 for (int j = 0; j < N; ++j)
538 24427 acc += arow[j] * Vp[static_cast<std::size_t>(j) * D + k];
539 2313 yr[k] = acc;
540 2313 }
541 66 });
542
543 // O = Y @ Wo^T, zero invalid query rows. Each i owns row i of O
544 // exclusively (Y/Wo read-only), so this parallelizes over i.
545
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 70 parallel_for(static_cast<std::size_t>(N), [&](std::size_t ii) {
546 63 const int i = static_cast<int>(ii);
547 63 float* orow = Op + static_cast<std::size_t>(i) * D;
548
4/4
✓ Branch 0 taken 30 times.
✓ Branch 1 taken 33 times.
✓ Branch 2 taken 16 times.
✓ Branch 3 taken 14 times.
 63 if (d_mask && d_mask[i] < 0.5f) {
549
2/2
✓ Branch 0 taken 548 times.
✓ Branch 1 taken 14 times.
 562 for (int c = 0; c < D; ++c) orow[c] = 0.0f;
550 14 return;
551 }
552 49 const float* yr = Yp + static_cast<std::size_t>(i) * D;
553
2/2
✓ Branch 0 taken 2048 times.
✓ Branch 1 taken 49 times.
 2097 for (int c = 0; c < D; ++c) {
554 2048 const float* wr = Wop + static_cast<std::size_t>(c) * D;
555 2048 float acc = 0.0f;
556
2/2
✓ Branch 0 taken 74786 times.
✓ Branch 1 taken 2048 times.
 76834 for (int k = 0; k < D; ++k) acc += yr[k] * wr[k];
557 2048 orow[c] = acc;
558 2048 }
559 63 });
560 7 }
561
562 7 void attention_backward(const ::brotensor::Tensor& dO,
563 const ::brotensor::Tensor& X,
564 const ::brotensor::Tensor& Q, const ::brotensor::Tensor& K,
565 const ::brotensor::Tensor& V, const ::brotensor::Tensor& Attn,
566 const ::brotensor::Tensor& Y_pre_Wo,
567 const ::brotensor::Tensor& Wq, const ::brotensor::Tensor& Wk,
568 const ::brotensor::Tensor& Wv, const ::brotensor::Tensor& Wo,
569 const float* d_mask,
570 ::brotensor::Tensor& dX,
571 ::brotensor::Tensor& dWq, ::brotensor::Tensor& dWk,
572 ::brotensor::Tensor& dWv, ::brotensor::Tensor& dWo) {
573 using ::brotensor::Dtype;
574 7 const int N = X.rows;
575 7 const int D = X.cols;
576
3/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 taken 7 times.
 7 if (dX.rows != N || dX.cols != D || dX.dtype != Dtype::FP32) {
577 ✗ dX.resize(N, D, Dtype::FP32);
578 ✗ }
579
2/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 7 times.
 7 if (N == 0 || D == 0) return;
580
581 7 const float inv_sqrtd = 1.0f / std::sqrt(static_cast<float>(D));
582
583 7 const float* dOp = dO.host_f32();
584 7 const float* Xp = X.host_f32();
585 7 const float* Qp = Q.host_f32();
586 7 const float* Kp = K.host_f32();
587 7 const float* Vp = V.host_f32();
588 7 const float* Ap = Attn.host_f32();
589 7 const float* Yp = Y_pre_Wo.host_f32();
590 7 const float* Wqp = Wq.host_f32();
591 7 const float* Wkp = Wk.host_f32();
592 7 const float* Wvp = Wv.host_f32();
593 7 const float* Wop = Wo.host_f32();
594 7 float* dXp = dX.host_f32_mut();
595 7 float* dWqp = dWq.host_f32_mut();
596 7 float* dWkp = dWk.host_f32_mut();
597 7 float* dWvp = dWv.host_f32_mut();
598 7 float* dWop = dWo.host_f32_mut();
599
600 7 std::vector<float> dY(static_cast<std::size_t>(N) * D, 0.0f);
601
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 std::vector<float> dAttn(static_cast<std::size_t>(N) * N, 0.0f);
602
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 std::vector<float> dV(static_cast<std::size_t>(N) * D, 0.0f);
603
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 std::vector<float> dScores(static_cast<std::size_t>(N) * N, 0.0f);
604
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 std::vector<float> dQ(static_cast<std::size_t>(N) * D, 0.0f);
605
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 std::vector<float> dK(static_cast<std::size_t>(N) * D, 0.0f);
606
607 // dY (overwrite, zero on invalid rows). Each i owns row i of dY
608 // exclusively (Wo/dO read-only), so this parallelizes over i.
609
2/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
 73 parallel_for(static_cast<std::size_t>(N), [&](std::size_t ii) {
610 66 const int i = static_cast<int>(ii);
611
2/2
✓ Branch 0 taken 34 times.
✓ Branch 1 taken 32 times.
 66 const bool valid = !(d_mask && d_mask[i] < 0.5f);
612
2/2
✓ Branch 0 taken 66 times.
✓ Branch 1 taken 2717 times.
 2783 for (int k = 0; k < D; ++k) {
613 2717 float acc = 0.0f;
614
2/2
✓ Branch 0 taken 2084 times.
✓ Branch 1 taken 633 times.
 2717 if (valid) {
615
2/2
✓ Branch 0 taken 95868 times.
✓ Branch 1 taken 2084 times.
 97952 for (int c = 0; c < D; ++c)
616 191736 acc += Wop[static_cast<std::size_t>(c) * D + k] *
617 95868 dOp[static_cast<std::size_t>(i) * D + c];
618 2084 }
619 2717 dY[static_cast<std::size_t>(i) * D + k] = acc;
620 2717 }
621 66 });
622 // dWo (accumulate). NOT parallelized: this is a genuine reduction across
623 // the whole token axis i into the single shared (D, D) dWo buffer (every
624 // i contributes to every (c, k) cell) — unlike the per-i loops above/
625 // below, no reordering of this token-outermost loop makes different i
626 // values touch disjoint memory, so splitting it across threads would race
627 // multiple i's on the same dWrow[k] accumulator. Left single-threaded
628 // (see also mha_backward's dWo/dbo, and cross_attention_backward's dWo,
629 // which already uses a c-outer/i-inner-private-sum order that IS
630 // parallel-safe and has been parallelized over c).
631
2/2
✓ Branch 0 taken 66 times.
✓ Branch 1 taken 7 times.
 73 for (int i = 0; i < N; ++i) {
632
4/4
✓ Branch 0 taken 32 times.
✓ Branch 1 taken 34 times.
✓ Branch 2 taken 18 times.
✓ Branch 3 taken 14 times.
 66 if (d_mask && d_mask[i] < 0.5f) continue;
633 52 const float* dOi = dOp + static_cast<std::size_t>(i) * D;
634 52 const float* Yi = Yp + static_cast<std::size_t>(i) * D;
635
2/2
✓ Branch 0 taken 2320 times.
✓ Branch 1 taken 52 times.
 2372 for (int c = 0; c < D; ++c) {
636 2320 const float dov = dOi[c];
637 2320 float* dWrow = dWop + static_cast<std::size_t>(c) * D;
638
2/2
✓ Branch 0 taken 132352 times.
✓ Branch 1 taken 2320 times.
 134672 for (int k = 0; k < D; ++k) {
639 132352 dWrow[k] += dov * Yi[k];
640 132352 }
641 2320 }
642 52 }
643
644 // dAttn(i,j) = sum_k dY(i,k) * V(j,k). Each i owns row i of dAttn
645 // exclusively (dY/V read-only), so this parallelizes over i.
646
2/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
 72 parallel_for(static_cast<std::size_t>(N), [&](std::size_t ii) {
647 65 const int i = static_cast<int>(ii);
648
2/2
✓ Branch 0 taken 753 times.
✓ Branch 1 taken 65 times.
 818 for (int j = 0; j < N; ++j) {
649 753 float acc = 0.0f;
650
2/2
✓ Branch 0 taken 39020 times.
✓ Branch 1 taken 753 times.
 39773 for (int k = 0; k < D; ++k)
651 78040 acc += dY[static_cast<std::size_t>(i) * D + k] *
652 39020 Vp[static_cast<std::size_t>(j) * D + k];
653 753 dAttn[static_cast<std::size_t>(i) * N + j] = acc;
654 753 }
655 65 });
656 // dV(j,k) = sum_i Attn(i,j) * dY(i,k). NOT parallelized: token-outermost
657 // reorder (same as dWo above) means every i accumulates into every dV
658 // row j — a genuine reduction across i into shared dV, not disjoint per
659 // outer-loop index. Left single-threaded. dV starts zero-filled
660 // (std::vector value-init), so accumulating with += is equivalent to the
661 // original overwrite.
662
2/2
✓ Branch 0 taken 66 times.
✓ Branch 1 taken 7 times.
 73 for (int i = 0; i < N; ++i) {
663 66 const float* Ai = Ap + static_cast<std::size_t>(i) * N;
664 66 const float* dYi = dY.data() + static_cast<std::size_t>(i) * D;
665
2/2
✓ Branch 0 taken 772 times.
✓ Branch 1 taken 66 times.
 838 for (int j = 0; j < N; ++j) {
666 772 const float av = Ai[j];
667 772 float* dVrow = dV.data() + static_cast<std::size_t>(j) * D;
668
2/2
✓ Branch 0 taken 40000 times.
✓ Branch 1 taken 772 times.
 40772 for (int k = 0; k < D; ++k) {
669 40000 dVrow[k] += av * dYi[k];
670 40000 }
671 772 }
672 66 }
673
674 // dScores via per-row softmax backward, scaled by inv_sqrtd. Each i owns
675 // row i of dScores exclusively, so this parallelizes over i.
676
2/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
 73 parallel_for(static_cast<std::size_t>(N), [&](std::size_t ii) {
677 66 const int i = static_cast<int>(ii);
678 66 const float* prow = Ap + static_cast<std::size_t>(i) * N;
679 66 const float* dprow = dAttn.data() + static_cast<std::size_t>(i) * N;
680 66 float* drow = dScores.data() + static_cast<std::size_t>(i) * N;
681
4/4
✓ Branch 0 taken 31 times.
✓ Branch 1 taken 35 times.
✓ Branch 2 taken 17 times.
✓ Branch 3 taken 14 times.
 66 if (d_mask && d_mask[i] < 0.5f) {
682
2/2
✓ Branch 0 taken 184 times.
✓ Branch 1 taken 14 times.
 198 for (int j = 0; j < N; ++j) drow[j] = 0.0f;
683 14 return;
684 }
685 52 float dot = 0.0f;
686
2/2
✓ Branch 0 taken 506 times.
✓ Branch 1 taken 52 times.
 558 for (int j = 0; j < N; ++j) dot += dprow[j] * prow[j];
687
2/2
✓ Branch 0 taken 52 times.
✓ Branch 1 taken 537 times.
 589 for (int j = 0; j < N; ++j) {
688
4/4
✓ Branch 0 taken 193 times.
✓ Branch 1 taken 344 times.
✓ Branch 2 taken 115 times.
✓ Branch 3 taken 78 times.
 537 if (d_mask && d_mask[j] < 0.5f) drow[j] = 0.0f;
689 459 else drow[j] = prow[j] * (dprow[j] - dot) * inv_sqrtd;
690 537 }
691 66 });
692
693 // dQ(i,k) = sum_j dScores(i,j) * K(j,k). Each i owns row i of dQ
694 // exclusively (dScores/K read-only), so this parallelizes over i.
695
2/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
 70 parallel_for(static_cast<std::size_t>(N), [&](std::size_t ii) {
696 63 const int i = static_cast<int>(ii);
697
2/2
✓ Branch 0 taken 2665 times.
✓ Branch 1 taken 63 times.
 2728 for (int k = 0; k < D; ++k) {
698 2665 float acc = 0.0f;
699
2/2
✓ Branch 0 taken 30818 times.
✓ Branch 1 taken 2665 times.
 33483 for (int j = 0; j < N; ++j)
700 61636 acc += dScores[static_cast<std::size_t>(i) * N + j] *
701 30818 Kp[static_cast<std::size_t>(j) * D + k];
702 2665 dQ[static_cast<std::size_t>(i) * D + k] = acc;
703 2665 }
704 63 });
705 // dK(j,k) = sum_i dScores(i,j) * Q(i,k). NOT parallelized: same
706 // token-outermost reduction hazard as dV above — every i accumulates
707 // into every dK row j. Left single-threaded (dScores row i is
708 // contiguous over j; dK starts zero-filled, so += matches the original
709 // overwrite).
710
2/2
✓ Branch 0 taken 66 times.
✓ Branch 1 taken 7 times.
 73 for (int i = 0; i < N; ++i) {
711 66 const float* dSi = dScores.data() + static_cast<std::size_t>(i) * N;
712 66 const float* Qi = Qp + static_cast<std::size_t>(i) * D;
713
2/2
✓ Branch 0 taken 772 times.
✓ Branch 1 taken 66 times.
 838 for (int j = 0; j < N; ++j) {
714 772 const float ds = dSi[j];
715 772 float* dKrow = dK.data() + static_cast<std::size_t>(j) * D;
716
2/2
✓ Branch 0 taken 40000 times.
✓ Branch 1 taken 772 times.
 40772 for (int k = 0; k < D; ++k) {
717 40000 dKrow[k] += ds * Qi[k];
718 40000 }
719 772 }
720 66 }
721
722 // dWq/dWk/dWv accumulate: dW(j,k) += sum_i dQ(i,j) * X(i,k). NOT
723 // parallelized: token-axis (i) outermost, same reduction hazard as
724 // dWo/dV/dK above — every i accumulates into the full (D,D) dWq/dWk/dWv,
725 // no per-i disjoint slice exists (single-head attention has no per-head
726 // row partition to hide the reduction behind, unlike mha_backward's
727 // per-head dWq/dWk/dWv below). Left single-threaded.
728
2/2
✓ Branch 0 taken 7 times.
✓ Branch 1 taken 66 times.
 73 for (int i = 0; i < N; ++i) {
729 66 const float* Xi = Xp + static_cast<std::size_t>(i) * D;
730 66 const float* dQi = dQ.data() + static_cast<std::size_t>(i) * D;
731 66 const float* dKi = dK.data() + static_cast<std::size_t>(i) * D;
732 66 const float* dVi = dV.data() + static_cast<std::size_t>(i) * D;
733
2/2
✓ Branch 0 taken 2976 times.
✓ Branch 1 taken 66 times.
 3042 for (int j = 0; j < D; ++j) {
734 2976 const float qj = dQi[j];
735 2976 const float kj = dKi[j];
736 2976 const float vj = dVi[j];
737 2976 float* dWqrow = dWqp + static_cast<std::size_t>(j) * D;
738 2976 float* dWkrow = dWkp + static_cast<std::size_t>(j) * D;
739 2976 float* dWvrow = dWvp + static_cast<std::size_t>(j) * D;
740
2/2
✓ Branch 0 taken 170496 times.
✓ Branch 1 taken 2976 times.
 173472 for (int k = 0; k < D; ++k) {
741 170496 const float xv = Xi[k];
742 170496 dWqrow[k] += qj * xv;
743 170496 dWkrow[k] += kj * xv;
744 170496 dWvrow[k] += vj * xv;
745 170496 }
746 2976 }
747 66 }
748
749 // dX(i,k) = sum_j dQ(i,j)*Wq(j,k) + dK(i,j)*Wk(j,k) + dV(i,j)*Wv(j,k).
750 // Each i owns row i of dX exclusively (dQ/dK/dV/Wq/Wk/Wv read-only), so
751 // this parallelizes over i.
752
2/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
 71 parallel_for(static_cast<std::size_t>(N), [&](std::size_t ii) {
753 64 const int i = static_cast<int>(ii);
754
2/2
✓ Branch 0 taken 2884 times.
✓ Branch 1 taken 64 times.
 2948 for (int k = 0; k < D; ++k) {
755 2884 float acc = 0.0f;
756
2/2
✓ Branch 0 taken 88922 times.
✓ Branch 1 taken 2884 times.
 91806 for (int j = 0; j < D; ++j) {
757 88922 const std::size_t widx = static_cast<std::size_t>(j) * D + k;
758 177844 acc += dQ[static_cast<std::size_t>(i) * D + j] * Wqp[widx]
759 88922 + dK[static_cast<std::size_t>(i) * D + j] * Wkp[widx]
760 88922 + dV[static_cast<std::size_t>(i) * D + j] * Wvp[widx];
761 88922 }
762 2884 dXp[static_cast<std::size_t>(i) * D + k] = acc;
763 2884 }
764 64 });
765 7 }
766
767 // ─── Multi-head self-attention ─────────────────────────────────────────────
768
769 26 void mha_forward(const ::brotensor::Tensor& X,
770 const ::brotensor::Tensor& Wq, const ::brotensor::Tensor& Wk,
771 const ::brotensor::Tensor& Wv, const ::brotensor::Tensor& Wo,
772 const ::brotensor::Tensor* bq, const ::brotensor::Tensor* bk,
773 const ::brotensor::Tensor* bv, const ::brotensor::Tensor* bo,
774 const float* d_mask, int num_heads,
775 ::brotensor::Tensor& Qh, ::brotensor::Tensor& Kh,
776 ::brotensor::Tensor& Vh, ::brotensor::Tensor& Attnh,
777 ::brotensor::Tensor& Yconcat, ::brotensor::Tensor& O) {
778 using ::brotensor::Dtype;
779 26 const int K = X.rows;
780 26 const int D = X.cols;
781 26 const int H = num_heads;
782
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 26 const int dh = (H > 0) ? D / H : 0;
783
4/6
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 14 times.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 12 times.
 26 if (Qh.rows != H * K || Qh.cols != dh || Qh.dtype != Dtype::FP32) Qh.resize(H * K, dh, Dtype::FP32);
784
4/6
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 14 times.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 12 times.
 26 if (Kh.rows != H * K || Kh.cols != dh || Kh.dtype != Dtype::FP32) Kh.resize(H * K, dh, Dtype::FP32);
785
4/6
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 14 times.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 12 times.
 26 if (Vh.rows != H * K || Vh.cols != dh || Vh.dtype != Dtype::FP32) Vh.resize(H * K, dh, Dtype::FP32);
786
4/6
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 14 times.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 12 times.
 26 if (Attnh.rows != H * K || Attnh.cols != K || Attnh.dtype != Dtype::FP32) Attnh.resize(H * K, K, Dtype::FP32);
787
4/6
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 14 times.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 12 times.
 26 if (Yconcat.rows != K || Yconcat.cols != D || Yconcat.dtype != Dtype::FP32) Yconcat.resize(K, D, Dtype::FP32);
788
4/6
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 14 times.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 12 times.
 26 if (O.rows != K || O.cols != D || O.dtype != Dtype::FP32) O.resize(K, D, Dtype::FP32);
789
3/6
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 26 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 26 times.
 26 if (K == 0 || D == 0 || H == 0) return;
790
791 26 const float* Xp = X.host_f32();
792 26 const float* Wqp = Wq.host_f32();
793 26 const float* Wkp = Wk.host_f32();
794 26 const float* Wvp = Wv.host_f32();
795 26 const float* Wop = Wo.host_f32();
796 26 float* Qp = Qh.host_f32_mut();
797 26 float* Kp = Kh.host_f32_mut();
798 26 float* Vp = Vh.host_f32_mut();
799 26 float* Ap = Attnh.host_f32_mut();
800 26 float* Yp = Yconcat.host_f32_mut();
801 26 float* Op = O.host_f32_mut();
802
803
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 20 times.
 26 const float* bqp = bq ? bq->host_f32() : nullptr;
804
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 20 times.
 26 const float* bkp = bk ? bk->host_f32() : nullptr;
805
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 20 times.
 26 const float* bvp = bv ? bv->host_f32() : nullptr;
806
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 20 times.
 26 const float* bop = bo ? bo->host_f32() : nullptr;
807
808 // Per-head Q/K/V projection: out(hh,i,j) = sum_k X(i,k) * W(hh*dh+j, k)
809 // + (b ? b[hh*dh+j] : 0). Biases are flat length-D vectors indexed by
810 // the full output column (hh*dh+j). Each head hh owns rows
811 // [hh*K, hh*K+K) of Qh/Kh/Vh exclusively (X/W*/b* read-only), so the
812 // head axis parallelizes with no cross-thread writes.
813
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 130 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
814 104 const int hh = static_cast<int>(hhi);
815 104 const int row_off = hh * dh;
816
2/2
✓ Branch 0 taken 1143 times.
✓ Branch 1 taken 104 times.
 1247 for (int i = 0; i < K; ++i) {
817 1143 const float* xr = Xp + static_cast<std::size_t>(i) * D;
818 1143 const std::size_t out_row =
819 1143 (static_cast<std::size_t>(hh) * K + i) * dh;
820
2/2
✓ Branch 0 taken 12129 times.
✓ Branch 1 taken 1143 times.
 13272 for (int j = 0; j < dh; ++j) {
821 12129 const int wrow = row_off + j;
822 12129 const float* wq = Wqp + static_cast<std::size_t>(wrow) * D;
823 12129 const float* wk = Wkp + static_cast<std::size_t>(wrow) * D;
824 12129 const float* wv = Wvp + static_cast<std::size_t>(wrow) * D;
825 12129 float aq = 0.0f, ak = 0.0f, av = 0.0f;
826
2/2
✓ Branch 0 taken 449563 times.
✓ Branch 1 taken 12129 times.
 461692 for (int k = 0; k < D; ++k) {
827 449563 aq += xr[k] * wq[k];
828 449563 ak += xr[k] * wk[k];
829 449563 av += xr[k] * wv[k];
830 449563 }
831
2/2
✓ Branch 0 taken 3082 times.
✓ Branch 1 taken 9047 times.
 12129 if (bqp) aq += bqp[wrow];
832
2/2
✓ Branch 0 taken 3070 times.
✓ Branch 1 taken 9059 times.
 12129 if (bkp) ak += bkp[wrow];
833
2/2
✓ Branch 0 taken 3030 times.
✓ Branch 1 taken 9099 times.
 12129 if (bvp) av += bvp[wrow];
834 12129 Qp[out_row + j] = aq;
835 12129 Kp[out_row + j] = ak;
836 12129 Vp[out_row + j] = av;
837 12129 }
838 1143 }
839 104 });
840
841 26 const float inv_sqrtdh = 1.0f / std::sqrt(static_cast<float>(dh));
842
843 // Per-head scores → masked row softmax → Attnh. Each hh owns rows
844 // [hh*K, hh*K+K) of Attnh exclusively, so parallelizes over hh.
845
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 127 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
846 101 const int hh = static_cast<int>(hhi);
847
2/2
✓ Branch 0 taken 1092 times.
✓ Branch 1 taken 101 times.
 1193 for (int i = 0; i < K; ++i) {
848 1092 float* arow = Ap + (static_cast<std::size_t>(hh) * K + i) * K;
849
4/4
✓ Branch 0 taken 326 times.
✓ Branch 1 taken 766 times.
✓ Branch 2 taken 147 times.
✓ Branch 3 taken 179 times.
 1092 if (d_mask && d_mask[i] < 0.5f) {
850
2/2
✓ Branch 0 taken 2512 times.
✓ Branch 1 taken 179 times.
 2691 for (int j = 0; j < K; ++j) arow[j] = 0.0f;
851 179 continue;
852 }
853 913 const float* qr = Qp + (static_cast<std::size_t>(hh) * K + i) * dh;
854 913 float m = -1e30f;
855
2/2
✓ Branch 0 taken 9315 times.
✓ Branch 1 taken 913 times.
 10228 for (int j = 0; j < K; ++j) {
856
4/4
✓ Branch 0 taken 1818 times.
✓ Branch 1 taken 7497 times.
✓ Branch 2 taken 740 times.
✓ Branch 3 taken 1078 times.
 9315 if (d_mask && d_mask[j] < 0.5f) continue;
857 8237 const float* kr = Kp + (static_cast<std::size_t>(hh) * K + j) * dh;
858 8237 float s = 0.0f;
859
2/2
✓ Branch 0 taken 81765 times.
✓ Branch 1 taken 8237 times.
 90002 for (int k = 0; k < dh; ++k) s += qr[k] * kr[k];
860 8237 s *= inv_sqrtdh;
861 8237 arow[j] = s;
862
2/2
✓ Branch 0 taken 5684 times.
✓ Branch 1 taken 2553 times.
 8237 if (s > m) m = s;
863 8237 }
864 913 float sum = 0.0f;
865
2/2
✓ Branch 0 taken 9887 times.
✓ Branch 1 taken 913 times.
 10800 for (int j = 0; j < K; ++j) {
866
4/4
✓ Branch 0 taken 1879 times.
✓ Branch 1 taken 8008 times.
✓ Branch 2 taken 841 times.
✓ Branch 3 taken 1038 times.
 9887 if (d_mask && d_mask[j] < 0.5f) { arow[j] = 0.0f; continue; }
867 8849 const float e = std::exp(arow[j] - m);
868 8849 arow[j] = e;
869 8849 sum += e;
870 8849 }
871
1/2
✓ Branch 0 taken 913 times.
✗ Branch 1 not taken.
 913 const float inv = sum > 0.0f ? 1.0f / sum : 0.0f;
872
2/2
✓ Branch 0 taken 10893 times.
✓ Branch 1 taken 913 times.
 11806 for (int j = 0; j < K; ++j) arow[j] *= inv;
873 913 }
874 101 });
875
876 // Per-head Attn @ V → Yconcat(i, hh*dh+k). Each hh writes only its own
877 // column range [hh*dh, hh*dh+dh) of every row of Yconcat — different hh
878 // never touch the same bytes, so this parallelizes over hh.
879
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 127 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
880 101 const int hh = static_cast<int>(hhi);
881
2/2
✓ Branch 0 taken 1109 times.
✓ Branch 1 taken 101 times.
 1210 for (int i = 0; i < K; ++i) {
882 1109 const float* arow = Ap + (static_cast<std::size_t>(hh) * K + i) * K;
883
2/2
✓ Branch 0 taken 9703 times.
✓ Branch 1 taken 1109 times.
 10812 for (int k = 0; k < dh; ++k) {
884 9703 float acc = 0.0f;
885
2/2
✓ Branch 0 taken 101254 times.
✓ Branch 1 taken 9703 times.
 110957 for (int j = 0; j < K; ++j) {
886 101254 const float vv = Vp[(static_cast<std::size_t>(hh) * K + j) * dh + k];
887 101254 acc += arow[j] * vv;
888 101254 }
889 9703 Yp[static_cast<std::size_t>(i) * D + (hh * dh + k)] = acc;
890 9703 }
891 1109 }
892 101 });
893
894 // Output projection O = Yconcat @ Wo^T + bo, zero invalid query rows.
895 // Each i owns row i of O exclusively (Yconcat/Wo/bo read-only), so this
896 // parallelizes over i.
897
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 268 parallel_for(static_cast<std::size_t>(K), [&](std::size_t ii) {
898 242 const int i = static_cast<int>(ii);
899 242 float* orow = Op + static_cast<std::size_t>(i) * D;
900
4/4
✓ Branch 0 taken 69 times.
✓ Branch 1 taken 173 times.
✓ Branch 2 taken 29 times.
✓ Branch 3 taken 40 times.
 242 if (d_mask && d_mask[i] < 0.5f) {
901
2/2
✓ Branch 0 taken 1771 times.
✓ Branch 1 taken 40 times.
 1811 for (int c = 0; c < D; ++c) orow[c] = 0.0f;
902 40 return;
903 }
904 202 const float* yr = Yp + static_cast<std::size_t>(i) * D;
905
2/2
✓ Branch 0 taken 8768 times.
✓ Branch 1 taken 202 times.
 8970 for (int c = 0; c < D; ++c) {
906 8768 const float* wr = Wop + static_cast<std::size_t>(c) * D;
907 8768 float acc = 0.0f;
908
2/2
✓ Branch 0 taken 318015 times.
✓ Branch 1 taken 8768 times.
 326783 for (int k = 0; k < D; ++k) acc += yr[k] * wr[k];
909
2/2
✓ Branch 0 taken 6674 times.
✓ Branch 1 taken 2094 times.
 8768 if (bop) acc += bop[c];
910 8768 orow[c] = acc;
911 8768 }
912 242 });
913 26 }
914
915 19 void mha_backward(const ::brotensor::Tensor& dO,
916 const ::brotensor::Tensor& X,
917 const ::brotensor::Tensor& Qh, const ::brotensor::Tensor& Kh,
918 const ::brotensor::Tensor& Vh, const ::brotensor::Tensor& Attnh,
919 const ::brotensor::Tensor& Yconcat,
920 const ::brotensor::Tensor& Wq, const ::brotensor::Tensor& Wk,
921 const ::brotensor::Tensor& Wv, const ::brotensor::Tensor& Wo,
922 const float* d_mask, int num_heads,
923 ::brotensor::Tensor& dX,
924 ::brotensor::Tensor& dWq, ::brotensor::Tensor& dWk,
925 ::brotensor::Tensor& dWv, ::brotensor::Tensor& dWo,
926 ::brotensor::Tensor* dbq, ::brotensor::Tensor* dbk,
927 ::brotensor::Tensor* dbv, ::brotensor::Tensor* dbo) {
928 using ::brotensor::Dtype;
929 19 const int K = X.rows;
930 19 const int D = X.cols;
931 19 const int H = num_heads;
932
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 19 times.
 19 const int dh = (H > 0) ? D / H : 0;
933
3/6
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 19 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 19 times.
 19 if (dX.rows != K || dX.cols != D || dX.dtype != Dtype::FP32) {
934 ✗ dX.resize(K, D, Dtype::FP32);
935 ✗ }
936
3/6
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 19 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 19 times.
 19 if (K == 0 || D == 0 || H == 0) return;
937
938 19 const float inv_sqrtdh = 1.0f / std::sqrt(static_cast<float>(dh));
939
940 19 const float* dOp = dO.host_f32();
941 19 const float* Xp = X.host_f32();
942 19 const float* Qp = Qh.host_f32();
943 19 const float* Kp = Kh.host_f32();
944 19 const float* Vp = Vh.host_f32();
945 19 const float* Ap = Attnh.host_f32();
946 19 const float* Yp = Yconcat.host_f32();
947 19 const float* Wqp = Wq.host_f32();
948 19 const float* Wkp = Wk.host_f32();
949 19 const float* Wvp = Wv.host_f32();
950 19 const float* Wop = Wo.host_f32();
951 19 float* dXp = dX.host_f32_mut();
952 19 float* dWqp = dWq.host_f32_mut();
953 19 float* dWkp = dWk.host_f32_mut();
954 19 float* dWvp = dWv.host_f32_mut();
955 19 float* dWop = dWo.host_f32_mut();
956
957 19 std::vector<float> dYc(static_cast<std::size_t>(K) * D, 0.0f);
958
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 std::vector<float> dAttn(static_cast<std::size_t>(H) * K * K, 0.0f);
959
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 std::vector<float> dVh(static_cast<std::size_t>(H) * K * dh, 0.0f);
960
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 std::vector<float> dScores(static_cast<std::size_t>(H) * K * K, 0.0f);
961
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 std::vector<float> dQh(static_cast<std::size_t>(H) * K * dh, 0.0f);
962
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 std::vector<float> dKh(static_cast<std::size_t>(H) * K * dh, 0.0f);
963
964 // dYconcat (overwrite, zero on invalid rows). Each i owns row i of dYc
965 // exclusively (Wo/dO read-only), so this parallelizes over i.
966
2/4
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 19 times.
✗ Branch 3 not taken.
 200 parallel_for(static_cast<std::size_t>(K), [&](std::size_t ii) {
967 181 const int i = static_cast<int>(ii);
968
2/2
✓ Branch 0 taken 122 times.
✓ Branch 1 taken 59 times.
 181 const bool valid = !(d_mask && d_mask[i] < 0.5f);
969
2/2
✓ Branch 0 taken 181 times.
✓ Branch 1 taken 7749 times.
 7930 for (int k = 0; k < D; ++k) {
970 7749 float acc = 0.0f;
971
2/2
✓ Branch 0 taken 6313 times.
✓ Branch 1 taken 1436 times.
 7749 if (valid) {
972
2/2
✓ Branch 0 taken 243146 times.
✓ Branch 1 taken 6313 times.
 249459 for (int c = 0; c < D; ++c)
973 486292 acc += Wop[static_cast<std::size_t>(c) * D + k] *
974 243146 dOp[static_cast<std::size_t>(i) * D + c];
975 6313 }
976 7749 dYc[static_cast<std::size_t>(i) * D + k] = acc;
977 7749 }
978 181 });
979 // dWo (accumulate). NOT parallelized: token axis (i) outermost, same
980 // GEMM-reorder as attention_backward's dWo — every i accumulates into
981 // the single shared (D, D) dWo (no per-head row partition here, unlike
982 // dWq/dWk/dWv below, since Wo mixes all heads' columns together), so
983 // splitting this loop across threads would race multiple i's on the same
984 // dWrow[k] accumulator. Left single-threaded.
985
2/2
✓ Branch 0 taken 186 times.
✓ Branch 1 taken 19 times.
 205 for (int i = 0; i < K; ++i) {
986
4/4
✓ Branch 0 taken 64 times.
✓ Branch 1 taken 122 times.
✓ Branch 2 taken 26 times.
✓ Branch 3 taken 38 times.
 186 if (d_mask && d_mask[i] < 0.5f) continue;
987 148 const float* dOi = dOp + static_cast<std::size_t>(i) * D;
988 148 const float* Yi = Yp + static_cast<std::size_t>(i) * D;
989
2/2
✓ Branch 0 taken 7360 times.
✓ Branch 1 taken 148 times.
 7508 for (int c = 0; c < D; ++c) {
990 7360 const float dov = dOi[c];
991 7360 float* dWrow = dWop + static_cast<std::size_t>(c) * D;
992
2/2
✓ Branch 0 taken 411136 times.
✓ Branch 1 taken 7360 times.
 418496 for (int k = 0; k < D; ++k) {
993 411136 dWrow[k] += dov * Yi[k];
994 411136 }
995 7360 }
996 148 }
997
998 // dbo[c] += sum over valid rows of dO[i,c] (caller zeros). NOT
999 // parallelized: same full-reduction-over-i hazard as dWo above (every i
1000 // accumulates into the same D-length dbo). Left single-threaded.
1001
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 15 times.
 19 if (dbo) {
1002
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 float* dbop = dbo->host_f32_mut();
1003
2/2
✓ Branch 0 taken 40 times.
✓ Branch 1 taken 4 times.
 44 for (int i = 0; i < K; ++i) {
1004
4/4
✓ Branch 0 taken 20 times.
✓ Branch 1 taken 20 times.
✓ Branch 2 taken 8 times.
✓ Branch 3 taken 12 times.
 40 if (d_mask && d_mask[i] < 0.5f) continue;
1005 28 const float* dOi = dOp + static_cast<std::size_t>(i) * D;
1006
2/2
✓ Branch 0 taken 1632 times.
✓ Branch 1 taken 28 times.
 1660 for (int c = 0; c < D; ++c) {
1007 1632 dbop[c] += dOi[c];
1008 1632 }
1009 28 }
1010 4 }
1011
1012 // Per-head dAttn and dVh. Each hh owns rows [hh*K, hh*K+K) of dAttn and
1013 // dVh exclusively; the inner reduction over i (for dVh) sums only into
1014 // this head's own private slice, never touched by another hh, so the
1015 // whole per-head block parallelizes over hh.
1016
2/4
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 19 times.
✗ Branch 3 not taken.
 89 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
1017 70 const int hh = static_cast<int>(hhi);
1018
2/2
✓ Branch 0 taken 768 times.
✓ Branch 1 taken 70 times.
 838 for (int i = 0; i < K; ++i) {
1019
2/2
✓ Branch 0 taken 8722 times.
✓ Branch 1 taken 768 times.
 9490 for (int j = 0; j < K; ++j) {
1020 8722 float acc = 0.0f;
1021
2/2
✓ Branch 0 taken 88125 times.
✓ Branch 1 taken 8722 times.
 96847 for (int k = 0; k < dh; ++k) {
1022 88125 const float dy = dYc[static_cast<std::size_t>(i) * D + (hh * dh + k)];
1023 88125 const float vv = Vp[(static_cast<std::size_t>(hh) * K + j) * dh + k];
1024 88125 acc += dy * vv;
1025 88125 }
1026 8722 dAttn[(static_cast<std::size_t>(hh) * K + i) * K + j] = acc;
1027 8722 }
1028 768 }
1029 // dVh: token axis (i) outermost. Ap's row i is contiguous over j;
1030 // dYc's per-head segment is contiguous over k. dVh starts
1031 // zero-filled, so += matches the original overwrite.
1032
2/2
✓ Branch 0 taken 779 times.
✓ Branch 1 taken 70 times.
 849 for (int i = 0; i < K; ++i) {
1033 779 const float* Arow = Ap + (static_cast<std::size_t>(hh) * K + i) * K;
1034 779 const float* dYseg = dYc.data() + static_cast<std::size_t>(i) * D + hh * dh;
1035
2/2
✓ Branch 0 taken 9934 times.
✓ Branch 1 taken 779 times.
 10713 for (int j = 0; j < K; ++j) {
1036 9934 const float av = Arow[j];
1037 9934 float* dVrow = dVh.data() + (static_cast<std::size_t>(hh) * K + j) * dh;
1038
2/2
✓ Branch 0 taken 127064 times.
✓ Branch 1 taken 9934 times.
 136998 for (int k = 0; k < dh; ++k) {
1039 127064 dVrow[k] += av * dYseg[k];
1040 127064 }
1041 9934 }
1042 779 }
1043 70 });
1044
1045 // Per-head softmax backward → dScores. Each hh owns rows
1046 // [hh*K, hh*K+K) of dScores exclusively, so parallelizes over hh.
1047
2/4
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 19 times.
✗ Branch 3 not taken.
 87 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
1048 68 const int hh = static_cast<int>(hhi);
1049
2/2
✓ Branch 0 taken 720 times.
✓ Branch 1 taken 68 times.
 788 for (int i = 0; i < K; ++i) {
1050 720 const float* prow = Ap + (static_cast<std::size_t>(hh) * K + i) * K;
1051 720 const float* dprow = dAttn.data() + (static_cast<std::size_t>(hh) * K + i) * K;
1052 720 float* drow = dScores.data() + (static_cast<std::size_t>(hh) * K + i) * K;
1053
4/4
✓ Branch 0 taken 266 times.
✓ Branch 1 taken 454 times.
✓ Branch 2 taken 116 times.
✓ Branch 3 taken 150 times.
 720 if (d_mask && d_mask[i] < 0.5f) {
1054
2/2
✓ Branch 0 taken 1967 times.
✓ Branch 1 taken 150 times.
 2117 for (int j = 0; j < K; ++j) drow[j] = 0.0f;
1055 150 continue;
1056 }
1057 570 float dot = 0.0f;
1058
2/2
✓ Branch 0 taken 6423 times.
✓ Branch 1 taken 570 times.
 6993 for (int j = 0; j < K; ++j) dot += dprow[j] * prow[j];
1059
2/2
✓ Branch 0 taken 6575 times.
✓ Branch 1 taken 570 times.
 7145 for (int j = 0; j < K; ++j) {
1060
4/4
✓ Branch 0 taken 1539 times.
✓ Branch 1 taken 5036 times.
✓ Branch 2 taken 581 times.
✓ Branch 3 taken 958 times.
 6575 if (d_mask && d_mask[j] < 0.5f) drow[j] = 0.0f;
1061 5617 else drow[j] = prow[j] * (dprow[j] - dot) * inv_sqrtdh;
1062 6575 }
1063 570 }
1064 68 });
1065
1066 // Per-head dQh and dKh. Each hh owns rows [hh*K, hh*K+K) of dQh/dKh
1067 // exclusively (the inner dKh reduction over i sums only into this
1068 // head's own private slice), so parallelizes over hh.
1069
2/4
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 19 times.
✗ Branch 3 not taken.
 89 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
1070 70 const int hh = static_cast<int>(hhi);
1071
2/2
✓ Branch 0 taken 746 times.
✓ Branch 1 taken 70 times.
 816 for (int i = 0; i < K; ++i) {
1072
2/2
✓ Branch 0 taken 8180 times.
✓ Branch 1 taken 746 times.
 8926 for (int k = 0; k < dh; ++k) {
1073 8180 float acc = 0.0f;
1074
2/2
✓ Branch 0 taken 71553 times.
✓ Branch 1 taken 8180 times.
 79733 for (int j = 0; j < K; ++j) {
1075 71553 const float ds = dScores[(static_cast<std::size_t>(hh) * K + i) * K + j];
1076 71553 const float kk = Kp[(static_cast<std::size_t>(hh) * K + j) * dh + k];
1077 71553 acc += ds * kk;
1078 71553 }
1079 8180 dQh[(static_cast<std::size_t>(hh) * K + i) * dh + k] = acc;
1080 8180 }
1081 746 }
1082 // dKh: token axis (i) outermost, same reorder as dVh above.
1083
2/2
✓ Branch 0 taken 772 times.
✓ Branch 1 taken 70 times.
 842 for (int i = 0; i < K; ++i) {
1084 772 const float* dSrow = dScores.data() + (static_cast<std::size_t>(hh) * K + i) * K;
1085 772 const float* Qseg = Qp + (static_cast<std::size_t>(hh) * K + i) * dh;
1086
2/2
✓ Branch 0 taken 8713 times.
✓ Branch 1 taken 772 times.
 9485 for (int j = 0; j < K; ++j) {
1087 8713 const float ds = dSrow[j];
1088 8713 float* dKrow = dKh.data() + (static_cast<std::size_t>(hh) * K + j) * dh;
1089
2/2
✓ Branch 0 taken 117099 times.
✓ Branch 1 taken 8713 times.
 125812 for (int k = 0; k < dh; ++k) {
1090 117099 dKrow[k] += ds * Qseg[k];
1091 117099 }
1092 8713 }
1093 772 }
1094 70 });
1095
1096 // dbq/dbk/dbv accumulate: db[hh*dh+j] += sum_i dQh/dKh/dVh(hh,i,j).
1097 // Caller zeros. Skipped when the matching pointer is null. Each hh
1098 // writes only its own [hh*dh, hh*dh+dh) slice of dbq/dbk/dbv — disjoint
1099 // across heads (D = H*dh total) — so parallelizes over hh even though
1100 // each head's own reduction over i is accumulated serially inside.
1101
4/6
✓ Branch 0 taken 15 times.
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 15 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 15 times.
 19 if (dbq || dbk || dbv) {
1102
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 float* dbqp = dbq ? dbq->host_f32_mut() : nullptr;
1103
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 float* dbkp = dbk ? dbk->host_f32_mut() : nullptr;
1104
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 float* dbvp = dbv ? dbv->host_f32_mut() : nullptr;
1105
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 14 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
1106 10 const int hh = static_cast<int>(hhi);
1107
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 float* dbq_row = dbqp ? dbqp + hh * dh : nullptr;
1108
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 float* dbk_row = dbkp ? dbkp + hh * dh : nullptr;
1109
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 float* dbv_row = dbvp ? dbvp + hh * dh : nullptr;
1110
2/2
✓ Branch 0 taken 130 times.
✓ Branch 1 taken 8 times.
 138 for (int i = 0; i < K; ++i) {
1111 130 const std::size_t base = (static_cast<std::size_t>(hh) * K + i) * dh;
1112
2/2
✓ Branch 0 taken 1685 times.
✓ Branch 1 taken 128 times.
 1813 for (int j = 0; j < dh; ++j) {
1113
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1685 times.
 1685 if (dbq_row) dbq_row[j] += dQh[base + j];
1114
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 1684 times.
 1685 if (dbk_row) dbk_row[j] += dKh[base + j];
1115
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 1684 times.
 1685 if (dbv_row) dbv_row[j] += dVh[base + j];
1116 1685 }
1117 128 }
1118 8 });
1119 4 }
1120
1121 // dWq/dWk/dWv accumulate: dW(hh*dh+j, k) += sum_i dQh(hh,i,j) * X(i,k).
1122 // Each hh writes only its own rows [hh*dh, hh*dh+dh) of dWq/dWk/dWv —
1123 // disjoint across heads (D = H*dh total rows) — so this parallelizes
1124 // over hh even though the reduction over i is accumulated serially
1125 // inside each head's own row range (unlike attention_backward's
1126 // single-head dWq/dWk/dWv, which has no such per-head partition and is
1127 // left single-threaded there).
1128
2/4
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 19 times.
✗ Branch 3 not taken.
 87 parallel_for(static_cast<std::size_t>(H), [&](std::size_t hhi) {
1129 68 const int hh = static_cast<int>(hhi);
1130
2/2
✓ Branch 0 taken 777 times.
✓ Branch 1 taken 68 times.
 845 for (int i = 0; i < K; ++i) {
1131 777 const float* Xi = Xp + static_cast<std::size_t>(i) * D;
1132 777 const std::size_t base = (static_cast<std::size_t>(hh) * K + i) * dh;
1133
2/2
✓ Branch 0 taken 9123 times.
✓ Branch 1 taken 777 times.
 9900 for (int j = 0; j < dh; ++j) {
1134 9123 const int wrow = hh * dh + j;
1135 9123 const float qv = dQh[base + j];
1136 9123 const float kv = dKh[base + j];
1137 9123 const float vv = dVh[base + j];
1138 9123 float* dWqrow = dWqp + static_cast<std::size_t>(wrow) * D;
1139 9123 float* dWkrow = dWkp + static_cast<std::size_t>(wrow) * D;
1140 9123 float* dWvrow = dWvp + static_cast<std::size_t>(wrow) * D;
1141
2/2
✓ Branch 0 taken 326528 times.
✓ Branch 1 taken 9123 times.
 335651 for (int k = 0; k < D; ++k) {
1142 326528 const float xv = Xi[k];
1143 326528 dWqrow[k] += qv * xv;
1144 326528 dWkrow[k] += kv * xv;
1145 326528 dWvrow[k] += vv * xv;
1146 326528 }
1147 9123 }
1148 777 }
1149 68 });
1150
1151 // dX(i,k) = sum over heads, j of dQh*Wq + dKh*Wk + dVh*Wv at (hh*dh+j, k).
1152 // Each i owns row i of dX exclusively (dQh/dKh/dVh/Wq/Wk/Wv read-only),
1153 // so this parallelizes over i.
1154
2/4
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 19 times.
✗ Branch 3 not taken.
 203 parallel_for(static_cast<std::size_t>(K), [&](std::size_t ii) {
1155 184 const int i = static_cast<int>(ii);
1156
2/2
✓ Branch 0 taken 9233 times.
✓ Branch 1 taken 184 times.
 9417 for (int k = 0; k < D; ++k) {
1157 9233 float acc = 0.0f;
1158
2/2
✓ Branch 0 taken 35689 times.
✓ Branch 1 taken 9233 times.
 44922 for (int hh = 0; hh < H; ++hh) {
1159
2/2
✓ Branch 0 taken 249495 times.
✓ Branch 1 taken 35689 times.
 285184 for (int j = 0; j < dh; ++j) {
1160 249495 const int wrow = hh * dh + j;
1161 249495 const std::size_t widx = static_cast<std::size_t>(wrow) * D + k;
1162 249495 const float gq = dQh[(static_cast<std::size_t>(hh) * K + i) * dh + j];
1163 249495 const float gk = dKh[(static_cast<std::size_t>(hh) * K + i) * dh + j];
1164 249495 const float gv = dVh[(static_cast<std::size_t>(hh) * K + i) * dh + j];
1165 249495 acc += gq * Wqp[widx] + gk * Wkp[widx] + gv * Wvp[widx];
1166 249495 }
1167 35689 }
1168 9233 dXp[static_cast<std::size_t>(i) * D + k] = acc;
1169 9233 }
1170 184 });
1171 19 }
1172
1173 // ─── Concat / split ────────────────────────────────────────────────────────
1174
1175 ✗ void concat_rows(const std::vector<const ::brotensor::Tensor*>& parts,
1176 ::brotensor::Tensor& out) {
1177 ✗ int total = 0;
1178 ✗ for (const auto* p : parts) {
1179 ✗ if (p) total += p->size();
1180 }
1181 ✗ if (out.rows != total || out.cols != 1 ||
1182 ✗ out.dtype != ::brotensor::Dtype::FP32) {
1183 ✗ out.resize(total, 1, ::brotensor::Dtype::FP32);
1184 ✗ }
1185 ✗ if (total == 0) return;
1186 ✗ float* op = out.host_f32_mut();
1187 ✗ int off = 0;
1188 ✗ for (const auto* p : parts) {
1189 ✗ if (!p) continue;
1190 ✗ const int n = p->size();
1191 ✗ if (n == 0) continue;
1192 ✗ const float* pp = p->host_f32();
1193 ✗ for (int i = 0; i < n; ++i) op[off + i] = pp[i];
1194 ✗ off += n;
1195 }
1196 ✗ }
1197
1198 ✗ void split_rows(const ::brotensor::Tensor& in,
1199 const std::vector<::brotensor::Tensor*>& parts) {
1200 ✗ const float* ip = in.host_f32();
1201 ✗ int off = 0;
1202 ✗ for (auto* p : parts) {
1203 ✗ if (!p) continue;
1204 ✗ const int n = p->size();
1205 ✗ if (n == 0) continue;
1206 ✗ float* pp = p->host_f32_mut();
1207 ✗ for (int i = 0; i < n; ++i) pp[i] = ip[off + i];
1208 ✗ off += n;
1209 }
1210 ✗ }
1211
1212 // ─── Batched elementwise / linear ──────────────────────────────────────────
1213
1214 ✗ void add_inplace_batched(::brotensor::Tensor& Y_BD,
1215 const ::brotensor::Tensor& X_BD) {
1216 ✗ const int n = Y_BD.size();
1217 ✗ assert(X_BD.size() == n);
1218 ✗ float* yp = Y_BD.host_f32_mut();
1219 ✗ const float* xp = X_BD.host_f32();
1220 ✗ for (int i = 0; i < n; ++i) yp[i] += xp[i];
1221 ✗ }
1222
1223 6 void linear_forward_batched(const ::brotensor::Tensor& W,
1224 const ::brotensor::Tensor& bias,
1225 const ::brotensor::Tensor& X_BD,
1226 ::brotensor::Tensor& Y_BD) {
1227 6 const int out_dim = W.rows;
1228 6 const int in_dim = W.cols;
1229 6 const int B = X_BD.rows;
1230
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 6 if (Y_BD.rows != B || Y_BD.cols != out_dim ||
1231 ✗ Y_BD.dtype != ::brotensor::Dtype::FP32) {
1232 6 Y_BD.resize(B, out_dim, ::brotensor::Dtype::FP32);
1233 6 }
1234
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 6 times.
 6 if (B == 0 || out_dim == 0) return;
1235
1236 6 const float* bp = bias.host_f32();
1237 6 const float* Xp = X_BD.host_f32();
1238 6 float* Yp = Y_BD.host_f32_mut();
1239
1240 // 16-bit weights (FP32 activations): widen per element. Correctness path —
1241 // the performance-relevant 16-bit-weight decode runs on the GPU backends.
1242
3/4
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 6 if (W.dtype == ::brotensor::Dtype::FP16 ||
1243 4 W.dtype == ::brotensor::Dtype::BF16) {
1244 6 const bool bf = (W.dtype == ::brotensor::Dtype::BF16);
1245 6 const std::uint16_t* Wp = static_cast<const std::uint16_t*>(W.data);
1246
2/2
✓ Branch 0 taken 136 times.
✓ Branch 1 taken 6 times.
 142 for (int b = 0; b < B; ++b) {
1247 136 const float* xr = Xp + static_cast<std::size_t>(b) * in_dim;
1248 136 float* yr = Yp + static_cast<std::size_t>(b) * out_dim;
1249
2/2
✓ Branch 0 taken 12582 times.
✓ Branch 1 taken 136 times.
 12718 for (int i = 0; i < out_dim; ++i) {
1250 12582 const std::uint16_t* wr = Wp + static_cast<std::size_t>(i) * in_dim;
1251 12582 float acc = bp[i];
1252
2/2
✓ Branch 0 taken 1674330 times.
✓ Branch 1 taken 12582 times.
 1686912 for (int j = 0; j < in_dim; ++j) {
1253
2/2
✓ Branch 0 taken 838746 times.
✓ Branch 1 taken 835584 times.
 1674330 const float w = bf ? ::brotensor::bf16_bits_to_fp32(wr[j])
1254 835584 : ::brotensor::fp16_bits_to_fp32(wr[j]);
1255 1674330 acc += w * xr[j];
1256 1674330 }
1257 12582 yr[i] = acc;
1258 12582 }
1259 136 }
1260 6 return;
1261 }
1262
1263 ✗ const float* Wp = W.host_f32();
1264 ✗ for (int b = 0; b < B; ++b) {
1265 ✗ const float* xr = Xp + static_cast<std::size_t>(b) * in_dim;
1266 ✗ float* yr = Yp + static_cast<std::size_t>(b) * out_dim;
1267 ✗ for (int i = 0; i < out_dim; ++i) {
1268 ✗ const float* wr = Wp + static_cast<std::size_t>(i) * in_dim;
1269 ✗ float acc = bp[i];
1270 ✗ for (int j = 0; j < in_dim; ++j) acc += wr[j] * xr[j];
1271 ✗ yr[i] = acc;
1272 ✗ }
1273 ✗ }
1274 6 }
1275
1276 3 void linear_backward_batched(const ::brotensor::Tensor& W,
1277 const ::brotensor::Tensor& X_BD,
1278 const ::brotensor::Tensor& dY_BD,
1279 ::brotensor::Tensor& dX_BD,
1280 ::brotensor::Tensor& dW,
1281 ::brotensor::Tensor& dB) {
1282 3 const int out_dim = W.rows;
1283 3 const int in_dim = W.cols;
1284 3 const int B = X_BD.rows;
1285
3/6
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 3 times.
 3 if (dX_BD.rows != B || dX_BD.cols != in_dim ||
1286 3 dX_BD.dtype != ::brotensor::Dtype::FP32) {
1287 ✗ dX_BD.resize(B, in_dim, ::brotensor::Dtype::FP32);
1288 ✗ }
1289
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (B == 0) return;
1290
1291 3 const float* Wp = W.host_f32();
1292 3 const float* Xp = X_BD.host_f32();
1293 3 const float* dYp = dY_BD.host_f32();
1294 3 float* dXp = dX_BD.host_f32_mut();
1295 3 float* dWp = dW.host_f32_mut();
1296 3 float* dBp = dB.host_f32_mut();
1297
1298 // dX[b] = W^T * dY[b] (overwrite).
1299
2/2
✓ Branch 0 taken 17 times.
✓ Branch 1 taken 3 times.
 20 for (int b = 0; b < B; ++b) {
1300 17 const float* dyr = dYp + static_cast<std::size_t>(b) * out_dim;
1301 17 float* dxr = dXp + static_cast<std::size_t>(b) * in_dim;
1302
2/2
✓ Branch 0 taken 552 times.
✓ Branch 1 taken 17 times.
 569 for (int j = 0; j < in_dim; ++j) dxr[j] = 0.0f;
1303
2/2
✓ Branch 0 taken 328 times.
✓ Branch 1 taken 17 times.
 345 for (int i = 0; i < out_dim; ++i) {
1304 328 const float g = dyr[i];
1305 328 const float* wr = Wp + static_cast<std::size_t>(i) * in_dim;
1306
2/2
✓ Branch 0 taken 16952 times.
✓ Branch 1 taken 328 times.
 17280 for (int j = 0; j < in_dim; ++j) dxr[j] += wr[j] * g;
1307 328 }
1308 17 }
1309 // dW += sum_b dY[b] * X[b]^T; dB += sum_b dY[b] (accumulate).
1310
2/2
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 17 times.
 20 for (int b = 0; b < B; ++b) {
1311 17 const float* dyr = dYp + static_cast<std::size_t>(b) * out_dim;
1312 17 const float* xr = Xp + static_cast<std::size_t>(b) * in_dim;
1313
2/2
✓ Branch 0 taken 328 times.
✓ Branch 1 taken 17 times.
 345 for (int i = 0; i < out_dim; ++i) {
1314 328 const float g = dyr[i];
1315 328 float* dwr = dWp + static_cast<std::size_t>(i) * in_dim;
1316
2/2
✓ Branch 0 taken 16952 times.
✓ Branch 1 taken 328 times.
 17280 for (int j = 0; j < in_dim; ++j) dwr[j] += g * xr[j];
1317 328 dBp[i] += g;
1318 328 }
1319 17 }
1320 3 }
1321
1322 ✗ void relu_forward_batched(const ::brotensor::Tensor& X_BD,
1323 ::brotensor::Tensor& Y_BD) {
1324 ✗ ensure_match(X_BD, Y_BD);
1325 ✗ const int n = X_BD.size();
1326 ✗ const float* xp = X_BD.host_f32();
1327 ✗ float* yp = Y_BD.host_f32_mut();
1328 ✗ for (int i = 0; i < n; ++i) yp[i] = xp[i] > 0.0f ? xp[i] : 0.0f;
1329 ✗ }
1330
1331 3 void relu_backward_batched(const ::brotensor::Tensor& X_BD,
1332 const ::brotensor::Tensor& dY_BD,
1333 ::brotensor::Tensor& dX_BD) {
1334 3 ensure_match(X_BD, dX_BD);
1335 3 const int n = X_BD.size();
1336 3 const float* xp = X_BD.host_f32();
1337 3 const float* dYp = dY_BD.host_f32();
1338 3 float* dXp = dX_BD.host_f32_mut();
1339
4/4
✓ Branch 0 taken 2592 times.
✓ Branch 1 taken 3 times.
✓ Branch 2 taken 1277 times.
✓ Branch 3 taken 1315 times.
 2595 for (int i = 0; i < n; ++i) dXp[i] = xp[i] > 0.0f ? dYp[i] : 0.0f;
1340 3 }
1341
1342 3 void tanh_forward_batched(const ::brotensor::Tensor& X_BD,
1343 ::brotensor::Tensor& Y_BD) {
1344 3 ensure_match(X_BD, Y_BD);
1345 3 const int n = X_BD.size();
1346 3 const float* xp = X_BD.host_f32();
1347 3 float* yp = Y_BD.host_f32_mut();
1348
2/2
✓ Branch 0 taken 2592 times.
✓ Branch 1 taken 3 times.
 2595 for (int i = 0; i < n; ++i) yp[i] = std::tanh(xp[i]);
1349 3 }
1350
1351 3 void tanh_backward_batched(const ::brotensor::Tensor& Y_BD,
1352 const ::brotensor::Tensor& dY_BD,
1353 ::brotensor::Tensor& dX_BD) {
1354 3 ensure_match(Y_BD, dX_BD);
1355 3 const int n = Y_BD.size();
1356 3 const float* yp = Y_BD.host_f32();
1357 3 const float* dYp = dY_BD.host_f32();
1358 3 float* dXp = dX_BD.host_f32_mut();
1359
2/2
✓ Branch 0 taken 2592 times.
✓ Branch 1 taken 3 times.
 2595 for (int i = 0; i < n; ++i) dXp[i] = dYp[i] * (1.0f - yp[i] * yp[i]);
1360 3 }
1361
1362 // ─── Batched per-sample losses ─────────────────────────────────────────────
1363
1364 3 void mse_vec_per_sample(const ::brotensor::Tensor& pred,
1365 const ::brotensor::Tensor& target,
1366 ::brotensor::Tensor& dPred,
1367 ::brotensor::Tensor& loss_per_sample) {
1368 3 const int B = pred.size();
1369
3/6
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 3 times.
 3 if (dPred.rows != pred.rows || dPred.cols != pred.cols ||
1370 3 dPred.dtype != ::brotensor::Dtype::FP32) {
1371 ✗ dPred.resize(pred.rows, pred.cols, ::brotensor::Dtype::FP32);
1372 ✗ }
1373
3/6
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 3 times.
 3 if (loss_per_sample.rows != B || loss_per_sample.cols != 1 ||
1374 3 loss_per_sample.dtype != ::brotensor::Dtype::FP32) {
1375 ✗ loss_per_sample.resize(B, 1, ::brotensor::Dtype::FP32);
1376 ✗ }
1377
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (B == 0) return;
1378 3 const float* pp = pred.host_f32();
1379 3 const float* tp = target.host_f32();
1380 3 float* dp = dPred.host_f32_mut();
1381 3 float* lp = loss_per_sample.host_f32_mut();
1382
2/2
✓ Branch 0 taken 73 times.
✓ Branch 1 taken 3 times.
 76 for (int b = 0; b < B; ++b) {
1383 73 const float d = pp[b] - tp[b];
1384 73 dp[b] = d;
1385 73 lp[b] = 0.5f * d * d;
1386 73 }
1387 3 }
1388
1389 6 void softmax_xent_fused_batched(const ::brotensor::Tensor& logits_BL,
1390 const ::brotensor::Tensor& target_BL,
1391 const float* d_mask_BL,
1392 const int* d_head_offsets,
1393 int n_heads,
1394 ::brotensor::Tensor& probs_BL,
1395 ::brotensor::Tensor& dLogits_BL,
1396 ::brotensor::Tensor& loss_per_sample) {
1397 6 const int B = logits_BL.rows;
1398 6 const int n_act = logits_BL.cols;
1399
3/6
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 6 times.
 6 if (probs_BL.rows != B || probs_BL.cols != n_act ||
1400 6 probs_BL.dtype != ::brotensor::Dtype::FP32) {
1401 ✗ probs_BL.resize(B, n_act, ::brotensor::Dtype::FP32);
1402 ✗ }
1403
3/6
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 6 times.
 6 if (dLogits_BL.rows != B || dLogits_BL.cols != n_act ||
1404 6 dLogits_BL.dtype != ::brotensor::Dtype::FP32) {
1405 ✗ dLogits_BL.resize(B, n_act, ::brotensor::Dtype::FP32);
1406 ✗ }
1407
3/6
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 6 times.
 6 if (loss_per_sample.rows != B || loss_per_sample.cols != 1 ||
1408 6 loss_per_sample.dtype != ::brotensor::Dtype::FP32) {
1409 ✗ loss_per_sample.resize(B, 1, ::brotensor::Dtype::FP32);
1410 ✗ }
1411
3/6
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 6 times.
 6 if (B == 0 || n_act == 0 || n_heads <= 0) return;
1412
1413 6 const float* lp = logits_BL.host_f32();
1414 6 const float* tp = target_BL.host_f32();
1415 6 float* pp = probs_BL.host_f32_mut();
1416 6 float* dp = dLogits_BL.host_f32_mut();
1417 6 float* lossp = loss_per_sample.host_f32_mut();
1418
1419
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 34 times.
 40 for (int b = 0; b < B; ++b) {
1420 // Overwrite loss before per-head accumulation.
1421 34 float sample_loss = 0.0f;
1422 34 const int row_off = b * n_act;
1423
2/2
✓ Branch 0 taken 69 times.
✓ Branch 1 taken 34 times.
 103 for (int h = 0; h < n_heads; ++h) {
1424 69 const int off = d_head_offsets[h];
1425 69 const int end = d_head_offsets[h + 1];
1426 69 const int len = end - off;
1427 69 const float* lr = lp + row_off + off;
1428 69 const float* tr = tp + row_off + off;
1429
2/2
✓ Branch 0 taken 36 times.
✓ Branch 1 taken 33 times.
 69 const float* mr = d_mask_BL ? (d_mask_BL + row_off + off) : nullptr;
1430 69 float* pr = pp + row_off + off;
1431 69 float* dr = dp + row_off + off;
1432
1433 69 float m = -1e30f;
1434
2/2
✓ Branch 0 taken 2112 times.
✓ Branch 1 taken 69 times.
 2181 for (int i = 0; i < len; ++i) {
1435
4/4
✓ Branch 0 taken 1056 times.
✓ Branch 1 taken 1056 times.
✓ Branch 2 taken 528 times.
✓ Branch 3 taken 528 times.
 2112 if (mr && mr[i] < 0.5f) continue;
1436
2/2
✓ Branch 0 taken 1334 times.
✓ Branch 1 taken 250 times.
 1584 if (lr[i] > m) m = lr[i];
1437 1584 }
1438 69 float sum = 0.0f;
1439
2/2
✓ Branch 0 taken 2112 times.
✓ Branch 1 taken 69 times.
 2181 for (int i = 0; i < len; ++i) {
1440
4/4
✓ Branch 0 taken 1056 times.
✓ Branch 1 taken 1056 times.
✓ Branch 2 taken 528 times.
✓ Branch 3 taken 528 times.
 2112 if (mr && mr[i] < 0.5f) { pr[i] = 0.0f; continue; }
1441 1584 const float e = std::exp(lr[i] - m);
1442 1584 pr[i] = e;
1443 1584 sum += e;
1444 1584 }
1445
1/2
✓ Branch 0 taken 69 times.
✗ Branch 1 not taken.
 69 const float inv = sum > 0.0f ? 1.0f / sum : 0.0f;
1446
2/2
✓ Branch 0 taken 2112 times.
✓ Branch 1 taken 69 times.
 2181 for (int i = 0; i < len; ++i) {
1447
4/4
✓ Branch 0 taken 1056 times.
✓ Branch 1 taken 1056 times.
✓ Branch 2 taken 528 times.
✓ Branch 3 taken 528 times.
 2112 if (mr && mr[i] < 0.5f) { dr[i] = 0.0f; continue; }
1448 1584 const float p = pr[i] * inv;
1449 1584 pr[i] = p;
1450 1584 const float t = tr[i];
1451
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1584 times.
 1584 if (t > 0.0f) {
1452
1/2
✓ Branch 0 taken 1584 times.
✗ Branch 1 not taken.
 1584 const float pc = p > 1e-12f ? p : 1e-12f;
1453 1584 sample_loss -= t * std::log(pc);
1454 1584 }
1455 1584 dr[i] = p - t;
1456 1584 }
1457 69 }
1458 34 lossp[b] = sample_loss;
1459 34 }
1460 6 }
1461
1462 7 void bce_with_logits_fused_batched(const ::brotensor::Tensor& logits_BL,
1463 const ::brotensor::Tensor& target_BL,
1464 const float* d_mask_BL,
1465 float pos_weight,
1466 ::brotensor::Tensor& probs_BL,
1467 ::brotensor::Tensor& dLogits_BL,
1468 ::brotensor::Tensor& loss_per_sample) {
1469 7 const int B = logits_BL.rows;
1470 7 const int L = logits_BL.cols;
1471
3/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 taken 7 times.
 7 if (probs_BL.rows != B || probs_BL.cols != L ||
1472 7 probs_BL.dtype != ::brotensor::Dtype::FP32) {
1473 ✗ probs_BL.resize(B, L, ::brotensor::Dtype::FP32);
1474 ✗ }
1475
3/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 taken 7 times.
 7 if (dLogits_BL.rows != B || dLogits_BL.cols != L ||
1476 7 dLogits_BL.dtype != ::brotensor::Dtype::FP32) {
1477 ✗ dLogits_BL.resize(B, L, ::brotensor::Dtype::FP32);
1478 ✗ }
1479
3/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 taken 7 times.
 7 if (loss_per_sample.rows != B || loss_per_sample.cols != 1 ||
1480 7 loss_per_sample.dtype != ::brotensor::Dtype::FP32) {
1481 ✗ loss_per_sample.resize(B, 1, ::brotensor::Dtype::FP32);
1482 ✗ }
1483
2/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 7 times.
 7 if (B == 0 || L == 0) return;
1484
1485 7 const float* lp = logits_BL.host_f32();
1486 7 const float* tp = target_BL.host_f32();
1487 7 float* pp = probs_BL.host_f32_mut();
1488 7 float* dp = dLogits_BL.host_f32_mut();
1489 7 float* lossp = loss_per_sample.host_f32_mut();
1490 7 const float w = pos_weight;
1491
1492
2/2
✓ Branch 0 taken 7 times.
✓ Branch 1 taken 25 times.
 32 for (int b = 0; b < B; ++b) {
1493 25 float sample_loss = 0.0f;
1494 25 const int row_off = b * L;
1495 25 const float* lr = lp + row_off;
1496 25 const float* tr = tp + row_off;
1497
2/2
✓ Branch 0 taken 9 times.
✓ Branch 1 taken 16 times.
 25 const float* mr = d_mask_BL ? (d_mask_BL + row_off) : nullptr;
1498 25 float* pr = pp + row_off;
1499 25 float* dr = dp + row_off;
1500
2/2
✓ Branch 0 taken 54 times.
✓ Branch 1 taken 25 times.
 79 for (int i = 0; i < L; ++i) {
1501
4/4
✓ Branch 0 taken 34 times.
✓ Branch 1 taken 20 times.
✓ Branch 2 taken 17 times.
✓ Branch 3 taken 17 times.
 54 if (mr && mr[i] < 0.5f) {
1502 17 pr[i] = 0.0f;
1503 17 dr[i] = 0.0f;
1504 17 continue;
1505 }
1506 37 const float z = lr[i];
1507 37 const float y = tr[i];
1508 // softplus(x) = max(x, 0) + log1p(exp(-|x|))
1509 37 const float az = std::fabs(z);
1510
2/2
✓ Branch 0 taken 18 times.
✓ Branch 1 taken 19 times.
 37 const float sp_neg = (z > 0.0f ? 0.0f : -z) + std::log1p(std::exp(-az));
1511
2/2
✓ Branch 0 taken 18 times.
✓ Branch 1 taken 19 times.
 37 const float sp_pos = (z > 0.0f ? z : 0.0f) + std::log1p(std::exp(-az));
1512 37 sample_loss += w * y * sp_neg + (1.0f - y) * sp_pos;
1513 // s = sigmoid(z), numerically stable.
1514
2/2
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 15 times.
 37 const float s = z >= 0.0f
1515 22 ? 1.0f / (1.0f + std::exp(-z))
1516 15 : std::exp(z) / (1.0f + std::exp(z));
1517 37 pr[i] = s;
1518 37 dr[i] = s * (w * y + 1.0f - y) - w * y;
1519 37 }
1520 25 lossp[b] = sample_loss;
1521 25 }
1522 7 }
1523
1524 // ─── Masked mean-pool ──────────────────────────────────────────────────────
1525
1526 4 void build_slot_mask(const ::brotensor::Tensor& x, int offset, int K, int stride,
1527 ::brotensor::Tensor& mask) {
1528
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 4 if (mask.rows != K || mask.cols != 1 ||
1529 ✗ mask.dtype != ::brotensor::Dtype::FP32) {
1530 4 mask.resize(K, 1, ::brotensor::Dtype::FP32);
1531 4 }
1532 4 const float* xp = x.host_f32();
1533 4 float* mp = mask.host_f32_mut();
1534
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 73 times.
 77 for (int k = 0; k < K; ++k) {
1535 73 mp[k] = (xp[offset + static_cast<std::size_t>(k) * stride] > 0.5f)
1536 ? 1.0f : 0.0f;
1537 73 }
1538 4 }
1539
1540 4 void copy_d2d(const ::brotensor::Tensor& src, int src_off,
1541 ::brotensor::Tensor& dst, int dst_off, int n) {
1542 4 const float* sp = src.host_f32();
1543 4 float* dp = dst.host_f32_mut();
1544 4 std::memcpy(dp + dst_off, sp + src_off, static_cast<std::size_t>(n) * sizeof(float));
1545 4 }
1546
1547 4 void copy_d2d_strided(const ::brotensor::Tensor& src, int src_off, int src_pitch,
1548 ::brotensor::Tensor& dst, int dst_off, int dst_pitch,
1549 int width, int height) {
1550 4 const float* sp = src.host_f32() + src_off;
1551 4 float* dp = dst.host_f32_mut() + dst_off;
1552
2/2
✓ Branch 0 taken 84 times.
✓ Branch 1 taken 4 times.
 88 for (int r = 0; r < height; ++r) {
1553 168 std::memcpy(dp + static_cast<std::size_t>(r) * dst_pitch,
1554 84 sp + static_cast<std::size_t>(r) * src_pitch,
1555 84 static_cast<std::size_t>(width) * sizeof(float));
1556 84 }
1557 4 }
1558
1559 6 void masked_mean_pool_forward(const ::brotensor::Tensor& X, const float* d_mask,
1560 ::brotensor::Tensor& y) {
1561 6 const int K = X.rows;
1562 6 const int D = X.cols;
1563
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 6 if (y.rows != D || y.cols != 1 || y.dtype != ::brotensor::Dtype::FP32) {
1564 6 y.resize(D, 1, ::brotensor::Dtype::FP32);
1565 6 }
1566 6 const float* Xp = X.host_f32();
1567 6 float* yp = y.host_f32_mut();
1568
2/2
✓ Branch 0 taken 384 times.
✓ Branch 1 taken 6 times.
 390 for (int j = 0; j < D; ++j) yp[j] = 0.0f;
1569 6 int num_valid = 0;
1570
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 162 times.
 168 for (int k = 0; k < K; ++k) {
1571
4/4
✓ Branch 0 taken 137 times.
✓ Branch 1 taken 25 times.
✓ Branch 2 taken 57 times.
✓ Branch 3 taken 80 times.
 162 if (d_mask && d_mask[k] < 0.5f) continue;
1572 82 const float* row = Xp + static_cast<std::size_t>(k) * D;
1573
2/2
✓ Branch 0 taken 6800 times.
✓ Branch 1 taken 82 times.
 6882 for (int j = 0; j < D; ++j) yp[j] += row[j];
1574 82 ++num_valid;
1575 82 }
1576
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 5 times.
 6 if (num_valid > 0) {
1577 5 const float inv = 1.0f / static_cast<float>(num_valid);
1578
2/2
✓ Branch 0 taken 352 times.
✓ Branch 1 taken 5 times.
 357 for (int j = 0; j < D; ++j) yp[j] *= inv;
1579 5 }
1580 6 }
1581
1582 4 void masked_mean_pool_backward(const ::brotensor::Tensor& dY, const float* d_mask,
1583 int K, ::brotensor::Tensor& dX) {
1584 4 const int D = dY.size();
1585
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 4 if (dX.rows != K || dX.cols != D || dX.dtype != ::brotensor::Dtype::FP32) {
1586 4 dX.resize(K, D, ::brotensor::Dtype::FP32);
1587 4 }
1588 4 const float* dYp = dY.host_f32();
1589 4 float* dXp = dX.host_f32_mut();
1590 4 int num_valid = 0;
1591
2/2
✓ Branch 0 taken 122 times.
✓ Branch 1 taken 4 times.
 126 for (int k = 0; k < K; ++k) {
1592
4/4
✓ Branch 0 taken 97 times.
✓ Branch 1 taken 25 times.
✓ Branch 2 taken 51 times.
✓ Branch 3 taken 46 times.
 122 if (!d_mask || d_mask[k] >= 0.5f) ++num_valid;
1593 122 }
1594
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 const float inv = num_valid > 0 ? 1.0f / static_cast<float>(num_valid) : 0.0f;
1595
2/2
✓ Branch 0 taken 122 times.
✓ Branch 1 taken 4 times.
 126 for (int k = 0; k < K; ++k) {
1596 122 float* row = dXp + static_cast<std::size_t>(k) * D;
1597
4/4
✓ Branch 0 taken 97 times.
✓ Branch 1 taken 25 times.
✓ Branch 2 taken 51 times.
✓ Branch 3 taken 46 times.
 122 const bool valid = (!d_mask || d_mask[k] >= 0.5f) && num_valid > 0;
1598
4/4
✓ Branch 0 taken 12064 times.
✓ Branch 1 taken 122 times.
✓ Branch 2 taken 7072 times.
✓ Branch 3 taken 4992 times.
 12186 for (int j = 0; j < D; ++j) row[j] = valid ? dYp[j] * inv : 0.0f;
1599 122 }
1600 4 }
1601
1602 } // namespace brotensor::detail::cpu
1603