GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 93.0% 452 / 0 / 486
Functions: 97.0% 64 / 0 / 66
Branches: 49.1% 110 / 0 / 224
src/ops.cpp
Line Branch Exec Source
1 // brotensor public op wrappers — core dense, elementwise, norm, losses, embedding, concat, optim.
2 //
3 // Spatial/conv/pool/StyleGAN wrappers live in src/ops_spatial.cpp;
4 // attention/matmul/RoPE/RMSNorm/reductions live in src/ops_attention.cpp;
5 // quantisation/diffusion/audio/sampling/noise live in src/ops_quant_audio.cpp.
6
7 #include <brotensor/ops.h>
8 #include <brotensor/tensor.h>
9 #include <brotensor/detail/dispatch.h>
10
11 #include <cstdint>
12 #include <stdexcept>
13 #include <string>
14 #include <vector>
15
16 namespace brotensor {
17
18 namespace {
19
20 ✗ [[noreturn]] void throw_empty_parts(const char* op) {
21 ✗ std::string m = "brotensor: ";
22 ✗ m += op;
23 ✗ m += ": parts must be non-empty";
24 ✗ throw std::runtime_error(m);
25 ✗ }
26
27 67 const detail::OpsVTable& vtable_from_parts(
28 const std::vector<const Tensor*>& parts, const char* op) {
29
1/2
✓ Branch 0 taken 67 times.
✗ Branch 1 not taken.
 67 if (parts.empty() || parts[0] == nullptr) throw_empty_parts(op);
30 67 Device d = parts[0]->device;
31
2/2
✓ Branch 0 taken 99 times.
✓ Branch 1 taken 67 times.
 166 for (std::size_t i = 1; i < parts.size(); ++i) {
32
2/4
✓ Branch 0 taken 99 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 99 times.
✗ Branch 3 not taken.
 99 if (parts[i] && parts[i]->device != d) {
33 ✗ std::string m = "brotensor: ";
34 ✗ m += op;
35 ✗ m += ": parts[";
36 ✗ m += std::to_string(i);
37 ✗ m += "] on different device";
38 ✗ throw std::runtime_error(m);
39 ✗ }
40 99 }
41 67 return detail::ops_for(d);
42 ✗ }
43
44 } // namespace
45
46 // ─── Dense layers + elementwise ────────────────────────────────────────────
47
48 1264 void linear_forward(const Tensor& W, const Tensor& b,
49 const Tensor& x, Tensor& y) {
50 1264 const auto& v = detail::dispatch(W, b, x, y);
51
1/2
✓ Branch 0 taken 1264 times.
✗ Branch 1 not taken.
 1264 if (!v.linear_forward) detail::throw_not_implemented("linear_forward", W.device);
52 1264 detail::adopt_output(y, W.device);
53 1264 v.linear_forward(W, b, x, y);
54 1264 }
55
56 30 void linear_backward(const Tensor& W, const Tensor& x, const Tensor& dY,
57 Tensor& dX, Tensor& dW, Tensor& dB) {
58 30 const auto& v = detail::dispatch(W, x, dY, dX, dW, dB);
59
1/2
✓ Branch 0 taken 30 times.
✗ Branch 1 not taken.
 30 if (!v.linear_backward) detail::throw_not_implemented("linear_backward", W.device);
60 30 detail::adopt_output(dX, W.device);
61 30 detail::adopt_output(dW, W.device);
62 30 detail::adopt_output(dB, W.device);
63 30 v.linear_backward(W, x, dY, dX, dW, dB);
64 30 }
65
66 96 void relu_forward(const Tensor& x, Tensor& y) {
67 96 const auto& v = detail::dispatch(x, y);
68
1/2
✓ Branch 0 taken 96 times.
✗ Branch 1 not taken.
 96 if (!v.relu_forward) detail::throw_not_implemented("relu_forward", x.device);
69 96 detail::adopt_output(y, x.device);
70 96 v.relu_forward(x, y);
71 96 }
72
73 11 void relu_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
74 11 const auto& v = detail::dispatch(x, dY, dX);
75
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.relu_backward) detail::throw_not_implemented("relu_backward", x.device);
76 11 detail::adopt_output(dX, x.device);
77 11 v.relu_backward(x, dY, dX);
78 11 }
79
80 99 void tanh_forward(const Tensor& x, Tensor& y) {
81 99 const auto& v = detail::dispatch(x, y);
82
1/2
✓ Branch 0 taken 99 times.
✗ Branch 1 not taken.
 99 if (!v.tanh_forward) detail::throw_not_implemented("tanh_forward", x.device);
83 99 detail::adopt_output(y, x.device);
84 99 v.tanh_forward(x, y);
85 99 }
86
87 11 void tanh_backward(const Tensor& y, const Tensor& dY, Tensor& dX) {
88 11 const auto& v = detail::dispatch(y, dY, dX);
89
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.tanh_backward) detail::throw_not_implemented("tanh_backward", y.device);
90 11 detail::adopt_output(dX, y.device);
91 11 v.tanh_backward(y, dY, dX);
92 11 }
93
94 26 void sigmoid_forward(const Tensor& x, Tensor& y) {
95 26 const auto& v = detail::dispatch(x, y);
96
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 26 if (!v.sigmoid_forward) detail::throw_not_implemented("sigmoid_forward", x.device);
97 26 detail::adopt_output(y, x.device);
98 26 v.sigmoid_forward(x, y);
99 26 }
100
101 11 void sigmoid_backward(const Tensor& y, const Tensor& dY, Tensor& dX) {
102 11 const auto& v = detail::dispatch(y, dY, dX);
103
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.sigmoid_backward) detail::throw_not_implemented("sigmoid_backward", y.device);
104 11 detail::adopt_output(dX, y.device);
105 11 v.sigmoid_backward(y, dY, dX);
106 11 }
107
108 1182 void add_inplace(Tensor& y, const Tensor& x) {
109 1182 const auto& v = detail::dispatch(y, x);
110
1/2
✓ Branch 0 taken 1182 times.
✗ Branch 1 not taken.
 1182 if (!v.add_inplace) detail::throw_not_implemented("add_inplace", y.device);
111 1182 detail::adopt_output(y, y.device);
112 1182 v.add_inplace(y, x);
113 1182 }
114
115 4 void axpby_inplace(Tensor& y, const Tensor& x, float a, float b) {
116 4 const auto& v = detail::dispatch(y, x);
117
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!v.axpby_inplace) detail::throw_not_implemented("axpby_inplace", y.device);
118 4 detail::adopt_output(y, y.device);
119 4 v.axpby_inplace(y, x, a, b);
120 4 }
121
122 16 void add_scalar_inplace(Tensor& y, float s) {
123 16 const auto& v = detail::dispatch(y);
124
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.add_scalar_inplace) detail::throw_not_implemented("add_scalar_inplace", y.device);
125 16 detail::adopt_output(y, y.device);
126 16 v.add_scalar_inplace(y, s);
127 16 }
128
129 39 void add_channel_bias_inplace(Tensor& y, const Tensor& bias, int C, int L) {
130 39 const auto& v = detail::dispatch(y);
131
1/2
✓ Branch 0 taken 39 times.
✗ Branch 1 not taken.
 39 if (!v.add_channel_bias_inplace)
132 ✗ detail::throw_not_implemented("add_channel_bias_inplace", y.device);
133 39 detail::adopt_output(y, y.device);
134 39 v.add_channel_bias_inplace(y, bias, C, L);
135 39 }
136
137 14 void add_row_bias_inplace(Tensor& Y, const Tensor& bias) {
138 14 const auto& v = detail::dispatch(Y, bias);
139
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.add_row_bias_inplace)
140 ✗ detail::throw_not_implemented("add_row_bias_inplace", Y.device);
141 14 detail::adopt_output(Y, Y.device);
142 14 v.add_row_bias_inplace(Y, bias);
143 14 }
144
145 301 void scale_inplace(Tensor& y, float s) {
146 301 const auto& v = detail::dispatch(y);
147
1/2
✓ Branch 0 taken 301 times.
✗ Branch 1 not taken.
 301 if (!v.scale_inplace) detail::throw_not_implemented("scale_inplace", y.device);
148 301 detail::adopt_output(y, y.device);
149 301 v.scale_inplace(y, s);
150 301 }
151
152 17 void clamp(Tensor& y, float lo, float hi) {
153
3/4
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 16 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
 17 if (detail::jit_clamp_inplace_hook && detail::jit_clamp_inplace_hook(y, lo, hi)) {
154 1 return;
155 }
156 16 const auto& v = detail::dispatch(y);
157
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.clamp) detail::throw_not_implemented("clamp", y.device);
158 16 detail::adopt_output(y, y.device);
159 16 v.clamp(y, lo, hi);
160 17 }
161
162 1263 void mul_inplace(Tensor& y, const Tensor& x) {
163 1263 const auto& v = detail::dispatch(y, x);
164
1/2
✓ Branch 0 taken 1263 times.
✗ Branch 1 not taken.
 1263 if (!v.mul_inplace) detail::throw_not_implemented("mul_inplace", y.device);
165 1263 detail::adopt_output(y, y.device);
166 1263 v.mul_inplace(y, x);
167 1263 }
168
169 ✗ void div_inplace(Tensor& y, const Tensor& x) {
170 ✗ const auto& v = detail::dispatch(y, x);
171 ✗ if (!v.div_inplace) detail::throw_not_implemented("div_inplace", y.device);
172 ✗ detail::adopt_output(y, y.device);
173 ✗ v.div_inplace(y, x);
174 ✗ }
175
176 732 void modulate(const Tensor& X, const Tensor& scale, const Tensor& shift,
177 Tensor& Y) {
178 732 const auto& v = detail::dispatch(X, scale, shift, Y);
179
1/2
✓ Branch 0 taken 732 times.
✗ Branch 1 not taken.
 732 if (!v.modulate) detail::throw_not_implemented("modulate", X.device);
180 732 detail::adopt_output(Y, X.device);
181 732 v.modulate(X, scale, shift, Y);
182 732 }
183
184 490 void broadcast_mul(const Tensor& X, const Tensor& vv, Tensor& Y) {
185 490 const auto& v = detail::dispatch(X, vv, Y);
186
1/2
✓ Branch 0 taken 490 times.
✗ Branch 1 not taken.
 490 if (!v.broadcast_mul) detail::throw_not_implemented("broadcast_mul", X.device);
187 490 detail::adopt_output(Y, X.device);
188 490 v.broadcast_mul(X, vv, Y);
189 490 }
190
191 8 void build_slot_mask(const Tensor& x, int offset, int K, int stride, Tensor& mask) {
192 8 const auto& v = detail::dispatch(x, mask);
193
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.build_slot_mask) detail::throw_not_implemented("build_slot_mask", x.device);
194 8 detail::adopt_output(mask, x.device);
195 8 v.build_slot_mask(x, offset, K, stride, mask);
196 8 }
197
198 // ─── Reductions / norm / softmax / attention (training) ────────────────────
199
200 42 void softmax_forward(const Tensor& logits, Tensor& probs, const float* mask) {
201 42 const auto& v = detail::dispatch(logits, probs);
202
1/2
✓ Branch 0 taken 42 times.
✗ Branch 1 not taken.
 42 if (!v.softmax_forward) detail::throw_not_implemented("softmax_forward", logits.device);
203 42 detail::adopt_output(probs, logits.device);
204 42 v.softmax_forward(logits, probs, mask);
205 42 }
206
207 46 void softmax_rows_forward(const Tensor& X, Tensor& Y, int rows, int cols) {
208 46 const auto& v = detail::dispatch(X, Y);
209
1/2
✓ Branch 0 taken 46 times.
✗ Branch 1 not taken.
 46 if (!v.softmax_rows_forward) detail::throw_not_implemented("softmax_rows_forward", X.device);
210 46 detail::adopt_output(Y, X.device);
211 46 v.softmax_rows_forward(X, Y, rows, cols);
212 46 }
213
214 21 void softmax_backward(const Tensor& probs, const Tensor& dProbs, Tensor& dLogits) {
215 21 const auto& v = detail::dispatch(probs, dProbs, dLogits);
216
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 if (!v.softmax_backward) detail::throw_not_implemented("softmax_backward", probs.device);
217 21 detail::adopt_output(dLogits, probs.device);
218 21 v.softmax_backward(probs, dProbs, dLogits);
219 21 }
220
221 8 void layernorm_forward(const Tensor& x, const Tensor& gamma, const Tensor& beta,
222 Tensor& y, Tensor& xhat,
223 float& mean_out, float& rstd_out, float eps) {
224 8 const auto& v = detail::dispatch(x, gamma, beta, y, xhat);
225
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.layernorm_forward) detail::throw_not_implemented("layernorm_forward", x.device);
226 8 detail::adopt_output(y, x.device);
227 8 detail::adopt_output(xhat, x.device);
228 8 v.layernorm_forward(x, gamma, beta, y, xhat, mean_out, rstd_out, eps);
229 8 }
230
231 14 void layernorm_backward(const Tensor& dY, const Tensor& xhat,
232 const Tensor& gamma, float rstd,
233 Tensor& dX, Tensor& dGamma, Tensor& dBeta) {
234 14 const auto& v = detail::dispatch(dY, xhat, gamma, dX, dGamma, dBeta);
235
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.layernorm_backward) detail::throw_not_implemented("layernorm_backward", dY.device);
236 14 detail::adopt_output(dX, dY.device);
237 14 detail::adopt_output(dGamma, dY.device);
238 14 detail::adopt_output(dBeta, dY.device);
239 14 v.layernorm_backward(dY, xhat, gamma, rstd, dX, dGamma, dBeta);
240 14 }
241
242 20 void attention_forward(const Tensor& X,
243 const Tensor& Wq, const Tensor& Wk,
244 const Tensor& Wv, const Tensor& Wo,
245 const float* d_mask,
246 Tensor& Q, Tensor& K, Tensor& V,
247 Tensor& Attn, Tensor& Y_pre_Wo, Tensor& O) {
248 20 const auto& v = detail::dispatch(X, Wq, Wk, Wv, Wo, Q, K, V);
249
1/2
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
 20 if (!v.attention_forward) detail::throw_not_implemented("attention_forward", X.device);
250 20 detail::adopt_output(Q, X.device);
251 20 detail::adopt_output(K, X.device);
252 20 detail::adopt_output(V, X.device);
253 20 detail::adopt_output(Attn, X.device);
254 20 detail::adopt_output(Y_pre_Wo, X.device);
255 20 detail::adopt_output(O, X.device);
256 20 v.attention_forward(X, Wq, Wk, Wv, Wo, d_mask, Q, K, V, Attn, Y_pre_Wo, O);
257 20 }
258
259 16 void attention_backward(const Tensor& dO, const Tensor& X,
260 const Tensor& Q, const Tensor& K,
261 const Tensor& V, const Tensor& Attn,
262 const Tensor& Y_pre_Wo,
263 const Tensor& Wq, const Tensor& Wk,
264 const Tensor& Wv, const Tensor& Wo,
265 const float* d_mask,
266 Tensor& dX,
267 Tensor& dWq, Tensor& dWk,
268 Tensor& dWv, Tensor& dWo) {
269 16 const auto& v = detail::dispatch(dO, X, Q, K, V, Attn, Y_pre_Wo);
270
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.attention_backward) detail::throw_not_implemented("attention_backward", dO.device);
271 16 detail::adopt_output(dX, dO.device);
272 16 detail::adopt_output(dWq, dO.device);
273 16 detail::adopt_output(dWk, dO.device);
274 16 detail::adopt_output(dWv, dO.device);
275 16 detail::adopt_output(dWo, dO.device);
276 32 v.attention_backward(dO, X, Q, K, V, Attn, Y_pre_Wo,
277 16 Wq, Wk, Wv, Wo, d_mask,
278 16 dX, dWq, dWk, dWv, dWo);
279 16 }
280
281 32 void mha_forward(const Tensor& X,
282 const Tensor& Wq, const Tensor& Wk,
283 const Tensor& Wv, const Tensor& Wo,
284 const Tensor* bq, const Tensor* bk,
285 const Tensor* bv, const Tensor* bo,
286 const float* d_mask, int num_heads,
287 Tensor& Qh, Tensor& Kh, Tensor& Vh,
288 Tensor& Attnh, Tensor& Yconcat, Tensor& O) {
289 32 const auto& v = detail::dispatch_with_opts(
290 32 X, Wq, {&Wk, &Wv, &Wo, bq, bk, bv, bo, &Qh, &Kh, &Vh});
291
1/2
✓ Branch 0 taken 32 times.
✗ Branch 1 not taken.
 32 if (!v.mha_forward) detail::throw_not_implemented("mha_forward", X.device);
292 32 detail::adopt_output(Qh, X.device);
293 32 detail::adopt_output(Kh, X.device);
294 32 detail::adopt_output(Vh, X.device);
295 32 detail::adopt_output(Attnh, X.device);
296 32 detail::adopt_output(Yconcat, X.device);
297 32 detail::adopt_output(O, X.device);
298 64 v.mha_forward(X, Wq, Wk, Wv, Wo, bq, bk, bv, bo, d_mask, num_heads,
299 32 Qh, Kh, Vh, Attnh, Yconcat, O);
300 32 }
301
302 30 void mha_backward(const Tensor& dO, const Tensor& X,
303 const Tensor& Qh, const Tensor& Kh,
304 const Tensor& Vh, const Tensor& Attnh,
305 const Tensor& Yconcat,
306 const Tensor& Wq, const Tensor& Wk,
307 const Tensor& Wv, const Tensor& Wo,
308 const float* d_mask, int num_heads,
309 Tensor& dX,
310 Tensor& dWq, Tensor& dWk,
311 Tensor& dWv, Tensor& dWo,
312 Tensor* dbq, Tensor* dbk,
313 Tensor* dbv, Tensor* dbo) {
314 30 const auto& v = detail::dispatch_with_opts(
315 30 dO, X, {&Qh, &Kh, &Vh, &Attnh, &Yconcat, dbq, dbk, dbv, dbo});
316
1/2
✓ Branch 0 taken 30 times.
✗ Branch 1 not taken.
 30 if (!v.mha_backward) detail::throw_not_implemented("mha_backward", dO.device);
317 30 detail::adopt_output(dX, dO.device);
318 30 detail::adopt_output(dWq, dO.device);
319 30 detail::adopt_output(dWk, dO.device);
320 30 detail::adopt_output(dWv, dO.device);
321 30 detail::adopt_output(dWo, dO.device);
322
2/2
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 8 times.
 30 if (dbq) detail::adopt_output(*dbq, dO.device);
323
2/2
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 8 times.
 30 if (dbk) detail::adopt_output(*dbk, dO.device);
324
2/2
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 8 times.
 30 if (dbv) detail::adopt_output(*dbv, dO.device);
325
2/2
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 8 times.
 30 if (dbo) detail::adopt_output(*dbo, dO.device);
326 60 v.mha_backward(dO, X, Qh, Kh, Vh, Attnh, Yconcat,
327 30 Wq, Wk, Wv, Wo, d_mask, num_heads,
328 30 dX, dWq, dWk, dWv, dWo,
329 30 dbq, dbk, dbv, dbo);
330 30 }
331
332 // ─── Pooling / losses / embedding / concat ─────────────────────────────────
333
334 22 void masked_mean_pool_forward(const Tensor& X, const float* d_mask, Tensor& y) {
335 22 const auto& v = detail::dispatch(X, y);
336
1/2
✓ Branch 0 taken 22 times.
✗ Branch 1 not taken.
 22 if (!v.masked_mean_pool_forward) detail::throw_not_implemented("masked_mean_pool_forward", X.device);
337 22 detail::adopt_output(y, X.device);
338 22 v.masked_mean_pool_forward(X, d_mask, y);
339 22 }
340
341 16 void masked_mean_pool_backward(const Tensor& dY, const float* d_mask,
342 int K, Tensor& dX) {
343 16 const auto& v = detail::dispatch(dY, dX);
344
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.masked_mean_pool_backward) detail::throw_not_implemented("masked_mean_pool_backward", dY.device);
345 16 detail::adopt_output(dX, dY.device);
346 16 v.masked_mean_pool_backward(dY, d_mask, K, dX);
347 16 }
348
349 9 float mse_vec_forward(const Tensor& pred, const Tensor& target) {
350 9 const auto& v = detail::dispatch(pred, target);
351
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.mse_vec_forward) detail::throw_not_implemented("mse_vec_forward", pred.device);
352 9 return v.mse_vec_forward(pred, target);
353 }
354
355 9 void mse_vec_backward(const Tensor& pred, const Tensor& target, Tensor& dPred) {
356 9 const auto& v = detail::dispatch(pred, target, dPred);
357
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.mse_vec_backward) detail::throw_not_implemented("mse_vec_backward", pred.device);
358 9 detail::adopt_output(dPred, pred.device);
359 9 v.mse_vec_backward(pred, target, dPred);
360 9 }
361
362 7 float mse_scalar(float pred, float target, float& dPred) {
363 7 const auto& v = detail::ops_for(Device::CPU);
364
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (!v.mse_scalar) detail::throw_not_implemented("mse_scalar", Device::CPU);
365 7 return v.mse_scalar(pred, target, dPred);
366 }
367
368 7 float softmax_xent(const Tensor& logits, const Tensor& target,
369 Tensor& probs, Tensor& dLogits, const float* mask) {
370 7 const auto& v = detail::dispatch(logits, target, probs, dLogits);
371
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (!v.softmax_xent) detail::throw_not_implemented("softmax_xent", logits.device);
372 7 detail::adopt_output(probs, logits.device);
373 7 detail::adopt_output(dLogits, logits.device);
374 7 return v.softmax_xent(logits, target, probs, dLogits, mask);
375 }
376
377 11 float softmax_xent_segment(const float* logits, const float* target,
378 float* probs, float* dLogits,
379 int n, const float* mask) {
380 11 const auto& v = detail::ops_for(Device::CPU);
381
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.softmax_xent_segment) detail::throw_not_implemented("softmax_xent_segment", Device::CPU);
382 11 return v.softmax_xent_segment(logits, target, probs, dLogits, n, mask);
383 }
384
385 19 float softmax_xent_fused(const Tensor& logits, const Tensor& target,
386 const float* d_mask,
387 Tensor& probs, Tensor& dLogits) {
388 19 const auto& v = detail::dispatch(logits, target, probs, dLogits);
389
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 if (!v.softmax_xent_fused) detail::throw_not_implemented("softmax_xent_fused", logits.device);
390 19 detail::adopt_output(probs, logits.device);
391 19 detail::adopt_output(dLogits, logits.device);
392 19 return v.softmax_xent_fused(logits, target, d_mask, probs, dLogits);
393 }
394
395 39 void embedding_lookup_forward(const Tensor& table, const int32_t* d_idx,
396 int B, Tensor& out) {
397 39 const auto& v = detail::dispatch(table, out);
398
1/2
✓ Branch 0 taken 39 times.
✗ Branch 1 not taken.
 39 if (!v.embedding_lookup_forward) detail::throw_not_implemented("embedding_lookup_forward", table.device);
399 39 detail::adopt_output(out, table.device);
400 39 v.embedding_lookup_forward(table, d_idx, B, out);
401 39 }
402
403 27 void embedding_lookup_backward(const Tensor& dOut, const int32_t* d_idx,
404 int B, Tensor& dTable) {
405 27 const auto& v = detail::dispatch(dOut, dTable);
406
1/2
✓ Branch 0 taken 27 times.
✗ Branch 1 not taken.
 27 if (!v.embedding_lookup_backward) detail::throw_not_implemented("embedding_lookup_backward", dOut.device);
407 27 detail::adopt_output(dTable, dOut.device);
408 27 v.embedding_lookup_backward(dOut, d_idx, B, dTable);
409 27 }
410
411 17 void concat_rows(const std::vector<const Tensor*>& parts, Tensor& out) {
412 17 const auto& v = vtable_from_parts(parts, "concat_rows");
413
3/4
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 15 times.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 17 if (out.data != nullptr && parts[0]->device != out.device) {
414 ✗ throw std::runtime_error("brotensor: concat_rows: out on different device");
415 }
416 17 detail::adopt_output(out, parts[0]->device);
417
1/2
✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
 17 if (!v.concat_rows) detail::throw_not_implemented("concat_rows", parts[0]->device);
418 17 v.concat_rows(parts, out);
419 17 }
420
421 12 void split_rows(const Tensor& in, const std::vector<Tensor*>& parts) {
422 12 Device d = in.device;
423
2/2
✓ Branch 0 taken 40 times.
✓ Branch 1 taken 12 times.
 52 for (std::size_t i = 0; i < parts.size(); ++i) {
424
3/6
✓ Branch 0 taken 40 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 40 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 40 times.
✗ Branch 5 not taken.
 40 if (parts[i] && parts[i]->data != nullptr && parts[i]->device != d) {
425 ✗ throw std::runtime_error("brotensor: split_rows: parts on different device");
426 }
427 40 }
428 12 const auto& v = detail::ops_for(d);
429
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.split_rows) detail::throw_not_implemented("split_rows", d);
430
3/4
✓ Branch 0 taken 40 times.
✓ Branch 1 taken 12 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 40 times.
 52 for (Tensor* _p : parts) if (_p) detail::adopt_output(*_p, d);
431 12 v.split_rows(in, parts);
432 12 }
433
434 23 void concat_batched_rows(const std::vector<const Tensor*>& parts, Tensor& out) {
435 23 const auto& v = vtable_from_parts(parts, "concat_batched_rows");
436
3/4
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 21 times.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 23 if (out.data != nullptr && parts[0]->device != out.device) {
437 ✗ throw std::runtime_error("brotensor: concat_batched_rows: out on different device");
438 }
439 23 detail::adopt_output(out, parts[0]->device);
440
1/2
✓ Branch 0 taken 23 times.
✗ Branch 1 not taken.
 23 if (!v.concat_batched_rows) detail::throw_not_implemented("concat_batched_rows", parts[0]->device);
441 23 v.concat_batched_rows(parts, out);
442 23 }
443
444 27 void concat_nchw_channels(const std::vector<const Tensor*>& parts,
445 int N, int H, int W,
446 const std::vector<int>& C_per_part,
447 Tensor& out) {
448 27 const auto& v = vtable_from_parts(parts, "concat_nchw_channels");
449
3/4
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 25 times.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 27 if (out.data != nullptr && parts[0]->device != out.device) {
450 ✗ throw std::runtime_error("brotensor: concat_nchw_channels: out on different device");
451 }
452 27 detail::adopt_output(out, parts[0]->device);
453
1/2
✓ Branch 0 taken 27 times.
✗ Branch 1 not taken.
 27 if (!v.concat_nchw_channels) detail::throw_not_implemented("concat_nchw_channels", parts[0]->device);
454 27 v.concat_nchw_channels(parts, N, H, W, C_per_part, out);
455 27 }
456
457 24 void concat_nchw_channels_backward(const Tensor& dY, int N, int H, int W,
458 const std::vector<int>& C_per_part,
459 const std::vector<Tensor*>& parts) {
460 24 Device d = dY.device;
461
2/2
✓ Branch 0 taken 48 times.
✓ Branch 1 taken 24 times.
 72 for (std::size_t i = 0; i < parts.size(); ++i) {
462
4/6
✓ Branch 0 taken 48 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✓ Branch 3 taken 44 times.
✓ Branch 4 taken 4 times.
✗ Branch 5 not taken.
 48 if (parts[i] && parts[i]->data != nullptr && parts[i]->device != d) {
463 ✗ throw std::runtime_error("brotensor: concat_nchw_channels_backward: parts on different device");
464 }
465 48 }
466 24 const auto& v = detail::ops_for(d);
467
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 if (!v.concat_nchw_channels_backward) detail::throw_not_implemented("concat_nchw_channels_backward", d);
468
3/4
✓ Branch 0 taken 48 times.
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 48 times.
 72 for (Tensor* _p : parts) if (_p) detail::adopt_output(*_p, d);
469 24 v.concat_nchw_channels_backward(dY, N, H, W, C_per_part, parts);
470 24 }
471
472 36 void copy_d2d(const Tensor& src, int src_off, Tensor& dst, int dst_off, int n) {
473 36 const auto& v = detail::dispatch(src, dst);
474
1/2
✓ Branch 0 taken 36 times.
✗ Branch 1 not taken.
 36 if (!v.copy_d2d) detail::throw_not_implemented("copy_d2d", src.device);
475 36 detail::adopt_output(dst, src.device);
476 36 v.copy_d2d(src, src_off, dst, dst_off, n);
477 36 }
478
479 10 void copy_d2d_strided(const Tensor& src, int src_off, int src_pitch,
480 Tensor& dst, int dst_off, int dst_pitch,
481 int width, int height) {
482
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (width > src_pitch || width > dst_pitch) {
483 ✗ throw std::runtime_error("brotensor: copy_d2d_strided: width exceeds pitch");
484 }
485 10 const auto& v = detail::dispatch(src, dst);
486
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.copy_d2d_strided) detail::throw_not_implemented("copy_d2d_strided", src.device);
487 10 detail::adopt_output(dst, src.device);
488 10 v.copy_d2d_strided(src, src_off, src_pitch, dst, dst_off, dst_pitch, width, height);
489 10 }
490
491 230 void cast(const Tensor& src, Tensor& dst, Dtype out_dtype) {
492
4/4
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 227 times.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 1 time.
 230 if (detail::jit_cast_out_hook && detail::jit_cast_out_hook(src, dst, out_dtype)) {
493 1 return;
494 }
495 229 const auto& v = detail::dispatch(src, dst);
496
1/2
✓ Branch 0 taken 229 times.
✗ Branch 1 not taken.
 229 if (!v.cast) detail::throw_not_implemented("cast", src.device);
497 229 detail::adopt_output(dst, src.device);
498 229 v.cast(src, dst, out_dtype);
499 230 }
500
501 // ─── Inference batched + optim ─────────────────────────────────────────────
502
503 758 void layernorm_forward_inference_batched(const Tensor& X_RD,
504 const Tensor& gamma,
505 const Tensor& beta,
506 Tensor& Y_RD, float eps) {
507 758 const auto& v = detail::dispatch(X_RD, gamma, beta, Y_RD);
508
1/2
✓ Branch 0 taken 758 times.
✗ Branch 1 not taken.
 758 if (!v.layernorm_forward_inference_batched)
509 ✗ detail::throw_not_implemented("layernorm_forward_inference_batched", X_RD.device);
510 758 detail::adopt_output(Y_RD, X_RD.device);
511 758 v.layernorm_forward_inference_batched(X_RD, gamma, &beta, Y_RD, eps);
512 758 }
513
514 7 void layernorm_forward_inference_batched(const Tensor& X_RD,
515 const Tensor& gamma,
516 Tensor& Y_RD, float eps) {
517 7 const auto& v = detail::dispatch(X_RD, gamma, Y_RD);
518
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (!v.layernorm_forward_inference_batched)
519 ✗ detail::throw_not_implemented("layernorm_forward_inference_batched", X_RD.device);
520 7 detail::adopt_output(Y_RD, X_RD.device);
521 7 v.layernorm_forward_inference_batched(X_RD, gamma, nullptr, Y_RD, eps);
522 7 }
523
524 18 void layernorm_forward_batched_with_caches(const Tensor& X_RD,
525 const Tensor& gamma,
526 const Tensor& beta,
527 Tensor& Y_RD, Tensor& Xhat_RD,
528 Tensor& Mean_R, Tensor& Rstd_R,
529 float eps) {
530 18 const auto& v = detail::dispatch(X_RD, gamma, beta, Y_RD, Xhat_RD, Mean_R, Rstd_R);
531
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 if (!v.layernorm_forward_batched_with_caches)
532 ✗ detail::throw_not_implemented("layernorm_forward_batched_with_caches", X_RD.device);
533 18 detail::adopt_output(Y_RD, X_RD.device);
534 18 detail::adopt_output(Xhat_RD, X_RD.device);
535 18 detail::adopt_output(Mean_R, X_RD.device);
536 18 detail::adopt_output(Rstd_R, X_RD.device);
537 36 v.layernorm_forward_batched_with_caches(X_RD, gamma, beta, Y_RD, Xhat_RD,
538 18 Mean_R, Rstd_R, eps);
539 18 }
540
541 18 void layernorm_backward_batched_with_caches(const Tensor& dY_RD,
542 const Tensor& Xhat_RD,
543 const Tensor& gamma,
544 const Tensor& Rstd_R,
545 Tensor& dX_RD,
546 Tensor& dGamma, Tensor& dBeta) {
547 18 const auto& v = detail::dispatch(dY_RD, Xhat_RD, gamma, Rstd_R, dX_RD, dGamma, dBeta);
548
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 if (!v.layernorm_backward_batched_with_caches)
549 ✗ detail::throw_not_implemented("layernorm_backward_batched_with_caches", dY_RD.device);
550 18 detail::adopt_output(dX_RD, dY_RD.device);
551 18 detail::adopt_output(dGamma, dY_RD.device);
552 18 detail::adopt_output(dBeta, dY_RD.device);
553 36 v.layernorm_backward_batched_with_caches(dY_RD, Xhat_RD, gamma, Rstd_R,
554 18 dX_RD, dGamma, dBeta);
555 18 }
556
557 5 void sgd_step(Tensor& param, Tensor& grad, Tensor& velocity,
558 float lr, float momentum) {
559 5 const auto& v = detail::dispatch(param, grad, velocity);
560
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 if (!v.sgd_step) detail::throw_not_implemented("sgd_step", param.device);
561 5 detail::adopt_output(param, param.device);
562 5 detail::adopt_output(grad, param.device);
563 5 detail::adopt_output(velocity, param.device);
564 5 v.sgd_step(param, grad, velocity, lr, momentum);
565 5 }
566
567 50 void adam_step(Tensor& param, const Tensor& grad,
568 Tensor& m, Tensor& v_buf,
569 float lr, float beta1, float beta2, float eps, int step) {
570 50 const auto& v = detail::dispatch(param, grad, m, v_buf);
571
1/2
✓ Branch 0 taken 50 times.
✗ Branch 1 not taken.
 50 if (!v.adam_step) detail::throw_not_implemented("adam_step", param.device);
572 50 detail::adopt_output(param, param.device);
573 50 detail::adopt_output(m, param.device);
574 50 detail::adopt_output(v_buf, param.device);
575 50 v.adam_step(param, grad, m, v_buf, lr, beta1, beta2, eps, step);
576 50 }
577
578 3 void xavier_init(Tensor& W, uint64_t& rng_state) {
579 3 const auto& v = detail::dispatch(W);
580
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.xavier_init) detail::throw_not_implemented("xavier_init", W.device);
581 3 detail::adopt_output(W, W.device);
582 3 v.xavier_init(W, rng_state);
583 3 }
584
585 // ─── Batched inference variants ────────────────────────────────────────────
586
587 97 void linear_forward_batched(const Tensor& W, const Tensor& bias,
588 const Tensor& X_BD, Tensor& Y_BD) {
589 97 const auto& v = detail::dispatch(W, bias, X_BD, Y_BD);
590
1/2
✓ Branch 0 taken 97 times.
✗ Branch 1 not taken.
 97 if (!v.linear_forward_batched) detail::throw_not_implemented("linear_forward_batched", W.device);
591 97 detail::adopt_output(Y_BD, W.device);
592 97 v.linear_forward_batched(W, bias, X_BD, Y_BD);
593 97 }
594
595 5 void relu_forward_batched(const Tensor& X_BD, Tensor& Y_BD) {
596 5 const auto& v = detail::dispatch(X_BD, Y_BD);
597
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 if (!v.relu_forward_batched) detail::throw_not_implemented("relu_forward_batched", X_BD.device);
598 5 detail::adopt_output(Y_BD, X_BD.device);
599 5 v.relu_forward_batched(X_BD, Y_BD);
600 5 }
601
602 11 void tanh_forward_batched(const Tensor& X_BD, Tensor& Y_BD) {
603 11 const auto& v = detail::dispatch(X_BD, Y_BD);
604
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.tanh_forward_batched) detail::throw_not_implemented("tanh_forward_batched", X_BD.device);
605 11 detail::adopt_output(Y_BD, X_BD.device);
606 11 v.tanh_forward_batched(X_BD, Y_BD);
607 11 }
608
609 5 void add_inplace_batched(Tensor& Y_BD, const Tensor& X_BD) {
610 5 const auto& v = detail::dispatch(Y_BD, X_BD);
611
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 if (!v.add_inplace_batched) detail::throw_not_implemented("add_inplace_batched", Y_BD.device);
612 5 detail::adopt_output(Y_BD, Y_BD.device);
613 5 v.add_inplace_batched(Y_BD, X_BD);
614 5 }
615
616 // ─── Batched backward variants ─────────────────────────────────────────────
617
618 55 void linear_backward_batched(const Tensor& W, const Tensor& X_BD,
619 const Tensor& dY_BD,
620 Tensor& dX_BD, Tensor& dW, Tensor& dB) {
621 55 const auto& v = detail::dispatch(W, X_BD, dY_BD, dX_BD, dW, dB);
622
1/2
✓ Branch 0 taken 55 times.
✗ Branch 1 not taken.
 55 if (!v.linear_backward_batched) detail::throw_not_implemented("linear_backward_batched", W.device);
623 55 detail::adopt_output(dX_BD, W.device);
624 55 detail::adopt_output(dW, W.device);
625 55 detail::adopt_output(dB, W.device);
626 55 v.linear_backward_batched(W, X_BD, dY_BD, dX_BD, dW, dB);
627 55 }
628
629 8 void relu_backward_batched(const Tensor& X_BD, const Tensor& dY_BD, Tensor& dX_BD) {
630 8 const auto& v = detail::dispatch(X_BD, dY_BD, dX_BD);
631
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.relu_backward_batched) detail::throw_not_implemented("relu_backward_batched", X_BD.device);
632 8 detail::adopt_output(dX_BD, X_BD.device);
633 8 v.relu_backward_batched(X_BD, dY_BD, dX_BD);
634 8 }
635
636 8 void tanh_backward_batched(const Tensor& Y_BD, const Tensor& dY_BD, Tensor& dX_BD) {
637 8 const auto& v = detail::dispatch(Y_BD, dY_BD, dX_BD);
638
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.tanh_backward_batched) detail::throw_not_implemented("tanh_backward_batched", Y_BD.device);
639 8 detail::adopt_output(dX_BD, Y_BD.device);
640 8 v.tanh_backward_batched(Y_BD, dY_BD, dX_BD);
641 8 }
642
643 // ─── Batched per-sample losses ─────────────────────────────────────────────
644
645 6 void mse_vec_per_sample(const Tensor& pred, const Tensor& target,
646 Tensor& dPred, Tensor& loss_per_sample) {
647 6 const auto& v = detail::dispatch(pred, target, dPred, loss_per_sample);
648
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.mse_vec_per_sample) detail::throw_not_implemented("mse_vec_per_sample", pred.device);
649 6 detail::adopt_output(dPred, pred.device);
650 6 detail::adopt_output(loss_per_sample, pred.device);
651 6 v.mse_vec_per_sample(pred, target, dPred, loss_per_sample);
652 6 }
653
654 12 void softmax_xent_fused_batched(const Tensor& logits_BL, const Tensor& target_BL,
655 const float* d_mask_BL,
656 const int* d_head_offsets, int n_heads,
657 Tensor& probs_BL, Tensor& dLogits_BL,
658 Tensor& loss_per_sample) {
659 12 const auto& v = detail::dispatch(logits_BL, target_BL, probs_BL, dLogits_BL, loss_per_sample);
660
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.softmax_xent_fused_batched)
661 ✗ detail::throw_not_implemented("softmax_xent_fused_batched", logits_BL.device);
662 12 detail::adopt_output(probs_BL, logits_BL.device);
663 12 detail::adopt_output(dLogits_BL, logits_BL.device);
664 12 detail::adopt_output(loss_per_sample, logits_BL.device);
665 24 v.softmax_xent_fused_batched(logits_BL, target_BL, d_mask_BL,
666 12 d_head_offsets, n_heads,
667 12 probs_BL, dLogits_BL, loss_per_sample);
668 12 }
669
670 17 void bce_with_logits_fused_batched(const Tensor& logits_BL, const Tensor& target_BL,
671 const float* d_mask_BL,
672 float pos_weight,
673 Tensor& probs_BL, Tensor& dLogits_BL,
674 Tensor& loss_per_sample) {
675 17 const auto& v = detail::dispatch(logits_BL, target_BL, probs_BL, dLogits_BL, loss_per_sample);
676
1/2
✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
 17 if (!v.bce_with_logits_fused_batched)
677 ✗ detail::throw_not_implemented("bce_with_logits_fused_batched", logits_BL.device);
678 17 detail::adopt_output(probs_BL, logits_BL.device);
679 17 detail::adopt_output(dLogits_BL, logits_BL.device);
680 17 detail::adopt_output(loss_per_sample, logits_BL.device);
681 34 v.bce_with_logits_fused_batched(logits_BL, target_BL, d_mask_BL,
682 17 pos_weight,
683 17 probs_BL, dLogits_BL, loss_per_sample);
684 17 }
685
686 } // namespace brotensor
687