GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 90.3% 1778 / 0 / 1968
Functions: 96.6% 254 / 0 / 263
Branches: 49.7% 330 / 0 / 664
src/ops.cpp
Line Branch Exec Source
1 // brotensor public op wrappers.
2 //
3 // One wrapper per public op declared in <brotensor/ops.h>. Each wrapper:
4 // 1. Validates that all operand tensors share a Device (via detail::dispatch
5 // overloads, which throw on mismatch).
6 // 2. Looks up the function pointer in the resolved OpsVTable.
7 // 3. If null, throws "not implemented on <device>".
8 // 4. Otherwise forwards the call.
9 //
10 // Special cases:
11 // * `mse_scalar`, `softmax_xent_segment` have no Tensor operands. They are
12 // host/CPU-only by nature; the wrapper routes directly to the CPU vtable.
13 // Backends other than CPU leave these slots null.
14 // * Ops with optional `const Tensor*` operands dispatch on the first
15 // guaranteed-non-null Tensor (typically the first positional arg).
16 // * Ops taking std::vector<const Tensor*> dispatch on parts[0]->device
17 // (or, for backward variants, on dY).
18 // * `quantize_int8_per_row_host` is a host helper over raw buffers — not in
19 // the vtable. Implemented inline here.
20
21 #include <brotensor/ops.h>
22 #include <brotensor/tensor.h>
23 #include <brotensor/detail/dispatch.h>
24 #include <brotensor/detail/cpu/thread_pool.h>
25
26 #include <algorithm>
27 #include <cmath>
28 #include <cstdint>
29 #include <stdexcept>
30 #include <string>
31 #include <vector>
32
33 namespace brotensor {
34
35 namespace {
36
37 ✗ [[noreturn]] void throw_empty_parts(const char* op) {
38 ✗ std::string m = "brotensor: ";
39 ✗ m += op;
40 ✗ m += ": parts must be non-empty";
41 ✗ throw std::runtime_error(m);
42 ✗ }
43
44 49 const detail::OpsVTable& vtable_from_parts(
45 const std::vector<const Tensor*>& parts, const char* op) {
46
1/2
✓ Branch 0 taken 49 times.
✗ Branch 1 not taken.
 49 if (parts.empty() || parts[0] == nullptr) throw_empty_parts(op);
47 49 Device d = parts[0]->device;
48
2/2
✓ Branch 0 taken 74 times.
✓ Branch 1 taken 49 times.
 123 for (std::size_t i = 1; i < parts.size(); ++i) {
49
2/4
✓ Branch 0 taken 74 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 74 times.
✗ Branch 3 not taken.
 74 if (parts[i] && parts[i]->device != d) {
50 ✗ std::string m = "brotensor: ";
51 ✗ m += op;
52 ✗ m += ": parts[";
53 ✗ m += std::to_string(i);
54 ✗ m += "] on different device";
55 ✗ throw std::runtime_error(m);
56 ✗ }
57 74 }
58 49 return detail::ops_for(d);
59 ✗ }
60
61 const detail::OpsVTable& vtable_from_parts(
62 const std::vector<Tensor*>& parts, const char* op) {
63 if (parts.empty() || parts[0] == nullptr) throw_empty_parts(op);
64 Device d = parts[0]->device;
65 for (std::size_t i = 1; i < parts.size(); ++i) {
66 if (parts[i] && parts[i]->device != d) {
67 std::string m = "brotensor: ";
68 m += op;
69 m += ": parts[";
70 m += std::to_string(i);
71 m += "] on different device";
72 throw std::runtime_error(m);
73 }
74 }
75 return detail::ops_for(d);
76 }
77
78 #define DISPATCH_REQUIRE(opname, vt) \
79 do { \
80 if (!(vt).opname) detail::throw_not_implemented(#opname, _disp_dev); \
81 } while (0)
82
83 } // namespace
84
85 // Helper to extract device for use after dispatch():
86 // const auto& ops = detail::dispatch(...);
87 // Device _disp_dev = X.device; // for DISPATCH_REQUIRE message
88 //
89 // We forgo that pattern and inline the throw with a fresh device argument
90 // where needed.
91
92 // ─── Dense layers + elementwise ────────────────────────────────────────────
93
94 1264 void linear_forward(const Tensor& W, const Tensor& b,
95 const Tensor& x, Tensor& y) {
96 1264 const auto& v = detail::dispatch(W, b, x, y);
97
1/2
✓ Branch 0 taken 1264 times.
✗ Branch 1 not taken.
 1264 if (!v.linear_forward) detail::throw_not_implemented("linear_forward", W.device);
98 1264 detail::adopt_output(y, W.device);
99 1264 v.linear_forward(W, b, x, y);
100 1264 }
101
102 30 void linear_backward(const Tensor& W, const Tensor& x, const Tensor& dY,
103 Tensor& dX, Tensor& dW, Tensor& dB) {
104 30 const auto& v = detail::dispatch(W, x, dY, dX, dW, dB);
105
1/2
✓ Branch 0 taken 30 times.
✗ Branch 1 not taken.
 30 if (!v.linear_backward) detail::throw_not_implemented("linear_backward", W.device);
106 30 detail::adopt_output(dX, W.device);
107 30 detail::adopt_output(dW, W.device);
108 30 detail::adopt_output(dB, W.device);
109 30 v.linear_backward(W, x, dY, dX, dW, dB);
110 30 }
111
112 93 void relu_forward(const Tensor& x, Tensor& y) {
113 93 const auto& v = detail::dispatch(x, y);
114
1/2
✓ Branch 0 taken 93 times.
✗ Branch 1 not taken.
 93 if (!v.relu_forward) detail::throw_not_implemented("relu_forward", x.device);
115 93 detail::adopt_output(y, x.device);
116 93 v.relu_forward(x, y);
117 93 }
118
119 11 void relu_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
120 11 const auto& v = detail::dispatch(x, dY, dX);
121
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.relu_backward) detail::throw_not_implemented("relu_backward", x.device);
122 11 detail::adopt_output(dX, x.device);
123 11 v.relu_backward(x, dY, dX);
124 11 }
125
126 97 void tanh_forward(const Tensor& x, Tensor& y) {
127 97 const auto& v = detail::dispatch(x, y);
128
1/2
✓ Branch 0 taken 97 times.
✗ Branch 1 not taken.
 97 if (!v.tanh_forward) detail::throw_not_implemented("tanh_forward", x.device);
129 97 detail::adopt_output(y, x.device);
130 97 v.tanh_forward(x, y);
131 97 }
132
133 11 void tanh_backward(const Tensor& y, const Tensor& dY, Tensor& dX) {
134 11 const auto& v = detail::dispatch(y, dY, dX);
135
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.tanh_backward) detail::throw_not_implemented("tanh_backward", y.device);
136 11 detail::adopt_output(dX, y.device);
137 11 v.tanh_backward(y, dY, dX);
138 11 }
139
140 26 void sigmoid_forward(const Tensor& x, Tensor& y) {
141 26 const auto& v = detail::dispatch(x, y);
142
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 26 if (!v.sigmoid_forward) detail::throw_not_implemented("sigmoid_forward", x.device);
143 26 detail::adopt_output(y, x.device);
144 26 v.sigmoid_forward(x, y);
145 26 }
146
147 11 void sigmoid_backward(const Tensor& y, const Tensor& dY, Tensor& dX) {
148 11 const auto& v = detail::dispatch(y, dY, dX);
149
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.sigmoid_backward) detail::throw_not_implemented("sigmoid_backward", y.device);
150 11 detail::adopt_output(dX, y.device);
151 11 v.sigmoid_backward(y, dY, dX);
152 11 }
153
154 363 void add_inplace(Tensor& y, const Tensor& x) {
155 363 const auto& v = detail::dispatch(y, x);
156
1/2
✓ Branch 0 taken 363 times.
✗ Branch 1 not taken.
 363 if (!v.add_inplace) detail::throw_not_implemented("add_inplace", y.device);
157 363 detail::adopt_output(y, y.device);
158 363 v.add_inplace(y, x);
159 363 }
160
161 2 void axpby_inplace(Tensor& y, const Tensor& x, float a, float b) {
162 2 const auto& v = detail::dispatch(y, x);
163
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!v.axpby_inplace) detail::throw_not_implemented("axpby_inplace", y.device);
164 2 detail::adopt_output(y, y.device);
165 2 v.axpby_inplace(y, x, a, b);
166 2 }
167
168 15 void add_scalar_inplace(Tensor& y, float s) {
169 15 const auto& v = detail::dispatch(y);
170
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 if (!v.add_scalar_inplace) detail::throw_not_implemented("add_scalar_inplace", y.device);
171 15 detail::adopt_output(y, y.device);
172 15 v.add_scalar_inplace(y, s);
173 15 }
174
175 36 void add_channel_bias_inplace(Tensor& y, const Tensor& bias, int C, int L) {
176 36 const auto& v = detail::dispatch(y);
177
1/2
✓ Branch 0 taken 36 times.
✗ Branch 1 not taken.
 36 if (!v.add_channel_bias_inplace)
178 ✗ detail::throw_not_implemented("add_channel_bias_inplace", y.device);
179 36 detail::adopt_output(y, y.device);
180 36 v.add_channel_bias_inplace(y, bias, C, L);
181 36 }
182
183 298 void scale_inplace(Tensor& y, float s) {
184 298 const auto& v = detail::dispatch(y);
185
1/2
✓ Branch 0 taken 298 times.
✗ Branch 1 not taken.
 298 if (!v.scale_inplace) detail::throw_not_implemented("scale_inplace", y.device);
186 298 detail::adopt_output(y, y.device);
187 298 v.scale_inplace(y, s);
188 298 }
189
190 12 void clamp(Tensor& y, float lo, float hi) {
191 12 const auto& v = detail::dispatch(y);
192
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.clamp) detail::throw_not_implemented("clamp", y.device);
193 12 detail::adopt_output(y, y.device);
194 12 v.clamp(y, lo, hi);
195 12 }
196
197 159 void mul_inplace(Tensor& y, const Tensor& x) {
198 159 const auto& v = detail::dispatch(y, x);
199
1/2
✓ Branch 0 taken 159 times.
✗ Branch 1 not taken.
 159 if (!v.mul_inplace) detail::throw_not_implemented("mul_inplace", y.device);
200 159 detail::adopt_output(y, y.device);
201 159 v.mul_inplace(y, x);
202 159 }
203
204 12 void modulate(const Tensor& X, const Tensor& scale, const Tensor& shift,
205 Tensor& Y) {
206 12 const auto& v = detail::dispatch(X, scale, shift, Y);
207
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.modulate) detail::throw_not_implemented("modulate", X.device);
208 12 detail::adopt_output(Y, X.device);
209 12 v.modulate(X, scale, shift, Y);
210 12 }
211
212 10 void broadcast_mul(const Tensor& X, const Tensor& vv, Tensor& Y) {
213 10 const auto& v = detail::dispatch(X, vv, Y);
214
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.broadcast_mul) detail::throw_not_implemented("broadcast_mul", X.device);
215 10 detail::adopt_output(Y, X.device);
216 10 v.broadcast_mul(X, vv, Y);
217 10 }
218
219 8 void build_slot_mask(const Tensor& x, int offset, int K, int stride, Tensor& mask) {
220 8 const auto& v = detail::dispatch(x, mask);
221
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);
222 8 detail::adopt_output(mask, x.device);
223 8 v.build_slot_mask(x, offset, K, stride, mask);
224 8 }
225
226 // ─── Reductions / norm / softmax / attention (training) ────────────────────
227
228 35 void softmax_forward(const Tensor& logits, Tensor& probs, const float* mask) {
229 35 const auto& v = detail::dispatch(logits, probs);
230
1/2
✓ Branch 0 taken 35 times.
✗ Branch 1 not taken.
 35 if (!v.softmax_forward) detail::throw_not_implemented("softmax_forward", logits.device);
231 35 detail::adopt_output(probs, logits.device);
232 35 v.softmax_forward(logits, probs, mask);
233 35 }
234
235 44 void softmax_rows_forward(const Tensor& X, Tensor& Y, int rows, int cols) {
236 44 const auto& v = detail::dispatch(X, Y);
237
1/2
✓ Branch 0 taken 44 times.
✗ Branch 1 not taken.
 44 if (!v.softmax_rows_forward) detail::throw_not_implemented("softmax_rows_forward", X.device);
238 44 detail::adopt_output(Y, X.device);
239 44 v.softmax_rows_forward(X, Y, rows, cols);
240 44 }
241
242 21 void softmax_backward(const Tensor& probs, const Tensor& dProbs, Tensor& dLogits) {
243 21 const auto& v = detail::dispatch(probs, dProbs, dLogits);
244
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 if (!v.softmax_backward) detail::throw_not_implemented("softmax_backward", probs.device);
245 21 detail::adopt_output(dLogits, probs.device);
246 21 v.softmax_backward(probs, dProbs, dLogits);
247 21 }
248
249 6 void layernorm_forward(const Tensor& x, const Tensor& gamma, const Tensor& beta,
250 Tensor& y, Tensor& xhat,
251 float& mean_out, float& rstd_out, float eps) {
252 6 const auto& v = detail::dispatch(x, gamma, beta, y, xhat);
253
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.layernorm_forward) detail::throw_not_implemented("layernorm_forward", x.device);
254 6 detail::adopt_output(y, x.device);
255 6 detail::adopt_output(xhat, x.device);
256 6 v.layernorm_forward(x, gamma, beta, y, xhat, mean_out, rstd_out, eps);
257 6 }
258
259 12 void layernorm_backward(const Tensor& dY, const Tensor& xhat,
260 const Tensor& gamma, float rstd,
261 Tensor& dX, Tensor& dGamma, Tensor& dBeta) {
262 12 const auto& v = detail::dispatch(dY, xhat, gamma, dX, dGamma, dBeta);
263
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.layernorm_backward) detail::throw_not_implemented("layernorm_backward", dY.device);
264 12 detail::adopt_output(dX, dY.device);
265 12 detail::adopt_output(dGamma, dY.device);
266 12 detail::adopt_output(dBeta, dY.device);
267 12 v.layernorm_backward(dY, xhat, gamma, rstd, dX, dGamma, dBeta);
268 12 }
269
270 14 void attention_forward(const Tensor& X,
271 const Tensor& Wq, const Tensor& Wk,
272 const Tensor& Wv, const Tensor& Wo,
273 const float* d_mask,
274 Tensor& Q, Tensor& K, Tensor& V,
275 Tensor& Attn, Tensor& Y_pre_Wo, Tensor& O) {
276 14 const auto& v = detail::dispatch(X, Wq, Wk, Wv, Wo, Q, K, V);
277
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.attention_forward) detail::throw_not_implemented("attention_forward", X.device);
278 14 detail::adopt_output(Q, X.device);
279 14 detail::adopt_output(K, X.device);
280 14 detail::adopt_output(V, X.device);
281 14 detail::adopt_output(Attn, X.device);
282 14 detail::adopt_output(Y_pre_Wo, X.device);
283 14 detail::adopt_output(O, X.device);
284 14 v.attention_forward(X, Wq, Wk, Wv, Wo, d_mask, Q, K, V, Attn, Y_pre_Wo, O);
285 14 }
286
287 14 void attention_backward(const Tensor& dO, const Tensor& X,
288 const Tensor& Q, const Tensor& K,
289 const Tensor& V, const Tensor& Attn,
290 const Tensor& Y_pre_Wo,
291 const Tensor& Wq, const Tensor& Wk,
292 const Tensor& Wv, const Tensor& Wo,
293 const float* d_mask,
294 Tensor& dX,
295 Tensor& dWq, Tensor& dWk,
296 Tensor& dWv, Tensor& dWo) {
297 14 const auto& v = detail::dispatch(dO, X, Q, K, V, Attn, Y_pre_Wo);
298
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.attention_backward) detail::throw_not_implemented("attention_backward", dO.device);
299 14 detail::adopt_output(dX, dO.device);
300 14 detail::adopt_output(dWq, dO.device);
301 14 detail::adopt_output(dWk, dO.device);
302 14 detail::adopt_output(dWv, dO.device);
303 14 detail::adopt_output(dWo, dO.device);
304 28 v.attention_backward(dO, X, Q, K, V, Attn, Y_pre_Wo,
305 14 Wq, Wk, Wv, Wo, d_mask,
306 14 dX, dWq, dWk, dWv, dWo);
307 14 }
308
309 26 void mha_forward(const Tensor& X,
310 const Tensor& Wq, const Tensor& Wk,
311 const Tensor& Wv, const Tensor& Wo,
312 const Tensor* bq, const Tensor* bk,
313 const Tensor* bv, const Tensor* bo,
314 const float* d_mask, int num_heads,
315 Tensor& Qh, Tensor& Kh, Tensor& Vh,
316 Tensor& Attnh, Tensor& Yconcat, Tensor& O) {
317 26 const auto& v = detail::dispatch_with_opts(
318 26 X, Wq, {&Wk, &Wv, &Wo, bq, bk, bv, bo, &Qh, &Kh, &Vh});
319
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 26 if (!v.mha_forward) detail::throw_not_implemented("mha_forward", X.device);
320 26 detail::adopt_output(Qh, X.device);
321 26 detail::adopt_output(Kh, X.device);
322 26 detail::adopt_output(Vh, X.device);
323 26 detail::adopt_output(Attnh, X.device);
324 26 detail::adopt_output(Yconcat, X.device);
325 26 detail::adopt_output(O, X.device);
326 52 v.mha_forward(X, Wq, Wk, Wv, Wo, bq, bk, bv, bo, d_mask, num_heads,
327 26 Qh, Kh, Vh, Attnh, Yconcat, O);
328 26 }
329
330 24 void mha_backward(const Tensor& dO, const Tensor& X,
331 const Tensor& Qh, const Tensor& Kh,
332 const Tensor& Vh, const Tensor& Attnh,
333 const Tensor& Yconcat,
334 const Tensor& Wq, const Tensor& Wk,
335 const Tensor& Wv, const Tensor& Wo,
336 const float* d_mask, int num_heads,
337 Tensor& dX,
338 Tensor& dWq, Tensor& dWk,
339 Tensor& dWv, Tensor& dWo,
340 Tensor* dbq, Tensor* dbk,
341 Tensor* dbv, Tensor* dbo) {
342 24 const auto& v = detail::dispatch_with_opts(
343 24 dO, X, {&Qh, &Kh, &Vh, &Attnh, &Yconcat, dbq, dbk, dbv, dbo});
344
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 if (!v.mha_backward) detail::throw_not_implemented("mha_backward", dO.device);
345 24 detail::adopt_output(dX, dO.device);
346 24 detail::adopt_output(dWq, dO.device);
347 24 detail::adopt_output(dWk, dO.device);
348 24 detail::adopt_output(dWv, dO.device);
349 24 detail::adopt_output(dWo, dO.device);
350
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 8 times.
 24 if (dbq) detail::adopt_output(*dbq, dO.device);
351
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 8 times.
 24 if (dbk) detail::adopt_output(*dbk, dO.device);
352
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 8 times.
 24 if (dbv) detail::adopt_output(*dbv, dO.device);
353
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 8 times.
 24 if (dbo) detail::adopt_output(*dbo, dO.device);
354 48 v.mha_backward(dO, X, Qh, Kh, Vh, Attnh, Yconcat,
355 24 Wq, Wk, Wv, Wo, d_mask, num_heads,
356 24 dX, dWq, dWk, dWv, dWo,
357 24 dbq, dbk, dbv, dbo);
358 24 }
359
360 // ─── Pooling / losses / embedding / concat ─────────────────────────────────
361
362 18 void masked_mean_pool_forward(const Tensor& X, const float* d_mask, Tensor& y) {
363 18 const auto& v = detail::dispatch(X, y);
364
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 if (!v.masked_mean_pool_forward) detail::throw_not_implemented("masked_mean_pool_forward", X.device);
365 18 detail::adopt_output(y, X.device);
366 18 v.masked_mean_pool_forward(X, d_mask, y);
367 18 }
368
369 14 void masked_mean_pool_backward(const Tensor& dY, const float* d_mask,
370 int K, Tensor& dX) {
371 14 const auto& v = detail::dispatch(dY, dX);
372
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.masked_mean_pool_backward) detail::throw_not_implemented("masked_mean_pool_backward", dY.device);
373 14 detail::adopt_output(dX, dY.device);
374 14 v.masked_mean_pool_backward(dY, d_mask, K, dX);
375 14 }
376
377 9 float mse_vec_forward(const Tensor& pred, const Tensor& target) {
378 9 const auto& v = detail::dispatch(pred, target);
379
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);
380 9 return v.mse_vec_forward(pred, target);
381 }
382
383 9 void mse_vec_backward(const Tensor& pred, const Tensor& target, Tensor& dPred) {
384 9 const auto& v = detail::dispatch(pred, target, dPred);
385
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);
386 9 detail::adopt_output(dPred, pred.device);
387 9 v.mse_vec_backward(pred, target, dPred);
388 9 }
389
390 // CPU-only host helper. No tensors involved — routes through the CPU vtable.
391 3 float mse_scalar(float pred, float target, float& dPred) {
392 3 const auto& v = detail::ops_for(Device::CPU);
393
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.mse_scalar) detail::throw_not_implemented("mse_scalar", Device::CPU);
394 3 return v.mse_scalar(pred, target, dPred);
395 }
396
397 3 float softmax_xent(const Tensor& logits, const Tensor& target,
398 Tensor& probs, Tensor& dLogits, const float* mask) {
399 3 const auto& v = detail::dispatch(logits, target, probs, dLogits);
400
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.softmax_xent) detail::throw_not_implemented("softmax_xent", logits.device);
401 3 detail::adopt_output(probs, logits.device);
402 3 detail::adopt_output(dLogits, logits.device);
403 3 return v.softmax_xent(logits, target, probs, dLogits, mask);
404 }
405
406 // CPU-only host helper. Raw pointers — CPU vtable directly.
407 7 float softmax_xent_segment(const float* logits, const float* target,
408 float* probs, float* dLogits,
409 int n, const float* mask) {
410 7 const auto& v = detail::ops_for(Device::CPU);
411
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (!v.softmax_xent_segment) detail::throw_not_implemented("softmax_xent_segment", Device::CPU);
412 7 return v.softmax_xent_segment(logits, target, probs, dLogits, n, mask);
413 }
414
415 16 float softmax_xent_fused(const Tensor& logits, const Tensor& target,
416 const float* d_mask,
417 Tensor& probs, Tensor& dLogits) {
418 16 const auto& v = detail::dispatch(logits, target, probs, dLogits);
419
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.softmax_xent_fused) detail::throw_not_implemented("softmax_xent_fused", logits.device);
420 16 detail::adopt_output(probs, logits.device);
421 16 detail::adopt_output(dLogits, logits.device);
422 16 return v.softmax_xent_fused(logits, target, d_mask, probs, dLogits);
423 }
424
425 24 void embedding_lookup_forward(const Tensor& table, const int32_t* d_idx,
426 int B, Tensor& out) {
427 24 const auto& v = detail::dispatch(table, out);
428
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 if (!v.embedding_lookup_forward) detail::throw_not_implemented("embedding_lookup_forward", table.device);
429 24 detail::adopt_output(out, table.device);
430 24 v.embedding_lookup_forward(table, d_idx, B, out);
431 24 }
432
433 25 void embedding_lookup_backward(const Tensor& dOut, const int32_t* d_idx,
434 int B, Tensor& dTable) {
435 25 const auto& v = detail::dispatch(dOut, dTable);
436
1/2
✓ Branch 0 taken 25 times.
✗ Branch 1 not taken.
 25 if (!v.embedding_lookup_backward) detail::throw_not_implemented("embedding_lookup_backward", dOut.device);
437 25 detail::adopt_output(dTable, dOut.device);
438 25 v.embedding_lookup_backward(dOut, d_idx, B, dTable);
439 25 }
440
441 9 void concat_rows(const std::vector<const Tensor*>& parts, Tensor& out) {
442 9 const auto& v = vtable_from_parts(parts, "concat_rows");
443
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 9 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 9 if (out.data != nullptr && parts[0]->device != out.device) {
444 ✗ throw std::runtime_error("brotensor: concat_rows: out on different device");
445 }
446 9 detail::adopt_output(out, parts[0]->device);
447
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.concat_rows) detail::throw_not_implemented("concat_rows", parts[0]->device);
448 9 v.concat_rows(parts, out);
449 9 }
450
451 9 void split_rows(const Tensor& in, const std::vector<Tensor*>& parts) {
452 9 Device d = in.device;
453
2/2
✓ Branch 0 taken 33 times.
✓ Branch 1 taken 9 times.
 42 for (std::size_t i = 0; i < parts.size(); ++i) {
454
3/6
✓ Branch 0 taken 33 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 33 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 33 times.
✗ Branch 5 not taken.
 33 if (parts[i] && parts[i]->data != nullptr && parts[i]->device != d) {
455 ✗ throw std::runtime_error("brotensor: split_rows: parts on different device");
456 }
457 33 }
458 9 const auto& v = detail::ops_for(d);
459
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.split_rows) detail::throw_not_implemented("split_rows", d);
460
3/4
✓ Branch 0 taken 33 times.
✓ Branch 1 taken 9 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 33 times.
 42 for (Tensor* _p : parts) if (_p) detail::adopt_output(*_p, d);
461 9 v.split_rows(in, parts);
462 9 }
463
464 16 void concat_batched_rows(const std::vector<const Tensor*>& parts, Tensor& out) {
465 16 const auto& v = vtable_from_parts(parts, "concat_batched_rows");
466
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 16 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 16 if (out.data != nullptr && parts[0]->device != out.device) {
467 ✗ throw std::runtime_error("brotensor: concat_batched_rows: out on different device");
468 }
469 16 detail::adopt_output(out, parts[0]->device);
470
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.concat_batched_rows) detail::throw_not_implemented("concat_batched_rows", parts[0]->device);
471 16 v.concat_batched_rows(parts, out);
472 16 }
473
474 24 void concat_nchw_channels(const std::vector<const Tensor*>& parts,
475 int N, int H, int W,
476 const std::vector<int>& C_per_part,
477 Tensor& out) {
478 24 const auto& v = vtable_from_parts(parts, "concat_nchw_channels");
479
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 24 if (out.data != nullptr && parts[0]->device != out.device) {
480 ✗ throw std::runtime_error("brotensor: concat_nchw_channels: out on different device");
481 }
482 24 detail::adopt_output(out, parts[0]->device);
483
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 if (!v.concat_nchw_channels) detail::throw_not_implemented("concat_nchw_channels", parts[0]->device);
484 24 v.concat_nchw_channels(parts, N, H, W, C_per_part, out);
485 24 }
486
487 20 void concat_nchw_channels_backward(const Tensor& dY, int N, int H, int W,
488 const std::vector<int>& C_per_part,
489 const std::vector<Tensor*>& parts) {
490 20 Device d = dY.device;
491
2/2
✓ Branch 0 taken 40 times.
✓ Branch 1 taken 20 times.
 60 for (std::size_t i = 0; i < parts.size(); ++i) {
492
2/6
✓ Branch 0 taken 40 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 40 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 40 if (parts[i] && parts[i]->data != nullptr && parts[i]->device != d) {
493 ✗ throw std::runtime_error("brotensor: concat_nchw_channels_backward: parts on different device");
494 }
495 40 }
496 20 const auto& v = detail::ops_for(d);
497
1/2
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
 20 if (!v.concat_nchw_channels_backward) detail::throw_not_implemented("concat_nchw_channels_backward", d);
498
3/4
✓ Branch 0 taken 40 times.
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 40 times.
 60 for (Tensor* _p : parts) if (_p) detail::adopt_output(*_p, d);
499 20 v.concat_nchw_channels_backward(dY, N, H, W, C_per_part, parts);
500 20 }
501
502 8 void copy_d2d(const Tensor& src, int src_off, Tensor& dst, int dst_off, int n) {
503 8 const auto& v = detail::dispatch(src, dst);
504
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.copy_d2d) detail::throw_not_implemented("copy_d2d", src.device);
505 8 detail::adopt_output(dst, src.device);
506 8 v.copy_d2d(src, src_off, dst, dst_off, n);
507 8 }
508
509 8 void copy_d2d_strided(const Tensor& src, int src_off, int src_pitch,
510 Tensor& dst, int dst_off, int dst_pitch,
511 int width, int height) {
512
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (width > src_pitch || width > dst_pitch) {
513 ✗ throw std::runtime_error("brotensor: copy_d2d_strided: width exceeds pitch");
514 }
515 8 const auto& v = detail::dispatch(src, dst);
516
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.copy_d2d_strided) detail::throw_not_implemented("copy_d2d_strided", src.device);
517 8 detail::adopt_output(dst, src.device);
518 8 v.copy_d2d_strided(src, src_off, src_pitch, dst, dst_off, dst_pitch, width, height);
519 8 }
520
521 24 void cast(const Tensor& src, Tensor& dst, Dtype out_dtype) {
522 24 const auto& v = detail::dispatch(src, dst);
523
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 if (!v.cast) detail::throw_not_implemented("cast", src.device);
524 24 detail::adopt_output(dst, src.device);
525 24 v.cast(src, dst, out_dtype);
526 24 }
527
528 // ─── Inference batched + optim ─────────────────────────────────────────────
529
530 26 void layernorm_forward_inference_batched(const Tensor& X_RD,
531 const Tensor& gamma,
532 const Tensor& beta,
533 Tensor& Y_RD, float eps) {
534 26 const auto& v = detail::dispatch(X_RD, gamma, beta, Y_RD);
535
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 26 if (!v.layernorm_forward_inference_batched)
536 ✗ detail::throw_not_implemented("layernorm_forward_inference_batched", X_RD.device);
537 26 detail::adopt_output(Y_RD, X_RD.device);
538 26 v.layernorm_forward_inference_batched(X_RD, gamma, beta, Y_RD, eps);
539 26 }
540
541 16 void layernorm_forward_batched_with_caches(const Tensor& X_RD,
542 const Tensor& gamma,
543 const Tensor& beta,
544 Tensor& Y_RD, Tensor& Xhat_RD,
545 Tensor& Mean_R, Tensor& Rstd_R,
546 float eps) {
547 16 const auto& v = detail::dispatch(X_RD, gamma, beta, Y_RD, Xhat_RD, Mean_R, Rstd_R);
548
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.layernorm_forward_batched_with_caches)
549 ✗ detail::throw_not_implemented("layernorm_forward_batched_with_caches", X_RD.device);
550 16 detail::adopt_output(Y_RD, X_RD.device);
551 16 detail::adopt_output(Xhat_RD, X_RD.device);
552 16 detail::adopt_output(Mean_R, X_RD.device);
553 16 detail::adopt_output(Rstd_R, X_RD.device);
554 32 v.layernorm_forward_batched_with_caches(X_RD, gamma, beta, Y_RD, Xhat_RD,
555 16 Mean_R, Rstd_R, eps);
556 16 }
557
558 16 void layernorm_backward_batched_with_caches(const Tensor& dY_RD,
559 const Tensor& Xhat_RD,
560 const Tensor& gamma,
561 const Tensor& Rstd_R,
562 Tensor& dX_RD,
563 Tensor& dGamma, Tensor& dBeta) {
564 16 const auto& v = detail::dispatch(dY_RD, Xhat_RD, gamma, Rstd_R, dX_RD, dGamma, dBeta);
565
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.layernorm_backward_batched_with_caches)
566 ✗ detail::throw_not_implemented("layernorm_backward_batched_with_caches", dY_RD.device);
567 16 detail::adopt_output(dX_RD, dY_RD.device);
568 16 detail::adopt_output(dGamma, dY_RD.device);
569 16 detail::adopt_output(dBeta, dY_RD.device);
570 32 v.layernorm_backward_batched_with_caches(dY_RD, Xhat_RD, gamma, Rstd_R,
571 16 dX_RD, dGamma, dBeta);
572 16 }
573
574 5 void sgd_step(Tensor& param, Tensor& grad, Tensor& velocity,
575 float lr, float momentum) {
576 5 const auto& v = detail::dispatch(param, grad, velocity);
577
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 if (!v.sgd_step) detail::throw_not_implemented("sgd_step", param.device);
578 5 detail::adopt_output(param, param.device);
579 5 detail::adopt_output(grad, param.device);
580 5 detail::adopt_output(velocity, param.device);
581 5 v.sgd_step(param, grad, velocity, lr, momentum);
582 5 }
583
584 50 void adam_step(Tensor& param, const Tensor& grad,
585 Tensor& m, Tensor& v_buf,
586 float lr, float beta1, float beta2, float eps, int step) {
587 50 const auto& v = detail::dispatch(param, grad, m, v_buf);
588
1/2
✓ Branch 0 taken 50 times.
✗ Branch 1 not taken.
 50 if (!v.adam_step) detail::throw_not_implemented("adam_step", param.device);
589 50 detail::adopt_output(param, param.device);
590 50 detail::adopt_output(m, param.device);
591 50 detail::adopt_output(v_buf, param.device);
592 50 v.adam_step(param, grad, m, v_buf, lr, beta1, beta2, eps, step);
593 50 }
594
595 3 void xavier_init(Tensor& W, uint64_t& rng_state) {
596 3 const auto& v = detail::dispatch(W);
597
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.xavier_init) detail::throw_not_implemented("xavier_init", W.device);
598 3 detail::adopt_output(W, W.device);
599 3 v.xavier_init(W, rng_state);
600 3 }
601
602 // ─── Batched inference variants ────────────────────────────────────────────
603
604 28 void linear_forward_batched(const Tensor& W, const Tensor& bias,
605 const Tensor& X_BD, Tensor& Y_BD) {
606 28 const auto& v = detail::dispatch(W, bias, X_BD, Y_BD);
607
1/2
✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
 28 if (!v.linear_forward_batched) detail::throw_not_implemented("linear_forward_batched", W.device);
608 28 detail::adopt_output(Y_BD, W.device);
609 28 v.linear_forward_batched(W, bias, X_BD, Y_BD);
610 28 }
611
612 3 void relu_forward_batched(const Tensor& X_BD, Tensor& Y_BD) {
613 3 const auto& v = detail::dispatch(X_BD, Y_BD);
614
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.relu_forward_batched) detail::throw_not_implemented("relu_forward_batched", X_BD.device);
615 3 detail::adopt_output(Y_BD, X_BD.device);
616 3 v.relu_forward_batched(X_BD, Y_BD);
617 3 }
618
619 9 void tanh_forward_batched(const Tensor& X_BD, Tensor& Y_BD) {
620 9 const auto& v = detail::dispatch(X_BD, Y_BD);
621
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.tanh_forward_batched) detail::throw_not_implemented("tanh_forward_batched", X_BD.device);
622 9 detail::adopt_output(Y_BD, X_BD.device);
623 9 v.tanh_forward_batched(X_BD, Y_BD);
624 9 }
625
626 3 void add_inplace_batched(Tensor& Y_BD, const Tensor& X_BD) {
627 3 const auto& v = detail::dispatch(Y_BD, X_BD);
628
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.add_inplace_batched) detail::throw_not_implemented("add_inplace_batched", Y_BD.device);
629 3 detail::adopt_output(Y_BD, Y_BD.device);
630 3 v.add_inplace_batched(Y_BD, X_BD);
631 3 }
632
633 // ─── Batched backward variants ─────────────────────────────────────────────
634
635 9 void linear_backward_batched(const Tensor& W, const Tensor& X_BD,
636 const Tensor& dY_BD,
637 Tensor& dX_BD, Tensor& dW, Tensor& dB) {
638 9 const auto& v = detail::dispatch(W, X_BD, dY_BD, dX_BD, dW, dB);
639
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.linear_backward_batched) detail::throw_not_implemented("linear_backward_batched", W.device);
640 9 detail::adopt_output(dX_BD, W.device);
641 9 detail::adopt_output(dW, W.device);
642 9 detail::adopt_output(dB, W.device);
643 9 v.linear_backward_batched(W, X_BD, dY_BD, dX_BD, dW, dB);
644 9 }
645
646 6 void relu_backward_batched(const Tensor& X_BD, const Tensor& dY_BD, Tensor& dX_BD) {
647 6 const auto& v = detail::dispatch(X_BD, dY_BD, dX_BD);
648
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.relu_backward_batched) detail::throw_not_implemented("relu_backward_batched", X_BD.device);
649 6 detail::adopt_output(dX_BD, X_BD.device);
650 6 v.relu_backward_batched(X_BD, dY_BD, dX_BD);
651 6 }
652
653 6 void tanh_backward_batched(const Tensor& Y_BD, const Tensor& dY_BD, Tensor& dX_BD) {
654 6 const auto& v = detail::dispatch(Y_BD, dY_BD, dX_BD);
655
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.tanh_backward_batched) detail::throw_not_implemented("tanh_backward_batched", Y_BD.device);
656 6 detail::adopt_output(dX_BD, Y_BD.device);
657 6 v.tanh_backward_batched(Y_BD, dY_BD, dX_BD);
658 6 }
659
660 // ─── Batched per-sample losses ─────────────────────────────────────────────
661
662 6 void mse_vec_per_sample(const Tensor& pred, const Tensor& target,
663 Tensor& dPred, Tensor& loss_per_sample) {
664 6 const auto& v = detail::dispatch(pred, target, dPred, loss_per_sample);
665
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);
666 6 detail::adopt_output(dPred, pred.device);
667 6 detail::adopt_output(loss_per_sample, pred.device);
668 6 v.mse_vec_per_sample(pred, target, dPred, loss_per_sample);
669 6 }
670
671 12 void softmax_xent_fused_batched(const Tensor& logits_BL, const Tensor& target_BL,
672 const float* d_mask_BL,
673 const int* d_head_offsets, int n_heads,
674 Tensor& probs_BL, Tensor& dLogits_BL,
675 Tensor& loss_per_sample) {
676 12 const auto& v = detail::dispatch(logits_BL, target_BL, probs_BL, dLogits_BL, loss_per_sample);
677
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.softmax_xent_fused_batched)
678 ✗ detail::throw_not_implemented("softmax_xent_fused_batched", logits_BL.device);
679 12 detail::adopt_output(probs_BL, logits_BL.device);
680 12 detail::adopt_output(dLogits_BL, logits_BL.device);
681 12 detail::adopt_output(loss_per_sample, logits_BL.device);
682 24 v.softmax_xent_fused_batched(logits_BL, target_BL, d_mask_BL,
683 12 d_head_offsets, n_heads,
684 12 probs_BL, dLogits_BL, loss_per_sample);
685 12 }
686
687 11 void bce_with_logits_fused_batched(const Tensor& logits_BL, const Tensor& target_BL,
688 const float* d_mask_BL,
689 float pos_weight,
690 Tensor& probs_BL, Tensor& dLogits_BL,
691 Tensor& loss_per_sample) {
692 11 const auto& v = detail::dispatch(logits_BL, target_BL, probs_BL, dLogits_BL, loss_per_sample);
693
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.bce_with_logits_fused_batched)
694 ✗ detail::throw_not_implemented("bce_with_logits_fused_batched", logits_BL.device);
695 11 detail::adopt_output(probs_BL, logits_BL.device);
696 11 detail::adopt_output(dLogits_BL, logits_BL.device);
697 11 detail::adopt_output(loss_per_sample, logits_BL.device);
698 22 v.bce_with_logits_fused_batched(logits_BL, target_BL, d_mask_BL,
699 11 pos_weight,
700 11 probs_BL, dLogits_BL, loss_per_sample);
701 11 }
702
703 // ─── Conv2d ────────────────────────────────────────────────────────────────
704
705 44 void conv2d_forward(const Tensor& X, const Tensor& Wt, const Tensor* bias,
706 int N, int C_in, int H, int W,
707 int C_out, int kH, int kW,
708 int stride_h, int stride_w,
709 int pad_h, int pad_w,
710 int dil_h, int dil_w,
711 int groups, Tensor& Y) {
712 44 const auto& v = detail::dispatch_with_opts(X, Wt, {bias, &Y});
713
1/2
✓ Branch 0 taken 44 times.
✗ Branch 1 not taken.
 44 if (!v.conv2d_forward) detail::throw_not_implemented("conv2d_forward", X.device);
714 44 detail::adopt_output(Y, X.device);
715 88 v.conv2d_forward(X, Wt, bias, N, C_in, H, W, C_out, kH, kW,
716 44 stride_h, stride_w, pad_h, pad_w, dil_h, dil_w, groups, Y);
717 44 }
718
719 28 void conv2d_backward_input(const Tensor& Wt, const Tensor& dY,
720 int N, int C_in, int H, int W,
721 int C_out, int kH, int kW,
722 int stride_h, int stride_w,
723 int pad_h, int pad_w,
724 int dil_h, int dil_w,
725 int groups, Tensor& dX) {
726 28 const auto& v = detail::dispatch(Wt, dY, dX);
727
1/2
✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
 28 if (!v.conv2d_backward_input) detail::throw_not_implemented("conv2d_backward_input", Wt.device);
728 28 detail::adopt_output(dX, Wt.device);
729 56 v.conv2d_backward_input(Wt, dY, N, C_in, H, W, C_out, kH, kW,
730 28 stride_h, stride_w, pad_h, pad_w, dil_h, dil_w, groups, dX);
731 28 }
732
733 26 void conv2d_backward_weight(const Tensor& X, const Tensor& dY,
734 int N, int C_in, int H, int W,
735 int C_out, int kH, int kW,
736 int stride_h, int stride_w,
737 int pad_h, int pad_w,
738 int dil_h, int dil_w,
739 int groups, Tensor& dWt) {
740 26 const auto& v = detail::dispatch(X, dY, dWt);
741
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 26 if (!v.conv2d_backward_weight) detail::throw_not_implemented("conv2d_backward_weight", X.device);
742 26 detail::adopt_output(dWt, X.device);
743 52 v.conv2d_backward_weight(X, dY, N, C_in, H, W, C_out, kH, kW,
744 26 stride_h, stride_w, pad_h, pad_w, dil_h, dil_w, groups, dWt);
745 26 }
746
747 16 void conv2d_backward_bias(const Tensor& dY,
748 int N, int C_out, int H_out, int W_out,
749 Tensor& dB) {
750 16 const auto& v = detail::dispatch(dY, dB);
751
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.conv2d_backward_bias) detail::throw_not_implemented("conv2d_backward_bias", dY.device);
752 16 detail::adopt_output(dB, dY.device);
753 16 v.conv2d_backward_bias(dY, N, C_out, H_out, W_out, dB);
754 16 }
755
756 12 void deform_conv2d_forward(const Tensor& X, const Tensor& offset,
757 const Tensor* mask, const Tensor& Wt,
758 const Tensor* bias,
759 int N, int C_in, int H, int W,
760 int C_out, int kH, int kW,
761 int stride_h, int stride_w,
762 int pad_h, int pad_w,
763 int dil_h, int dil_w,
764 int groups, int deform_groups, Tensor& Y) {
765 12 const auto& v = detail::dispatch_with_opts(X, Wt, {&offset, mask, bias, &Y});
766
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.deform_conv2d_forward) detail::throw_not_implemented("deform_conv2d_forward", X.device);
767 12 detail::adopt_output(Y, X.device);
768 24 v.deform_conv2d_forward(X, offset, mask, Wt, bias, N, C_in, H, W,
769 12 C_out, kH, kW, stride_h, stride_w, pad_h, pad_w,
770 12 dil_h, dil_w, groups, deform_groups, Y);
771 12 }
772
773 // ─── Conv3d ────────────────────────────────────────────────────────────────
774
775 37 void conv3d_forward(const Tensor& X, const Tensor& Wt, const Tensor* bias,
776 int N, int C_in, int T, int H, int W,
777 int C_out, int kT, int kH, int kW,
778 int stride_t, int stride_h, int stride_w,
779 int pad_t, int pad_h, int pad_w,
780 int dil_t, int dil_h, int dil_w,
781 int groups, Tensor& Y) {
782 37 const auto& v = detail::dispatch_with_opts(X, Wt, {bias, &Y});
783
1/2
✓ Branch 0 taken 37 times.
✗ Branch 1 not taken.
 37 if (!v.conv3d_forward) detail::throw_not_implemented("conv3d_forward", X.device);
784 37 detail::adopt_output(Y, X.device);
785 74 v.conv3d_forward(X, Wt, bias, N, C_in, T, H, W, C_out, kT, kH, kW,
786 37 stride_t, stride_h, stride_w,
787 37 pad_t, pad_h, pad_w,
788 37 dil_t, dil_h, dil_w, groups, Y);
789 37 }
790
791 6 void conv3d_int8w_fp16_forward(const Tensor& X, const Tensor& W_int8,
792 const Tensor& scales, const Tensor* bias,
793 int N, int C_in, int T, int H, int W,
794 int C_out, int kT, int kH, int kW,
795 int stride_t, int stride_h, int stride_w,
796 int pad_t, int pad_h, int pad_w,
797 int dil_t, int dil_h, int dil_w,
798 int groups, Tensor& Y) {
799 6 const auto& v = detail::dispatch_with_opts(X, W_int8, {&scales, bias, &Y});
800
2/2
✓ Branch 0 taken 5 times.
✓ Branch 1 taken 1 time.
 6 if (!v.conv3d_int8w_fp16_forward)
801 1 detail::throw_not_implemented("conv3d_int8w_fp16_forward", X.device);
802 5 detail::adopt_output(Y, X.device);
803 10 v.conv3d_int8w_fp16_forward(X, W_int8, scales, bias,
804 5 N, C_in, T, H, W, C_out, kT, kH, kW,
805 5 stride_t, stride_h, stride_w,
806 5 pad_t, pad_h, pad_w,
807 5 dil_t, dil_h, dil_w, groups, Y);
808 5 }
809
810 // ─── GroupNorm ─────────────────────────────────────────────────────────────
811
812 18 void group_norm_forward(const Tensor& X, const Tensor& gamma, const Tensor& beta,
813 int N, int C, int H, int W, int num_groups,
814 float eps, Tensor& Y) {
815 18 const auto& v = detail::dispatch(X, gamma, beta, Y);
816
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 if (!v.group_norm_forward) detail::throw_not_implemented("group_norm_forward", X.device);
817 18 detail::adopt_output(Y, X.device);
818 18 v.group_norm_forward(X, gamma, beta, N, C, H, W, num_groups, eps, Y);
819 18 }
820
821 18 void group_norm_backward(const Tensor& X, const Tensor& gamma, const Tensor& dY,
822 int N, int C, int H, int W, int num_groups, float eps,
823 Tensor& dX, Tensor& dGamma, Tensor& dBeta) {
824 18 const auto& v = detail::dispatch(X, gamma, dY, dX, dGamma, dBeta);
825
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 if (!v.group_norm_backward) detail::throw_not_implemented("group_norm_backward", X.device);
826 18 detail::adopt_output(dX, X.device);
827 18 detail::adopt_output(dGamma, X.device);
828 18 detail::adopt_output(dBeta, X.device);
829 36 v.group_norm_backward(X, gamma, dY, N, C, H, W, num_groups, eps,
830 18 dX, dGamma, dBeta);
831 18 }
832
833 // ─── Activations ───────────────────────────────────────────────────────────
834
835 10 void silu_forward(const Tensor& x, Tensor& y) {
836 10 const auto& v = detail::dispatch(x, y);
837
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.silu_forward) detail::throw_not_implemented("silu_forward", x.device);
838 10 detail::adopt_output(y, x.device);
839 10 v.silu_forward(x, y);
840 10 }
841 6 void silu_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
842 6 const auto& v = detail::dispatch(x, dY, dX);
843
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.silu_backward) detail::throw_not_implemented("silu_backward", x.device);
844 6 detail::adopt_output(dX, x.device);
845 6 v.silu_backward(x, dY, dX);
846 6 }
847 10 void gelu_forward(const Tensor& x, Tensor& y) {
848 10 const auto& v = detail::dispatch(x, y);
849
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.gelu_forward) detail::throw_not_implemented("gelu_forward", x.device);
850 10 detail::adopt_output(y, x.device);
851 10 v.gelu_forward(x, y);
852 10 }
853 6 void gelu_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
854 6 const auto& v = detail::dispatch(x, dY, dX);
855
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.gelu_backward) detail::throw_not_implemented("gelu_backward", x.device);
856 6 detail::adopt_output(dX, x.device);
857 6 v.gelu_backward(x, dY, dX);
858 6 }
859 8 void gelu_exact_forward(const Tensor& x, Tensor& y) {
860 8 const auto& v = detail::dispatch(x, y);
861
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.gelu_exact_forward) detail::throw_not_implemented("gelu_exact_forward", x.device);
862 8 detail::adopt_output(y, x.device);
863 8 v.gelu_exact_forward(x, y);
864 8 }
865 6 void gelu_exact_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
866 6 const auto& v = detail::dispatch(x, dY, dX);
867
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.gelu_exact_backward) detail::throw_not_implemented("gelu_exact_backward", x.device);
868 6 detail::adopt_output(dX, x.device);
869 6 v.gelu_exact_backward(x, dY, dX);
870 6 }
871 8 void quick_gelu_forward(const Tensor& x, Tensor& y) {
872 8 const auto& v = detail::dispatch(x, y);
873
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.quick_gelu_forward) detail::throw_not_implemented("quick_gelu_forward", x.device);
874 8 detail::adopt_output(y, x.device);
875 8 v.quick_gelu_forward(x, y);
876 8 }
877 6 void quick_gelu_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
878 6 const auto& v = detail::dispatch(x, dY, dX);
879
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.quick_gelu_backward) detail::throw_not_implemented("quick_gelu_backward", x.device);
880 6 detail::adopt_output(dX, x.device);
881 6 v.quick_gelu_backward(x, dY, dX);
882 6 }
883
884 // ─── Resample ──────────────────────────────────────────────────────────────
885
886 15 void upsample_nearest_2x(const Tensor& X, int N, int C, int H, int W, Tensor& Y) {
887 15 const auto& v = detail::dispatch(X, Y);
888
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 if (!v.upsample_nearest_2x) detail::throw_not_implemented("upsample_nearest_2x", X.device);
889 15 detail::adopt_output(Y, X.device);
890 15 v.upsample_nearest_2x(X, N, C, H, W, Y);
891 15 }
892 15 void upsample_bilinear_2x(const Tensor& X, int N, int C, int H, int W, Tensor& Y) {
893 15 const auto& v = detail::dispatch(X, Y);
894
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 if (!v.upsample_bilinear_2x) detail::throw_not_implemented("upsample_bilinear_2x", X.device);
895 15 detail::adopt_output(Y, X.device);
896 15 v.upsample_bilinear_2x(X, N, C, H, W, Y);
897 15 }
898 14 void downsample_avg_2x(const Tensor& X, int N, int C, int H, int W, Tensor& Y) {
899 14 const auto& v = detail::dispatch(X, Y);
900
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.downsample_avg_2x) detail::throw_not_implemented("downsample_avg_2x", X.device);
901 14 detail::adopt_output(Y, X.device);
902 14 v.downsample_avg_2x(X, N, C, H, W, Y);
903 14 }
904 12 void upsample_nearest_2x_backward(const Tensor& dY, int N, int C, int H, int W, Tensor& dX) {
905 12 const auto& v = detail::dispatch(dY, dX);
906
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.upsample_nearest_2x_backward) detail::throw_not_implemented("upsample_nearest_2x_backward", dY.device);
907 12 detail::adopt_output(dX, dY.device);
908 12 v.upsample_nearest_2x_backward(dY, N, C, H, W, dX);
909 12 }
910 12 void upsample_bilinear_2x_backward(const Tensor& dY, int N, int C, int H, int W, Tensor& dX) {
911 12 const auto& v = detail::dispatch(dY, dX);
912
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.upsample_bilinear_2x_backward) detail::throw_not_implemented("upsample_bilinear_2x_backward", dY.device);
913 12 detail::adopt_output(dX, dY.device);
914 12 v.upsample_bilinear_2x_backward(dY, N, C, H, W, dX);
915 12 }
916 12 void downsample_avg_2x_backward(const Tensor& dY, int N, int C, int H, int W, Tensor& dX) {
917 12 const auto& v = detail::dispatch(dY, dX);
918
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.downsample_avg_2x_backward) detail::throw_not_implemented("downsample_avg_2x_backward", dY.device);
919 12 detail::adopt_output(dX, dY.device);
920 12 v.downsample_avg_2x_backward(dY, N, C, H, W, dX);
921 12 }
922
923 76 void interp2d_forward(const Tensor& X,
924 int N, int C, int H_in, int W_in, int H_out, int W_out,
925 int mode, Tensor& Y) {
926 76 const auto& v = detail::dispatch(X, Y);
927
1/2
✓ Branch 0 taken 76 times.
✗ Branch 1 not taken.
 76 if (!v.interp2d_forward) detail::throw_not_implemented("interp2d_forward", X.device);
928 76 detail::adopt_output(Y, X.device);
929 76 v.interp2d_forward(X, N, C, H_in, W_in, H_out, W_out, mode, Y);
930 76 }
931 13 void interp2d_backward(const Tensor& dY,
932 int N, int C, int H_in, int W_in, int H_out, int W_out,
933 int mode, Tensor& dX) {
934 13 const auto& v = detail::dispatch(dY, dX);
935
1/2
✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
 13 if (!v.interp2d_backward) detail::throw_not_implemented("interp2d_backward", dY.device);
936 13 detail::adopt_output(dX, dY.device);
937 13 v.interp2d_backward(dY, N, C, H_in, W_in, H_out, W_out, mode, dX);
938 13 }
939 18 void interp2d_align_corners_forward(const Tensor& X,
940 int N, int C, int H_in, int W_in,
941 int H_out, int W_out, int mode, Tensor& Y) {
942 18 const auto& v = detail::dispatch(X, Y);
943
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 if (!v.interp2d_align_corners_forward)
944 ✗ detail::throw_not_implemented("interp2d_align_corners_forward", X.device);
945 18 detail::adopt_output(Y, X.device);
946 18 v.interp2d_align_corners_forward(X, N, C, H_in, W_in, H_out, W_out, mode, Y);
947 18 }
948
949 61 void pad2d_forward(const Tensor& X, int N, int C, int H, int W,
950 int pad_top, int pad_bottom, int pad_left, int pad_right,
951 int mode, Tensor& Y) {
952 61 const auto& v = detail::dispatch(X, Y);
953
1/2
✓ Branch 0 taken 61 times.
✗ Branch 1 not taken.
 61 if (!v.pad2d_forward) detail::throw_not_implemented("pad2d_forward", X.device);
954 61 detail::adopt_output(Y, X.device);
955 61 v.pad2d_forward(X, N, C, H, W, pad_top, pad_bottom, pad_left, pad_right, mode, Y);
956 61 }
957 11 void pad2d_backward(const Tensor& dY, int N, int C, int H, int W,
958 int pad_top, int pad_bottom, int pad_left, int pad_right,
959 int mode, Tensor& dX) {
960 11 const auto& v = detail::dispatch(dY, dX);
961
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.pad2d_backward) detail::throw_not_implemented("pad2d_backward", dY.device);
962 11 detail::adopt_output(dX, dY.device);
963 11 v.pad2d_backward(dY, N, C, H, W, pad_top, pad_bottom, pad_left, pad_right, mode, dX);
964 11 }
965
966 12 void slice2d_forward(const Tensor& X, int N, int C, int H, int W,
967 int h0, int w0, int H_out, int W_out, Tensor& Y) {
968 12 const auto& v = detail::dispatch(X, Y);
969
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.slice2d_forward) detail::throw_not_implemented("slice2d_forward", X.device);
970 12 detail::adopt_output(Y, X.device);
971 12 v.slice2d_forward(X, N, C, H, W, h0, w0, H_out, W_out, Y);
972 12 }
973 6 void slice2d_backward(const Tensor& dY, int N, int C, int H, int W,
974 int h0, int w0, int H_out, int W_out, Tensor& dX) {
975 6 const auto& v = detail::dispatch(dY, dX);
976
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.slice2d_backward) detail::throw_not_implemented("slice2d_backward", dY.device);
977 6 detail::adopt_output(dX, dY.device);
978 6 v.slice2d_backward(dY, N, C, H, W, h0, w0, H_out, W_out, dX);
979 6 }
980
981 10 void unfold2d_forward(const Tensor& X, int N, int C, int H, int W,
982 int kH, int kW, int stride_h, int stride_w,
983 int pad_top, int pad_bottom, int pad_left, int pad_right,
984 int mode, Tensor& Y) {
985 10 const auto& v = detail::dispatch(X, Y);
986
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.unfold2d_forward) detail::throw_not_implemented("unfold2d_forward", X.device);
987 10 detail::adopt_output(Y, X.device);
988 20 v.unfold2d_forward(X, N, C, H, W, kH, kW, stride_h, stride_w,
989 10 pad_top, pad_bottom, pad_left, pad_right, mode, Y);
990 10 }
991
992 10 void l2_normalize_nchw_forward(const Tensor& X, int N, int C, int H, int W,
993 float eps, Tensor& Y) {
994 10 const auto& v = detail::dispatch(X, Y);
995
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.l2_normalize_nchw_forward)
996 ✗ detail::throw_not_implemented("l2_normalize_nchw_forward", X.device);
997 10 detail::adopt_output(Y, X.device);
998 10 v.l2_normalize_nchw_forward(X, N, C, H, W, eps, Y);
999 10 }
1000
1001 10 void convex_upsample_forward(const Tensor& X, const Tensor& Mask,
1002 int N, int C, int H, int W, int scale, Tensor& Y) {
1003 10 const auto& v = detail::dispatch(X, Mask, Y);
1004
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.convex_upsample_forward)
1005 ✗ detail::throw_not_implemented("convex_upsample_forward", X.device);
1006 10 detail::adopt_output(Y, X.device);
1007 10 v.convex_upsample_forward(X, Mask, N, C, H, W, scale, Y);
1008 10 }
1009
1010 17 void top_k_rows(const Tensor& X, int k, Tensor& Vals, Tensor& Idx) {
1011 17 const auto& v = detail::dispatch(X, Vals, Idx);
1012
1/2
✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
 17 if (!v.top_k_rows) detail::throw_not_implemented("top_k_rows", X.device);
1013 17 detail::adopt_output(Vals, X.device);
1014 17 detail::adopt_output(Idx, X.device);
1015 17 v.top_k_rows(X, k, Vals, Idx);
1016 17 }
1017
1018 44 void adaptive_avg_pool2d_forward(const Tensor& X, int N, int C, int H, int W,
1019 int H_out, int W_out, Tensor& Y) {
1020 44 const auto& v = detail::dispatch(X, Y);
1021
1/2
✓ Branch 0 taken 44 times.
✗ Branch 1 not taken.
 44 if (!v.adaptive_avg_pool2d_forward)
1022 ✗ detail::throw_not_implemented("adaptive_avg_pool2d_forward", X.device);
1023 44 detail::adopt_output(Y, X.device);
1024 44 v.adaptive_avg_pool2d_forward(X, N, C, H, W, H_out, W_out, Y);
1025 44 }
1026 6 void adaptive_avg_pool2d_backward(const Tensor& dY, int N, int C, int H, int W,
1027 int H_out, int W_out, Tensor& dX) {
1028 6 const auto& v = detail::dispatch(dY, dX);
1029
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.adaptive_avg_pool2d_backward)
1030 ✗ detail::throw_not_implemented("adaptive_avg_pool2d_backward", dY.device);
1031 6 detail::adopt_output(dX, dY.device);
1032 6 v.adaptive_avg_pool2d_backward(dY, N, C, H, W, H_out, W_out, dX);
1033 6 }
1034
1035 19 void max_pool2d_forward(const Tensor& X, int N, int C, int H, int W,
1036 int kH, int kW, int stride_h, int stride_w,
1037 int pad_h, int pad_w, Tensor& Y, Tensor& Idx) {
1038 19 const auto& v = detail::dispatch(X, Y, Idx);
1039
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 if (!v.max_pool2d_forward)
1040 ✗ detail::throw_not_implemented("max_pool2d_forward", X.device);
1041 19 detail::adopt_output(Y, X.device);
1042 19 detail::adopt_output(Idx, X.device);
1043 38 v.max_pool2d_forward(X, N, C, H, W, kH, kW, stride_h, stride_w,
1044 19 pad_h, pad_w, Y, Idx);
1045 19 }
1046 8 void max_pool2d_backward(const Tensor& dY, const Tensor& Idx,
1047 int N, int C, int H, int W, int H_out, int W_out,
1048 Tensor& dX) {
1049 8 const auto& v = detail::dispatch(dY, Idx, dX);
1050
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.max_pool2d_backward)
1051 ✗ detail::throw_not_implemented("max_pool2d_backward", dY.device);
1052 8 detail::adopt_output(dX, dY.device);
1053 8 v.max_pool2d_backward(dY, Idx, N, C, H, W, H_out, W_out, dX);
1054 8 }
1055
1056 12 void gather_rows(const Tensor& X, const Tensor& Idx, Tensor& Y) {
1057 12 const auto& v = detail::dispatch(X, Idx, Y);
1058
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.gather_rows) detail::throw_not_implemented("gather_rows", X.device);
1059 12 detail::adopt_output(Y, X.device);
1060 12 v.gather_rows(X, Idx, Y);
1061 12 }
1062 9 void scatter_rows_add(const Tensor& dY, const Tensor& Idx, int R, Tensor& dX) {
1063 9 const auto& v = detail::dispatch(dY, Idx, dX);
1064
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.scatter_rows_add)
1065 ✗ detail::throw_not_implemented("scatter_rows_add", dY.device);
1066 9 detail::adopt_output(dX, dY.device);
1067 9 v.scatter_rows_add(dY, Idx, R, dX);
1068 9 }
1069 4 void scatter_rows(const Tensor& Y, const Tensor& Idx, Tensor& X) {
1070 4 const auto& v = detail::dispatch(Y, Idx, X);
1071
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!v.scatter_rows) detail::throw_not_implemented("scatter_rows", Y.device);
1072 // X is an in-place destination (existing rows are preserved), so it is
1073 // never adopted/resized here — the backend validates its shape instead.
1074 4 v.scatter_rows(Y, Idx, X);
1075 4 }
1076
1077 82 void conv_transpose2d_forward(const Tensor& X, const Tensor& Wt,
1078 const Tensor* bias,
1079 int N, int C_in, int H, int W,
1080 int C_out, int kH, int kW,
1081 int stride_h, int stride_w,
1082 int pad_h, int pad_w,
1083 int output_padding_h, int output_padding_w,
1084 int dil_h, int dil_w, int groups,
1085 Tensor& Y) {
1086 82 const auto& v = detail::dispatch(X, Wt, Y);
1087
1/2
✓ Branch 0 taken 82 times.
✗ Branch 1 not taken.
 82 if (!v.conv_transpose2d_forward)
1088 ✗ detail::throw_not_implemented("conv_transpose2d_forward", X.device);
1089 82 detail::adopt_output(Y, X.device);
1090 164 v.conv_transpose2d_forward(X, Wt, bias, N, C_in, H, W, C_out, kH, kW,
1091 82 stride_h, stride_w, pad_h, pad_w,
1092 82 output_padding_h, output_padding_w,
1093 82 dil_h, dil_w, groups, Y);
1094 82 }
1095 8 void conv_transpose2d_backward_input(const Tensor& Wt, const Tensor& dY,
1096 int N, int C_in, int H, int W,
1097 int C_out, int kH, int kW,
1098 int stride_h, int stride_w,
1099 int pad_h, int pad_w,
1100 int output_padding_h, int output_padding_w,
1101 int dil_h, int dil_w, int groups,
1102 Tensor& dX) {
1103 8 const auto& v = detail::dispatch(Wt, dY, dX);
1104
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.conv_transpose2d_backward_input)
1105 ✗ detail::throw_not_implemented("conv_transpose2d_backward_input", Wt.device);
1106 8 detail::adopt_output(dX, Wt.device);
1107 16 v.conv_transpose2d_backward_input(Wt, dY, N, C_in, H, W, C_out, kH, kW,
1108 8 stride_h, stride_w, pad_h, pad_w,
1109 8 output_padding_h, output_padding_w,
1110 8 dil_h, dil_w, groups, dX);
1111 8 }
1112 7 void conv_transpose2d_backward_weight(const Tensor& X, const Tensor& dY,
1113 int N, int C_in, int H, int W,
1114 int C_out, int kH, int kW,
1115 int stride_h, int stride_w,
1116 int pad_h, int pad_w,
1117 int output_padding_h, int output_padding_w,
1118 int dil_h, int dil_w, int groups,
1119 Tensor& dWt) {
1120 7 const auto& v = detail::dispatch(X, dY, dWt);
1121
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (!v.conv_transpose2d_backward_weight)
1122 ✗ detail::throw_not_implemented("conv_transpose2d_backward_weight", X.device);
1123 7 detail::adopt_output(dWt, X.device);
1124 14 v.conv_transpose2d_backward_weight(X, dY, N, C_in, H, W, C_out, kH, kW,
1125 7 stride_h, stride_w, pad_h, pad_w,
1126 7 output_padding_h, output_padding_w,
1127 7 dil_h, dil_w, groups, dWt);
1128 7 }
1129 3 void conv_transpose2d_backward_bias(const Tensor& dY, int N, int C_out,
1130 int H_out, int W_out, Tensor& dB) {
1131 3 const auto& v = detail::dispatch(dY, dB);
1132
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.conv_transpose2d_backward_bias)
1133 ✗ detail::throw_not_implemented("conv_transpose2d_backward_bias", dY.device);
1134 3 detail::adopt_output(dB, dY.device);
1135 3 v.conv_transpose2d_backward_bias(dY, N, C_out, H_out, W_out, dB);
1136 3 }
1137
1138 14 void window_partition_forward(const Tensor& X, int N, int C, int H, int W,
1139 int window, Tensor& Y) {
1140 14 const auto& v = detail::dispatch(X, Y);
1141
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.window_partition_forward)
1142 ✗ detail::throw_not_implemented("window_partition_forward", X.device);
1143 14 detail::adopt_output(Y, X.device);
1144 14 v.window_partition_forward(X, N, C, H, W, window, Y);
1145 14 }
1146 8 void window_reverse_forward(const Tensor& X, int N, int C, int H, int W,
1147 int window, Tensor& Y) {
1148 8 const auto& v = detail::dispatch(X, Y);
1149
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.window_reverse_forward)
1150 ✗ detail::throw_not_implemented("window_reverse_forward", X.device);
1151 8 detail::adopt_output(Y, X.device);
1152 8 v.window_reverse_forward(X, N, C, H, W, window, Y);
1153 8 }
1154
1155 // ─── FP16 linear + GEGLU ───────────────────────────────────────────────────
1156
1157 6 void linear_forward_batched_fp16(const Tensor& W, const Tensor* bias,
1158 const Tensor& X_BD, Tensor& Y_BD) {
1159 6 const auto& v = detail::dispatch_with_opts(W, X_BD, {bias, &Y_BD});
1160
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.linear_forward_batched_fp16) detail::throw_not_implemented("linear_forward_batched_fp16", W.device);
1161 6 detail::adopt_output(Y_BD, W.device);
1162 6 v.linear_forward_batched_fp16(W, bias, X_BD, Y_BD);
1163 6 }
1164
1165 8 void linear_forward_batched_fp16_act(const Tensor& W, const Tensor* bias,
1166 const Tensor& X_BD, int act, Tensor& Y_BD) {
1167 8 const auto& v = detail::dispatch_with_opts(W, X_BD, {bias, &Y_BD});
1168
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.linear_forward_batched_fp16_act)
1169 ✗ detail::throw_not_implemented("linear_forward_batched_fp16_act", W.device);
1170 8 detail::adopt_output(Y_BD, W.device);
1171 8 v.linear_forward_batched_fp16_act(W, bias, X_BD, act, Y_BD);
1172 8 }
1173
1174 10 void geglu_forward(const Tensor& X, Tensor& Y) {
1175 10 const auto& v = detail::dispatch(X, Y);
1176
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.geglu_forward) detail::throw_not_implemented("geglu_forward", X.device);
1177 10 detail::adopt_output(Y, X.device);
1178 10 v.geglu_forward(X, Y);
1179 10 }
1180 6 void geglu_backward(const Tensor& X, const Tensor& dY, Tensor& dX) {
1181 6 const auto& v = detail::dispatch(X, dY, dX);
1182
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.geglu_backward) detail::throw_not_implemented("geglu_backward", X.device);
1183 6 detail::adopt_output(dX, X.device);
1184 6 v.geglu_backward(X, dY, dX);
1185 6 }
1186 8 void geglu_exact_forward(const Tensor& X, Tensor& Y) {
1187 8 const auto& v = detail::dispatch(X, Y);
1188
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.geglu_exact_forward) detail::throw_not_implemented("geglu_exact_forward", X.device);
1189 8 detail::adopt_output(Y, X.device);
1190 8 v.geglu_exact_forward(X, Y);
1191 8 }
1192 6 void geglu_exact_backward(const Tensor& X, const Tensor& dY, Tensor& dX) {
1193 6 const auto& v = detail::dispatch(X, dY, dX);
1194
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.geglu_exact_backward) detail::throw_not_implemented("geglu_exact_backward", X.device);
1195 6 detail::adopt_output(dX, X.device);
1196 6 v.geglu_exact_backward(X, dY, dX);
1197 6 }
1198
1199 // ─── Causal mask helper ────────────────────────────────────────────────────
1200
1201 10 void build_causal_mask_row(int L, int q, Tensor& mask) {
1202 10 const auto& v = detail::dispatch(mask);
1203
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.build_causal_mask_row) detail::throw_not_implemented("build_causal_mask_row", mask.device);
1204 10 detail::adopt_output(mask, mask.device);
1205 10 v.build_causal_mask_row(L, q, mask);
1206 10 }
1207
1208 // ─── Cross-attention family ────────────────────────────────────────────────
1209
1210 16 void cross_attention_forward(const Tensor& X, const Tensor& Ctx,
1211 const Tensor& Wq, const Tensor& Wk,
1212 const Tensor& Wv, const Tensor& Wo,
1213 const float* d_mask, int num_heads, Tensor& O) {
1214 16 const auto& v = detail::dispatch(X, Ctx, Wq, Wk, Wv, Wo, O);
1215
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.cross_attention_forward) detail::throw_not_implemented("cross_attention_forward", X.device);
1216 16 detail::adopt_output(O, X.device);
1217 16 v.cross_attention_forward(X, Ctx, Wq, Wk, Wv, Wo, d_mask, num_heads, O);
1218 16 }
1219
1220 10 void cross_attention_forward_with_attn(const Tensor& X, const Tensor& Ctx,
1221 const Tensor& Wq, const Tensor& Wk,
1222 const Tensor& Wv, const Tensor& Wo,
1223 const float* d_mask,
1224 const Tensor* attn_logit_bias,
1225 int num_heads, Tensor& O, Tensor& AttnAvg) {
1226 10 const auto& v = detail::dispatch_with_opts(X, Ctx, {&Wq, &Wk, &Wv, &Wo, attn_logit_bias, &O, &AttnAvg});
1227
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.cross_attention_forward_with_attn)
1228 ✗ detail::throw_not_implemented("cross_attention_forward_with_attn", X.device);
1229 10 detail::adopt_output(O, X.device);
1230 10 detail::adopt_output(AttnAvg, X.device);
1231 20 v.cross_attention_forward_with_attn(X, Ctx, Wq, Wk, Wv, Wo, d_mask,
1232 10 attn_logit_bias, num_heads, O, AttnAvg);
1233 10 }
1234
1235 16 void self_attention_forward_train(const Tensor& X,
1236 const Tensor& Wq, const Tensor& Wk,
1237 const Tensor& Wv, const Tensor& Wo,
1238 const float* d_mask, int num_heads,
1239 Tensor& Qh, Tensor& Kh, Tensor& Vh,
1240 Tensor& Attnh, Tensor& Yconcat, Tensor& O) {
1241 16 const auto& v = detail::dispatch(X, Wq, Wk, Wv, Wo, Qh, Kh, Vh);
1242
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.self_attention_forward_train)
1243 ✗ detail::throw_not_implemented("self_attention_forward_train", X.device);
1244 16 detail::adopt_output(Qh, X.device);
1245 16 detail::adopt_output(Kh, X.device);
1246 16 detail::adopt_output(Vh, X.device);
1247 16 detail::adopt_output(Attnh, X.device);
1248 16 detail::adopt_output(Yconcat, X.device);
1249 16 detail::adopt_output(O, X.device);
1250 32 v.self_attention_forward_train(X, Wq, Wk, Wv, Wo, d_mask, num_heads,
1251 16 Qh, Kh, Vh, Attnh, Yconcat, O);
1252 16 }
1253
1254 14 void self_attention_backward(const Tensor& dO, const Tensor& X,
1255 const Tensor& Qh, const Tensor& Kh,
1256 const Tensor& Vh, const Tensor& Attnh,
1257 const Tensor& Yconcat,
1258 const Tensor& Wq, const Tensor& Wk,
1259 const Tensor& Wv, const Tensor& Wo,
1260 const float* d_mask, int num_heads,
1261 Tensor& dX,
1262 Tensor& dWq, Tensor& dWk,
1263 Tensor& dWv, Tensor& dWo) {
1264 14 const auto& v = detail::dispatch(dO, X, Qh, Kh, Vh, Attnh, Yconcat);
1265
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.self_attention_backward)
1266 ✗ detail::throw_not_implemented("self_attention_backward", dO.device);
1267 14 detail::adopt_output(dX, dO.device);
1268 14 detail::adopt_output(dWq, dO.device);
1269 14 detail::adopt_output(dWk, dO.device);
1270 14 detail::adopt_output(dWv, dO.device);
1271 14 detail::adopt_output(dWo, dO.device);
1272 28 v.self_attention_backward(dO, X, Qh, Kh, Vh, Attnh, Yconcat,
1273 14 Wq, Wk, Wv, Wo, d_mask, num_heads,
1274 14 dX, dWq, dWk, dWv, dWo);
1275 14 }
1276
1277 10 void attention_token_moments(const Tensor& Attn, int h_lat, int w_lat,
1278 Tensor& mass, Tensor& centroid) {
1279 10 const auto& v = detail::dispatch(Attn, mass, centroid);
1280
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.attention_token_moments) detail::throw_not_implemented("attention_token_moments", Attn.device);
1281 10 detail::adopt_output(mass, Attn.device);
1282 10 detail::adopt_output(centroid, Attn.device);
1283 10 v.attention_token_moments(Attn, h_lat, w_lat, mass, centroid);
1284 10 }
1285
1286 10 void cross_attention_forward_train(const Tensor& X, const Tensor& Ctx,
1287 const Tensor& Wq, const Tensor& Wk,
1288 const Tensor& Wv, const Tensor& Wo,
1289 const float* d_mask, int num_heads,
1290 Tensor& Qh, Tensor& Kh, Tensor& Vh,
1291 Tensor& Attnh, Tensor& Yconcat, Tensor& O) {
1292 10 const auto& v = detail::dispatch(X, Ctx, Wq, Wk, Wv, Wo, Qh, Kh);
1293
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.cross_attention_forward_train)
1294 ✗ detail::throw_not_implemented("cross_attention_forward_train", X.device);
1295 10 detail::adopt_output(Qh, X.device);
1296 10 detail::adopt_output(Kh, X.device);
1297 10 detail::adopt_output(Vh, X.device);
1298 10 detail::adopt_output(Attnh, X.device);
1299 10 detail::adopt_output(Yconcat, X.device);
1300 10 detail::adopt_output(O, X.device);
1301 20 v.cross_attention_forward_train(X, Ctx, Wq, Wk, Wv, Wo, d_mask, num_heads,
1302 10 Qh, Kh, Vh, Attnh, Yconcat, O);
1303 10 }
1304
1305 10 void cross_attention_backward(const Tensor& dO, const Tensor& X, const Tensor& Ctx,
1306 const Tensor& Qh, const Tensor& Kh,
1307 const Tensor& Vh, const Tensor& Attnh,
1308 const Tensor& Yconcat,
1309 const Tensor& Wq, const Tensor& Wk,
1310 const Tensor& Wv, const Tensor& Wo,
1311 const float* d_mask, int num_heads,
1312 Tensor& dX, Tensor& dCtx,
1313 Tensor& dWq, Tensor& dWk,
1314 Tensor& dWv, Tensor& dWo) {
1315 10 const auto& v = detail::dispatch(dO, X, Ctx, Qh, Kh, Vh, Attnh, Yconcat);
1316
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.cross_attention_backward)
1317 ✗ detail::throw_not_implemented("cross_attention_backward", dO.device);
1318 10 detail::adopt_output(dX, dO.device);
1319 10 detail::adopt_output(dCtx, dO.device);
1320 10 detail::adopt_output(dWq, dO.device);
1321 10 detail::adopt_output(dWk, dO.device);
1322 10 detail::adopt_output(dWv, dO.device);
1323 10 detail::adopt_output(dWo, dO.device);
1324 20 v.cross_attention_backward(dO, X, Ctx, Qh, Kh, Vh, Attnh, Yconcat,
1325 10 Wq, Wk, Wv, Wo, d_mask, num_heads,
1326 10 dX, dCtx, dWq, dWk, dWv, dWo);
1327 10 }
1328
1329 // ─── FP16 LN inference + FP16 self-attention ───────────────────────────────
1330
1331 ✗ void layernorm_forward_inference_batched_fp16(const Tensor& X_RD,
1332 const Tensor& gamma,
1333 const Tensor& beta,
1334 Tensor& Y_RD, float eps) {
1335 ✗ const auto& v = detail::dispatch(X_RD, gamma, beta, Y_RD);
1336 ✗ if (!v.layernorm_forward_inference_batched_fp16)
1337 ✗ detail::throw_not_implemented("layernorm_forward_inference_batched_fp16", X_RD.device);
1338 ✗ detail::adopt_output(Y_RD, X_RD.device);
1339 ✗ v.layernorm_forward_inference_batched_fp16(X_RD, gamma, beta, Y_RD, eps);
1340 ✗ }
1341
1342 12 void self_attention_forward(const Tensor& X,
1343 const Tensor& Wq, const Tensor& Wk,
1344 const Tensor& Wv, const Tensor& Wo,
1345 const float* d_mask, int num_heads, Tensor& O) {
1346 12 const auto& v = detail::dispatch(X, Wq, Wk, Wv, Wo, O);
1347
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.self_attention_forward) detail::throw_not_implemented("self_attention_forward", X.device);
1348 12 detail::adopt_output(O, X.device);
1349 12 v.self_attention_forward(X, Wq, Wk, Wv, Wo, d_mask, num_heads, O);
1350 12 }
1351
1352 39 void self_attention_bias_forward(const Tensor& X,
1353 const Tensor& Wq, const Tensor& Wk,
1354 const Tensor& Wv, const Tensor& Wo,
1355 const Tensor* bq, const Tensor* bk,
1356 const Tensor* bv, const Tensor* bo,
1357 const float* d_mask,
1358 const Tensor* attn_bias,
1359 int num_heads, float scale, Tensor& O) {
1360 39 const auto& v = detail::dispatch_with_opts(X, Wq, {&Wk, &Wv, &Wo, bq, bk, bv, bo, attn_bias, &O});
1361
1/2
✓ Branch 0 taken 39 times.
✗ Branch 1 not taken.
 39 if (!v.self_attention_bias_forward)
1362 ✗ detail::throw_not_implemented("self_attention_bias_forward", X.device);
1363 39 detail::adopt_output(O, X.device);
1364 78 v.self_attention_bias_forward(X, Wq, Wk, Wv, Wo, bq, bk, bv, bo, d_mask, attn_bias,
1365 39 num_heads, scale, O);
1366 39 }
1367
1368 6 void rel_pos_bias_xl_forward(const Tensor& Qv, const Tensor& Pk,
1369 int num_heads, int head_dim, Tensor& Bias) {
1370 6 const auto& v = detail::dispatch_with_opts(Qv, Pk, {&Bias});
1371
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.rel_pos_bias_xl_forward)
1372 ✗ detail::throw_not_implemented("rel_pos_bias_xl_forward", Qv.device);
1373 6 detail::adopt_output(Bias, Qv.device);
1374 6 v.rel_pos_bias_xl_forward(Qv, Pk, num_heads, head_dim, Bias);
1375 6 }
1376
1377 23 void self_attention_decomposed_rel_pos_forward(
1378 const Tensor& X,
1379 const Tensor& Wq, const Tensor* bq,
1380 const Tensor& Wk, const Tensor* bk,
1381 const Tensor& Wv, const Tensor* bv,
1382 const Tensor& Wo, const Tensor* bo,
1383 const Tensor& rel_pos_h, const Tensor& rel_pos_w,
1384 int num_heads, int grid_h, int grid_w, float scale, Tensor& O) {
1385 23 const auto& v = detail::dispatch_with_opts(
1386 23 X, Wq, {bq, &Wk, bk, &Wv, bv, &Wo, bo, &rel_pos_h, &rel_pos_w, &O});
1387
1/2
✓ Branch 0 taken 23 times.
✗ Branch 1 not taken.
 23 if (!v.self_attention_decomposed_rel_pos_forward)
1388 ✗ detail::throw_not_implemented("self_attention_decomposed_rel_pos_forward",
1389 ✗ X.device);
1390 23 detail::adopt_output(O, X.device);
1391 46 v.self_attention_decomposed_rel_pos_forward(
1392 23 X, Wq, bq, Wk, bk, Wv, bv, Wo, bo, rel_pos_h, rel_pos_w,
1393 23 num_heads, grid_h, grid_w, scale, O);
1394 23 }
1395
1396 15 void self_attention_decomposed_rel_pos_windowed_forward(
1397 const Tensor& X,
1398 const Tensor& Wq, const Tensor* bq,
1399 const Tensor& Wk, const Tensor* bk,
1400 const Tensor& Wv, const Tensor* bv,
1401 const Tensor& Wo, const Tensor* bo,
1402 const Tensor& rel_pos_h, const Tensor& rel_pos_w,
1403 int num_heads, int grid_h, int grid_w, int window, float scale,
1404 Tensor& O) {
1405 15 const auto& v = detail::dispatch_with_opts(
1406 15 X, Wq, {bq, &Wk, bk, &Wv, bv, &Wo, bo, &rel_pos_h, &rel_pos_w, &O});
1407
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 if (!v.self_attention_decomposed_rel_pos_windowed_forward)
1408 ✗ detail::throw_not_implemented(
1409 ✗ "self_attention_decomposed_rel_pos_windowed_forward", X.device);
1410 15 detail::adopt_output(O, X.device);
1411 30 v.self_attention_decomposed_rel_pos_windowed_forward(
1412 15 X, Wq, bq, Wk, bk, Wv, bv, Wo, bo, rel_pos_h, rel_pos_w,
1413 15 num_heads, grid_h, grid_w, window, scale, O);
1414 15 }
1415
1416 5 void self_attention_bias_int8w_fp16(const Tensor& X,
1417 const Tensor& Wq_int8, const Tensor& sq,
1418 const Tensor& Wk_int8, const Tensor& sk,
1419 const Tensor& Wv_int8, const Tensor& sv,
1420 const Tensor& Wo_int8, const Tensor& so,
1421 const float* d_mask,
1422 const Tensor* attn_bias,
1423 int num_heads, float scale, Tensor& O) {
1424 5 const auto& v = detail::dispatch_with_opts(
1425 15 X, Wq_int8, {&sq, &Wk_int8, &sk, &Wv_int8, &sv, &Wo_int8, &so,
1426 10 attn_bias, &O});
1427
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 if (!v.self_attention_bias_int8w_fp16)
1428 ✗ detail::throw_not_implemented("self_attention_bias_int8w_fp16", X.device);
1429 5 detail::adopt_output(O, X.device);
1430 10 v.self_attention_bias_int8w_fp16(X, Wq_int8, sq, Wk_int8, sk,
1431 5 Wv_int8, sv, Wo_int8, so,
1432 5 d_mask, attn_bias, num_heads, scale, O);
1433 5 }
1434
1435 // ─── Flash attention family ────────────────────────────────────────────────
1436
1437 65 void flash_attention_forward(const Tensor& Q, const Tensor& K, const Tensor& V,
1438 const float* d_mask, int num_heads, bool causal,
1439 Tensor& O) {
1440 65 const auto& v = detail::dispatch(Q, K, V, O);
1441
1/2
✓ Branch 0 taken 65 times.
✗ Branch 1 not taken.
 65 if (!v.flash_attention_forward) detail::throw_not_implemented("flash_attention_forward", Q.device);
1442 65 detail::adopt_output(O, Q.device);
1443 65 v.flash_attention_forward(Q, K, V, d_mask, num_heads, causal, O);
1444 65 }
1445
1446 15 void flash_attention_gqa_forward(const Tensor& Q, const Tensor& K, const Tensor& V,
1447 const float* d_mask, int num_q_heads, int num_kv_heads,
1448 bool causal, Tensor& O) {
1449 15 const auto& v = detail::dispatch(Q, K, V, O);
1450
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 if (!v.flash_attention_gqa_forward)
1451 ✗ detail::throw_not_implemented("flash_attention_gqa_forward", Q.device);
1452 15 detail::adopt_output(O, Q.device);
1453 15 v.flash_attention_gqa_forward(Q, K, V, d_mask, num_q_heads, num_kv_heads, causal, O);
1454 15 }
1455
1456 ✗ void flash_attention_windowed_forward(const Tensor& Q, const Tensor& K, const Tensor& V,
1457 const float* d_mask, int num_heads, int window,
1458 Tensor& O) {
1459 ✗ const auto& v = detail::dispatch(Q, K, V, O);
1460 ✗ if (!v.flash_attention_windowed_forward)
1461 ✗ detail::throw_not_implemented("flash_attention_windowed_forward", Q.device);
1462 ✗ detail::adopt_output(O, Q.device);
1463 ✗ v.flash_attention_windowed_forward(Q, K, V, d_mask, num_heads, window, O);
1464 ✗ }
1465
1466 27 void flash_attention_varlen_forward(const Tensor& Q, const Tensor& K, const Tensor& V,
1467 const int32_t* cu_seqlens_q,
1468 const int32_t* cu_seqlens_k,
1469 int batch_size, int max_seqlen_q, int max_seqlen_k,
1470 int num_heads, int head_dim, bool causal,
1471 Tensor& O) {
1472 27 const auto& v = detail::dispatch(Q, K, V, O);
1473
1/2
✓ Branch 0 taken 27 times.
✗ Branch 1 not taken.
 27 if (!v.flash_attention_varlen_forward)
1474 ✗ detail::throw_not_implemented("flash_attention_varlen_forward", Q.device);
1475 27 detail::adopt_output(O, Q.device);
1476 54 v.flash_attention_varlen_forward(Q, K, V, cu_seqlens_q, cu_seqlens_k,
1477 27 batch_size, max_seqlen_q, max_seqlen_k,
1478 27 num_heads, head_dim, causal, O);
1479 27 }
1480
1481 27 void flash_attention_varlen_backward(const Tensor& Q, const Tensor& K, const Tensor& V,
1482 const Tensor& O, const Tensor& dO,
1483 const int32_t* cu_seqlens_q,
1484 const int32_t* cu_seqlens_k,
1485 int batch_size, int max_seqlen_q, int max_seqlen_k,
1486 int num_heads, int head_dim, bool causal,
1487 Tensor& dQ, Tensor& dK, Tensor& dV) {
1488 27 const auto& v = detail::dispatch(Q, K, V, O, dO, dQ, dK, dV);
1489
1/2
✓ Branch 0 taken 27 times.
✗ Branch 1 not taken.
 27 if (!v.flash_attention_varlen_backward)
1490 ✗ detail::throw_not_implemented("flash_attention_varlen_backward", Q.device);
1491 27 detail::adopt_output(dQ, Q.device);
1492 27 detail::adopt_output(dK, Q.device);
1493 27 detail::adopt_output(dV, Q.device);
1494 54 v.flash_attention_varlen_backward(Q, K, V, O, dO,
1495 27 cu_seqlens_q, cu_seqlens_k,
1496 27 batch_size, max_seqlen_q, max_seqlen_k,
1497 27 num_heads, head_dim, causal,
1498 27 dQ, dK, dV);
1499 27 }
1500
1501 24 void flash_attention_qkvo_forward(const Tensor& X, const Tensor* Ctx,
1502 const Tensor& Wq, const Tensor* bq,
1503 const Tensor& Wk, const Tensor* bk,
1504 const Tensor& Wv, const Tensor* bv,
1505 const Tensor& Wo, const Tensor* bo,
1506 const float* d_mask, int num_heads,
1507 bool causal, Tensor& O) {
1508 24 const auto& v = detail::dispatch_with_opts(
1509 24 X, Wq, {Ctx, bq, &Wk, bk, &Wv, bv, &Wo, bo, &O});
1510
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 if (!v.flash_attention_qkvo_forward)
1511 ✗ detail::throw_not_implemented("flash_attention_qkvo_forward", X.device);
1512 24 detail::adopt_output(O, X.device);
1513 48 v.flash_attention_qkvo_forward(X, Ctx, Wq, bq, Wk, bk, Wv, bv, Wo, bo,
1514 24 d_mask, num_heads, causal, O);
1515 24 }
1516
1517 18 void flash_attention_qkvo_backward(
1518 const Tensor& X, const Tensor* Ctx,
1519 const Tensor& Wq, const Tensor* bq,
1520 const Tensor& Wk, const Tensor* bk,
1521 const Tensor& Wv, const Tensor* bv,
1522 const Tensor& Wo, const Tensor* bo,
1523 const float* d_mask, int num_heads, bool causal,
1524 const Tensor& dO,
1525 Tensor& dX, Tensor* dCtx,
1526 Tensor& dWq, Tensor* dbq,
1527 Tensor& dWk, Tensor* dbk,
1528 Tensor& dWv, Tensor* dbv,
1529 Tensor& dWo, Tensor* dbo) {
1530 18 const auto& v = detail::dispatch_with_opts(
1531 198 X, Wq, {Ctx, bq, &Wk, bk, &Wv, bv, &Wo, bo, &dO,
1532 180 &dX, dCtx, &dWq, dbq, &dWk, dbk, &dWv, dbv, &dWo, dbo});
1533
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 if (!v.flash_attention_qkvo_backward)
1534 ✗ detail::throw_not_implemented("flash_attention_qkvo_backward", X.device);
1535 18 detail::adopt_output(dX, X.device);
1536
2/2
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 6 times.
 18 if (dCtx) detail::adopt_output(*dCtx, X.device);
1537 18 detail::adopt_output(dWq, X.device);
1538
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 12 times.
 18 if (dbq) detail::adopt_output(*dbq, X.device);
1539 18 detail::adopt_output(dWk, X.device);
1540
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 12 times.
 18 if (dbk) detail::adopt_output(*dbk, X.device);
1541 18 detail::adopt_output(dWv, X.device);
1542
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 12 times.
 18 if (dbv) detail::adopt_output(*dbv, X.device);
1543 18 detail::adopt_output(dWo, X.device);
1544
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 12 times.
 18 if (dbo) detail::adopt_output(*dbo, X.device);
1545 36 v.flash_attention_qkvo_backward(X, Ctx, Wq, bq, Wk, bk, Wv, bv, Wo, bo,
1546 18 d_mask, num_heads, causal, dO,
1547 18 dX, dCtx, dWq, dbq, dWk, dbk,
1548 18 dWv, dbv, dWo, dbo);
1549 18 }
1550
1551 16 void flash_attention_backward(const Tensor& Q, const Tensor& K, const Tensor& V,
1552 const Tensor& O, const Tensor& dO,
1553 const float* d_mask, int num_heads, bool causal,
1554 Tensor& dQ, Tensor& dK, Tensor& dV) {
1555 16 const auto& v = detail::dispatch(Q, K, V, O, dO, dQ, dK, dV);
1556
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.flash_attention_backward) detail::throw_not_implemented("flash_attention_backward", Q.device);
1557 16 detail::adopt_output(dQ, Q.device);
1558 16 detail::adopt_output(dK, Q.device);
1559 16 detail::adopt_output(dV, Q.device);
1560 32 v.flash_attention_backward(Q, K, V, O, dO, d_mask, num_heads, causal,
1561 16 dQ, dK, dV);
1562 16 }
1563
1564 8 void flash_attention_project_kv(const Tensor& ctx,
1565 const Tensor& Wk, const Tensor* bk,
1566 const Tensor& Wv, const Tensor* bv,
1567 Tensor& K_out, Tensor& V_out) {
1568 8 const auto& v = detail::dispatch_with_opts(ctx, Wk, {bk, &Wv, bv, &K_out, &V_out});
1569
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.flash_attention_project_kv)
1570 ✗ detail::throw_not_implemented("flash_attention_project_kv", ctx.device);
1571 8 detail::adopt_output(K_out, ctx.device);
1572 8 detail::adopt_output(V_out, ctx.device);
1573 8 v.flash_attention_project_kv(ctx, Wk, bk, Wv, bv, K_out, V_out);
1574 8 }
1575
1576 12 void flash_attention_q_with_kv_cached_forward(const Tensor& X,
1577 const Tensor& K, const Tensor& V,
1578 const Tensor& Wq, const Tensor* bq,
1579 const Tensor& Wo, const Tensor* bo,
1580 const float* d_mask, int num_heads,
1581 bool causal, Tensor& O) {
1582 12 const auto& v = detail::dispatch_with_opts(X, K, {&V, &Wq, bq, &Wo, bo, &O});
1583
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.flash_attention_q_with_kv_cached_forward)
1584 ✗ detail::throw_not_implemented("flash_attention_q_with_kv_cached_forward", X.device);
1585 12 detail::adopt_output(O, X.device);
1586 24 v.flash_attention_q_with_kv_cached_forward(X, K, V, Wq, bq, Wo, bo,
1587 12 d_mask, num_heads, causal, O);
1588 12 }
1589
1590 // ─── NCHW <-> sequence ─────────────────────────────────────────────────────
1591
1592 19 void nchw_to_sequence(const Tensor& X, int N, int C, int H, int W, Tensor& Y) {
1593 19 const auto& v = detail::dispatch(X, Y);
1594
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 if (!v.nchw_to_sequence) detail::throw_not_implemented("nchw_to_sequence", X.device);
1595 19 detail::adopt_output(Y, X.device);
1596 19 v.nchw_to_sequence(X, N, C, H, W, Y);
1597 19 }
1598
1599 19 void sequence_to_nchw(const Tensor& X, int N, int C, int H, int W, Tensor& Y) {
1600 19 const auto& v = detail::dispatch(X, Y);
1601
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 if (!v.sequence_to_nchw) detail::throw_not_implemented("sequence_to_nchw", X.device);
1602 19 detail::adopt_output(Y, X.device);
1603 19 v.sequence_to_nchw(X, N, C, H, W, Y);
1604 19 }
1605
1606 34 void spatial_merge_2x2_forward(const Tensor& X, int N, int C, int H, int W,
1607 bool channel_major, Tensor& Y) {
1608 34 const auto& v = detail::dispatch(X, Y);
1609
1/2
✓ Branch 0 taken 34 times.
✗ Branch 1 not taken.
 34 if (!v.spatial_merge_2x2_forward)
1610 ✗ detail::throw_not_implemented("spatial_merge_2x2_forward", X.device);
1611 34 detail::adopt_output(Y, X.device);
1612 34 v.spatial_merge_2x2_forward(X, N, C, H, W, channel_major, Y);
1613 34 }
1614
1615 47 void pixel_shuffle_upsample_2x_forward(const Tensor& X, int N, int C_in,
1616 int H, int W, int C_out, Tensor& Y) {
1617 47 const auto& v = detail::dispatch(X, Y);
1618
1/2
✓ Branch 0 taken 47 times.
✗ Branch 1 not taken.
 47 if (!v.pixel_shuffle_upsample_2x_forward)
1619 ✗ detail::throw_not_implemented("pixel_shuffle_upsample_2x_forward", X.device);
1620 47 detail::adopt_output(Y, X.device);
1621 47 v.pixel_shuffle_upsample_2x_forward(X, N, C_in, H, W, C_out, Y);
1622 47 }
1623
1624 12 void patch_unpack_forward(const Tensor& tokens, int hp, int wp, int P,
1625 int C_total, int C_keep, bool channel_major,
1626 Tensor& Y) {
1627 12 const auto& v = detail::dispatch(tokens, Y);
1628
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.patch_unpack_forward)
1629 ✗ detail::throw_not_implemented("patch_unpack_forward", tokens.device);
1630 12 detail::adopt_output(Y, tokens.device);
1631 12 v.patch_unpack_forward(tokens, hp, wp, P, C_total, C_keep, channel_major, Y);
1632 12 }
1633
1634 // ─── ResBlock ──────────────────────────────────────────────────────────────
1635
1636 16 void resblock_forward(const Tensor& X,
1637 const Tensor& gamma1, const Tensor& beta1,
1638 const Tensor& W1, const Tensor* b1,
1639 const Tensor* t_emb_shift,
1640 const Tensor& gamma2, const Tensor& beta2,
1641 const Tensor& W2, const Tensor* b2,
1642 const Tensor* Wskip, const Tensor* bskip,
1643 int N, int C_in, int C_out, int H, int W,
1644 int num_groups, float eps, Tensor& Y) {
1645 16 const auto& v = detail::dispatch_with_opts(
1646 64 X, gamma1, {&beta1, &W1, b1, t_emb_shift, &gamma2, &beta2, &W2, b2,
1647 48 Wskip, bskip, &Y});
1648
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.resblock_forward) detail::throw_not_implemented("resblock_forward", X.device);
1649 16 detail::adopt_output(Y, X.device);
1650 32 v.resblock_forward(X, gamma1, beta1, W1, b1, t_emb_shift,
1651 16 gamma2, beta2, W2, b2, Wskip, bskip,
1652 16 N, C_in, C_out, H, W, num_groups, eps, Y);
1653 16 }
1654
1655 ✗ void resblock_forward_int8w_fp16(const Tensor& X,
1656 const Tensor& gamma1, const Tensor& beta1,
1657 const Tensor& W1_int8, const Tensor& s1,
1658 const Tensor* b1,
1659 const Tensor* t_emb_shift,
1660 const Tensor& gamma2, const Tensor& beta2,
1661 const Tensor& W2_int8, const Tensor& s2,
1662 const Tensor* b2,
1663 const Tensor* Wskip_int8, const Tensor* sskip,
1664 const Tensor* bskip,
1665 int N, int C_in, int C_out, int H, int W,
1666 int num_groups, float eps, Tensor& Y) {
1667 ✗ const auto& v = detail::dispatch_with_opts(
1668 ✗ X, gamma1, {&beta1, &W1_int8, &s1, b1, t_emb_shift,
1669 ✗ &gamma2, &beta2, &W2_int8, &s2, b2,
1670 ✗ Wskip_int8, sskip, bskip, &Y});
1671 ✗ if (!v.resblock_forward_int8w_fp16)
1672 ✗ detail::throw_not_implemented("resblock_forward_int8w_fp16", X.device);
1673 ✗ detail::adopt_output(Y, X.device);
1674 ✗ v.resblock_forward_int8w_fp16(X, gamma1, beta1, W1_int8, s1, b1, t_emb_shift,
1675 ✗ gamma2, beta2, W2_int8, s2, b2,
1676 ✗ Wskip_int8, sskip, bskip,
1677 ✗ N, C_in, C_out, H, W, num_groups, eps, Y);
1678 ✗ }
1679
1680 16 void resblock_backward(const Tensor& X,
1681 const Tensor& gamma1, const Tensor& beta1,
1682 const Tensor& W1, const Tensor* b1,
1683 const Tensor* t_emb_shift,
1684 const Tensor& gamma2, const Tensor& beta2,
1685 const Tensor& W2, const Tensor* b2,
1686 const Tensor* Wskip, const Tensor* bskip,
1687 int N, int C_in, int C_out, int H, int W,
1688 int num_groups, float eps,
1689 const Tensor& dY,
1690 Tensor& dX,
1691 Tensor& dGamma1, Tensor& dBeta1,
1692 Tensor& dW1, Tensor* db1,
1693 Tensor* dt_emb_shift,
1694 Tensor& dGamma2, Tensor& dBeta2,
1695 Tensor& dW2, Tensor* db2,
1696 Tensor* dWskip, Tensor* dbskip) {
1697 16 const auto& v = detail::dispatch_with_opts(
1698 256 X, gamma1, {&beta1, &W1, b1, t_emb_shift, &gamma2, &beta2, &W2, b2,
1699 128 Wskip, bskip, &dY, &dX, &dGamma1, &dBeta1, &dW1, db1,
1700 112 dt_emb_shift, &dGamma2, &dBeta2, &dW2, db2, dWskip, dbskip});
1701
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.resblock_backward) detail::throw_not_implemented("resblock_backward", X.device);
1702 16 detail::adopt_output(dX, X.device);
1703 16 detail::adopt_output(dGamma1, X.device);
1704 16 detail::adopt_output(dBeta1, X.device);
1705 16 detail::adopt_output(dW1, X.device);
1706
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 16 times.
 16 if (db1) detail::adopt_output(*db1, X.device);
1707
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 8 times.
 16 if (dt_emb_shift) detail::adopt_output(*dt_emb_shift, X.device);
1708 16 detail::adopt_output(dGamma2, X.device);
1709 16 detail::adopt_output(dBeta2, X.device);
1710 16 detail::adopt_output(dW2, X.device);
1711
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 16 times.
 16 if (db2) detail::adopt_output(*db2, X.device);
1712
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 8 times.
 16 if (dWskip) detail::adopt_output(*dWskip, X.device);
1713
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 8 times.
 16 if (dbskip) detail::adopt_output(*dbskip, X.device);
1714 32 v.resblock_backward(X, gamma1, beta1, W1, b1, t_emb_shift,
1715 16 gamma2, beta2, W2, b2, Wskip, bskip,
1716 16 N, C_in, C_out, H, W, num_groups, eps,
1717 16 dY, dX, dGamma1, dBeta1, dW1, db1, dt_emb_shift,
1718 16 dGamma2, dBeta2, dW2, db2, dWskip, dbskip);
1719 16 }
1720
1721 // ─── Matmul + RoPE + RMSNorm + SwiGLU + KV-cache + Llama ───────────────────
1722
1723 60 void matmul(const Tensor& A, const Tensor& B, Tensor& C) {
1724 60 const auto& v = detail::dispatch(A, B, C);
1725
1/2
✓ Branch 0 taken 60 times.
✗ Branch 1 not taken.
 60 if (!v.matmul) detail::throw_not_implemented("matmul", A.device);
1726 60 detail::adopt_output(C, A.device);
1727 60 v.matmul(A, B, C);
1728 60 }
1729
1730 62 void matmul_abt(const Tensor& A, const Tensor& B, Tensor& C,
1731 int batch, int M, int N, int K,
1732 long long strideA, long long strideB, long long strideC,
1733 const Tensor* bias, int act) {
1734 62 const auto& v = detail::dispatch(A, B, C);
1735
1/2
✓ Branch 0 taken 62 times.
✗ Branch 1 not taken.
 62 if (!v.matmul_abt) detail::throw_not_implemented("matmul_abt", A.device);
1736 62 detail::adopt_output(C, A.device);
1737 62 v.matmul_abt(A, B, C, batch, M, N, K, strideA, strideB, strideC, bias, act);
1738 62 }
1739
1740 16 void matmul_backward(const Tensor& A, const Tensor& B, const Tensor& dC,
1741 Tensor& dA, Tensor& dB) {
1742 16 const auto& v = detail::dispatch(A, B, dC, dA, dB);
1743
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.matmul_backward) detail::throw_not_implemented("matmul_backward", A.device);
1744 16 detail::adopt_output(dA, A.device);
1745 16 detail::adopt_output(dB, A.device);
1746 16 v.matmul_backward(A, B, dC, dA, dB);
1747 16 }
1748
1749 907 void lstm_forward_train(const Tensor& X, const Tensor& W_ih, const Tensor& W_hh,
1750 const Tensor* b_ih, const Tensor* b_hh,
1751 const Tensor* h0, const Tensor* c0, int T, int B,
1752 Tensor& Y, Tensor& gates, Tensor& C,
1753 Tensor* hT, Tensor* cT) {
1754 907 const auto& v = detail::dispatch(X);
1755
1/2
✓ Branch 0 taken 907 times.
✗ Branch 1 not taken.
 907 if (!v.lstm_forward_train) detail::throw_not_implemented("lstm_forward_train", X.device);
1756 907 detail::adopt_output(Y, X.device);
1757 907 detail::adopt_output(gates, X.device);
1758 907 detail::adopt_output(C, X.device);
1759
2/2
✓ Branch 0 taken 895 times.
✓ Branch 1 taken 12 times.
 907 if (hT) detail::adopt_output(*hT, X.device);
1760
2/2
✓ Branch 0 taken 895 times.
✓ Branch 1 taken 12 times.
 907 if (cT) detail::adopt_output(*cT, X.device);
1761 907 v.lstm_forward_train(X, W_ih, W_hh, b_ih, b_hh, h0, c0, T, B, Y, gates, C, hT, cT);
1762 907 }
1763
1764 413 void lstm_backward(const Tensor& X, const Tensor& W_ih, const Tensor& W_hh,
1765 const Tensor* h0, const Tensor* c0,
1766 const Tensor& Y, const Tensor& gates, const Tensor& C,
1767 const Tensor& dY, int T, int B,
1768 Tensor& dX, Tensor& dW_ih, Tensor& dW_hh,
1769 Tensor* db_ih, Tensor* db_hh,
1770 Tensor* dh0, Tensor* dc0) {
1771 413 const auto& v = detail::dispatch(X);
1772
1/2
✓ Branch 0 taken 413 times.
✗ Branch 1 not taken.
 413 if (!v.lstm_backward) detail::throw_not_implemented("lstm_backward", X.device);
1773 413 detail::adopt_output(dX, X.device);
1774 413 detail::adopt_output(dW_ih, X.device);
1775 413 detail::adopt_output(dW_hh, X.device);
1776
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 413 times.
 413 if (db_ih) detail::adopt_output(*db_ih, X.device);
1777
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 413 times.
 413 if (db_hh) detail::adopt_output(*db_hh, X.device);
1778
2/2
✓ Branch 0 taken 400 times.
✓ Branch 1 taken 13 times.
 413 if (dh0) detail::adopt_output(*dh0, X.device);
1779
2/2
✓ Branch 0 taken 400 times.
✓ Branch 1 taken 13 times.
 413 if (dc0) detail::adopt_output(*dc0, X.device);
1780 826 v.lstm_backward(X, W_ih, W_hh, h0, c0, Y, gates, C, dY, T, B,
1781 413 dX, dW_ih, dW_hh, db_ih, db_hh, dh0, dc0);
1782 413 }
1783
1784 14 void rope_forward(const Tensor& X, int head_dim, int num_heads,
1785 int seq_offset, float theta_base, Tensor& Y) {
1786 14 const auto& v = detail::dispatch(X, Y);
1787
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.rope_forward) detail::throw_not_implemented("rope_forward", X.device);
1788 14 detail::adopt_output(Y, X.device);
1789 14 v.rope_forward(X, head_dim, num_heads, seq_offset, theta_base, Y);
1790 14 }
1791
1792 8 void rope_backward(const Tensor& dY, int head_dim, int num_heads,
1793 int seq_offset, float theta_base, Tensor& dX) {
1794 8 const auto& v = detail::dispatch(dY, dX);
1795
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.rope_backward) detail::throw_not_implemented("rope_backward", dY.device);
1796 8 detail::adopt_output(dX, dY.device);
1797 8 v.rope_backward(dY, head_dim, num_heads, seq_offset, theta_base, dX);
1798 8 }
1799
1800 26 void rope_apply(const Tensor& X, const Tensor& cos_tbl, const Tensor& sin_tbl,
1801 int head_dim, int num_heads, Tensor& Y) {
1802 26 const auto& v = detail::dispatch(X, cos_tbl, sin_tbl, Y);
1803
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 26 if (!v.rope_apply) detail::throw_not_implemented("rope_apply", X.device);
1804 26 detail::adopt_output(Y, X.device);
1805 26 v.rope_apply(X, cos_tbl, sin_tbl, head_dim, num_heads, Y);
1806 26 }
1807
1808 2 void rope_apply_perhead(const Tensor& X, const Tensor& cos_tbl,
1809 const Tensor& sin_tbl, int head_dim, int num_heads,
1810 Tensor& Y) {
1811 2 const auto& v = detail::dispatch(X, cos_tbl, sin_tbl, Y);
1812
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!v.rope_apply_perhead)
1813 ✗ detail::throw_not_implemented("rope_apply_perhead", X.device);
1814 2 detail::adopt_output(Y, X.device);
1815 2 v.rope_apply_perhead(X, cos_tbl, sin_tbl, head_dim, num_heads, Y);
1816 2 }
1817
1818 4 void rope_apply_backward(const Tensor& dY, const Tensor& cos_tbl,
1819 const Tensor& sin_tbl, int head_dim, int num_heads,
1820 Tensor& dX) {
1821 4 const auto& v = detail::dispatch(dY, cos_tbl, sin_tbl, dX);
1822
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!v.rope_apply_backward)
1823 ✗ detail::throw_not_implemented("rope_apply_backward", dY.device);
1824 4 detail::adopt_output(dX, dY.device);
1825 4 v.rope_apply_backward(dY, cos_tbl, sin_tbl, head_dim, num_heads, dX);
1826 4 }
1827
1828 11 void rope_apply_mrope(const Tensor& X,
1829 const Tensor& cos_t, const Tensor& sin_t,
1830 const Tensor& cos_h, const Tensor& sin_h,
1831 const Tensor& cos_w, const Tensor& sin_w,
1832 const int32_t* pos_t, const int32_t* pos_h,
1833 const int32_t* pos_w,
1834 int head_dim, int num_heads,
1835 int d_t, int d_h, int d_w,
1836 Tensor& Y) {
1837 22 const auto& v = detail::dispatch(X, cos_t, sin_t, cos_h, sin_h, cos_w,
1838 11 sin_w, Y);
1839
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.rope_apply_mrope)
1840 ✗ detail::throw_not_implemented("rope_apply_mrope", X.device);
1841 11 detail::adopt_output(Y, X.device);
1842 22 v.rope_apply_mrope(X, cos_t, sin_t, cos_h, sin_h, cos_w, sin_w,
1843 11 pos_t, pos_h, pos_w, head_dim, num_heads,
1844 11 d_t, d_h, d_w, Y);
1845 11 }
1846
1847 18 void rms_norm_forward(const Tensor& X, const Tensor& gamma, float eps, Tensor& Y) {
1848 18 const auto& v = detail::dispatch(X, gamma, Y);
1849
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 if (!v.rms_norm_forward) detail::throw_not_implemented("rms_norm_forward", X.device);
1850 18 detail::adopt_output(Y, X.device);
1851 18 v.rms_norm_forward(X, gamma, eps, Y);
1852 18 }
1853
1854 14 void rms_norm_backward(const Tensor& X, const Tensor& gamma, const Tensor& dY,
1855 float eps, Tensor& dX, Tensor& dGamma) {
1856 14 const auto& v = detail::dispatch(X, gamma, dY, dX, dGamma);
1857
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.rms_norm_backward) detail::throw_not_implemented("rms_norm_backward", X.device);
1858 14 detail::adopt_output(dX, X.device);
1859 14 detail::adopt_output(dGamma, X.device);
1860 14 v.rms_norm_backward(X, gamma, dY, eps, dX, dGamma);
1861 14 }
1862
1863 10 void swiglu_forward(const Tensor& X, Tensor& Y) {
1864 10 const auto& v = detail::dispatch(X, Y);
1865
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.swiglu_forward) detail::throw_not_implemented("swiglu_forward", X.device);
1866 10 detail::adopt_output(Y, X.device);
1867 10 v.swiglu_forward(X, Y);
1868 10 }
1869 10 void swiglu_backward(const Tensor& X, const Tensor& dY, Tensor& dX) {
1870 10 const auto& v = detail::dispatch(X, dY, dX);
1871
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.swiglu_backward) detail::throw_not_implemented("swiglu_backward", X.device);
1872 10 detail::adopt_output(dX, X.device);
1873 10 v.swiglu_backward(X, dY, dX);
1874 10 }
1875
1876 28 void kv_cache_append(const Tensor& K_new, const Tensor& V_new, int cur_len,
1877 Tensor& K_cache, Tensor& V_cache) {
1878 28 const auto& v = detail::dispatch(K_new, V_new, K_cache, V_cache);
1879
1/2
✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
 28 if (!v.kv_cache_append) detail::throw_not_implemented("kv_cache_append", K_new.device);
1880 28 detail::adopt_output(K_cache, K_new.device);
1881 28 detail::adopt_output(V_cache, K_new.device);
1882 28 v.kv_cache_append(K_new, V_new, cur_len, K_cache, V_cache);
1883 28 }
1884
1885 51 void flash_attention_decode(const Tensor& Q,
1886 const Tensor& K_cache, const Tensor& V_cache,
1887 int valid_len, int num_q_heads, int num_kv_heads,
1888 Tensor& O, float attn_softcap, int window) {
1889 51 const auto& v = detail::dispatch(Q, K_cache, V_cache, O);
1890
1/2
✓ Branch 0 taken 51 times.
✗ Branch 1 not taken.
 51 if (!v.flash_attention_decode) detail::throw_not_implemented("flash_attention_decode", Q.device);
1891 51 detail::adopt_output(O, Q.device);
1892 102 v.flash_attention_decode(Q, K_cache, V_cache, valid_len,
1893 51 num_q_heads, num_kv_heads, O, attn_softcap, window);
1894 51 }
1895
1896 12 void flash_attention_decode_masked(const Tensor& Q,
1897 const Tensor& K_cache,
1898 const Tensor& V_cache,
1899 const float* d_mask,
1900 int num_q_heads, int num_kv_heads,
1901 Tensor& O, float attn_softcap, int window) {
1902 12 const auto& v = detail::dispatch(Q, K_cache, V_cache, O);
1903
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.flash_attention_decode_masked) {
1904 ✗ detail::throw_not_implemented("flash_attention_decode_masked", Q.device);
1905 }
1906 12 detail::adopt_output(O, Q.device);
1907 24 v.flash_attention_decode_masked(Q, K_cache, V_cache, d_mask,
1908 12 num_q_heads, num_kv_heads, O,
1909 12 attn_softcap, window);
1910 12 }
1911
1912 // ─── Public reductions ─────────────────────────────────────────────────────
1913
1914 15 void sum_rows(const Tensor& X, Tensor& Y) {
1915 15 const auto& v = detail::dispatch(X, Y);
1916
1/2
✓ Branch 0 taken 15 times.
✗ Branch 1 not taken.
 15 if (!v.sum_rows) detail::throw_not_implemented("sum_rows", X.device);
1917 15 detail::adopt_output(Y, X.device);
1918 15 v.sum_rows(X, Y);
1919 15 }
1920 14 void sum_cols(const Tensor& X, Tensor& Y) {
1921 14 const auto& v = detail::dispatch(X, Y);
1922
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.sum_cols) detail::throw_not_implemented("sum_cols", X.device);
1923 14 detail::adopt_output(Y, X.device);
1924 14 v.sum_cols(X, Y);
1925 14 }
1926 17 void argmax_rows(const Tensor& X, Tensor& Idx) {
1927 17 const auto& v = detail::dispatch(X, Idx);
1928
1/2
✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
 17 if (!v.argmax_rows) detail::throw_not_implemented("argmax_rows", X.device);
1929 17 detail::adopt_output(Idx, X.device);
1930 17 v.argmax_rows(X, Idx);
1931 17 }
1932 11 void rows_count_above(const Tensor& X, float t_lo, float t_hi, Tensor& counts) {
1933 11 const auto& v = detail::dispatch(X, counts);
1934
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.rows_count_above)
1935 ✗ detail::throw_not_implemented("rows_count_above", X.device);
1936 11 detail::adopt_output(counts, X.device);
1937 11 v.rows_count_above(X, t_lo, t_hi, counts);
1938 11 }
1939 9 void threshold_u8(const Tensor& X, float t, Tensor& Y) {
1940 9 const auto& v = detail::dispatch(X, Y);
1941
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.threshold_u8) detail::throw_not_implemented("threshold_u8", X.device);
1942 9 detail::adopt_output(Y, X.device);
1943 9 v.threshold_u8(X, t, Y);
1944 9 }
1945
1946 // ─── Diffusion sampler steps + timestep embedding ──────────────────────────
1947
1948 14 void ddim_step(const Tensor& x_t, const Tensor& eps_pred,
1949 float alpha_t, float alpha_prev, float sigma_t, Tensor& x_prev) {
1950 14 const auto& v = detail::dispatch(x_t, eps_pred, x_prev);
1951
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.ddim_step) detail::throw_not_implemented("ddim_step", x_t.device);
1952 14 detail::adopt_output(x_prev, x_t.device);
1953 14 v.ddim_step(x_t, eps_pred, alpha_t, alpha_prev, sigma_t, x_prev);
1954 14 }
1955
1956 12 void euler_step(const Tensor& x_t, const Tensor& eps_pred,
1957 float sigma_t, float sigma_prev, Tensor& x_prev) {
1958 12 const auto& v = detail::dispatch(x_t, eps_pred, x_prev);
1959
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.euler_step) detail::throw_not_implemented("euler_step", x_t.device);
1960 12 detail::adopt_output(x_prev, x_t.device);
1961 12 v.euler_step(x_t, eps_pred, sigma_t, sigma_prev, x_prev);
1962 12 }
1963
1964 12 void dpmpp_2m_step(const Tensor& x_t, const Tensor& eps_pred,
1965 const Tensor& x0_prev, float sigma_t,
1966 float c_xt, float c_x0t, float c_x0prev,
1967 Tensor& x_prev, Tensor& x0_out) {
1968 12 const auto& v = detail::dispatch(x_t, eps_pred, x0_prev, x_prev, x0_out);
1969
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.dpmpp_2m_step) detail::throw_not_implemented("dpmpp_2m_step", x_t.device);
1970 12 detail::adopt_output(x_prev, x_t.device);
1971 12 detail::adopt_output(x0_out, x_t.device);
1972 24 v.dpmpp_2m_step(x_t, eps_pred, x0_prev, sigma_t, c_xt, c_x0t, c_x0prev,
1973 12 x_prev, x0_out);
1974 12 }
1975
1976 8 void timestep_embedding(const Tensor& timesteps, int dim, float max_period,
1977 Tensor& Y) {
1978 8 const auto& v = detail::dispatch(timesteps, Y);
1979
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.timestep_embedding) detail::throw_not_implemented("timestep_embedding", timesteps.device);
1980 8 detail::adopt_output(Y, timesteps.device);
1981 8 v.timestep_embedding(timesteps, dim, max_period, Y);
1982 8 }
1983
1984 // ─── INT8 weight-only quantisation (W8A16) ─────────────────────────────────
1985
1986 // Host helper — pure host buffers, no device dispatch. Rows are independent,
1987 // so they fan out over the CPU thread pool: the multi-GB checkpoint quantise
1988 // during a quantized model load is this function's dominant caller, and the
1989 // serial version left it minutes-long on 10B+-parameter models.
1990 20 void quantize_int8_per_row_host(const uint16_t* W_fp16,
1991 int out, int in,
1992 int8_t* W_int8_out,
1993 float* scales_out) {
1994
2/4
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 20 times.
 20 if (out <= 0 || in <= 0) {
1995 ✗ for (int r = 0; r < out; ++r) scales_out[r] = 0.0f;
1996 ✗ return;
1997 }
1998
1/2
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
 845 detail::cpu::parallel_for(static_cast<std::size_t>(out), [&](std::size_t r) {
1999 825 const uint16_t* row = W_fp16 + r * static_cast<std::size_t>(in);
2000 825 float amax = 0.0f;
2001
2/2
✓ Branch 0 taken 33215 times.
✓ Branch 1 taken 825 times.
 34040 for (int c = 0; c < in; ++c) {
2002 33215 const float v = fp16_bits_to_fp32(row[c]);
2003 33215 const float a = std::fabs(v);
2004
2/2
✓ Branch 0 taken 29764 times.
✓ Branch 1 taken 3451 times.
 33215 if (a > amax) amax = a;
2005 33215 }
2006
1/2
✓ Branch 0 taken 825 times.
✗ Branch 1 not taken.
 825 const float scale = (amax > 0.0f) ? (amax / 127.0f) : 0.0f;
2007
1/2
✓ Branch 0 taken 825 times.
✗ Branch 1 not taken.
 825 const float inv = (scale > 0.0f) ? (1.0f / scale) : 0.0f;
2008 825 scales_out[r] = scale;
2009 825 int8_t* dst = W_int8_out + r * static_cast<std::size_t>(in);
2010
2/2
✓ Branch 0 taken 33510 times.
✓ Branch 1 taken 825 times.
 34335 for (int c = 0; c < in; ++c) {
2011 33510 const float v = fp16_bits_to_fp32(row[c]);
2012 33510 int q = static_cast<int>(std::lrint(v * inv));
2013
1/2
✓ Branch 0 taken 33510 times.
✗ Branch 1 not taken.
 33510 if (q < -127) q = -127;
2014
1/2
✓ Branch 0 taken 33510 times.
✗ Branch 1 not taken.
 33510 if (q > 127) q = 127;
2015 33510 dst[c] = static_cast<int8_t>(q);
2016 33510 }
2017 825 });
2018 20 }
2019
2020 1 void matmul_int8w_fp16(const Tensor& W_int8, const Tensor& scales,
2021 const Tensor& X, Tensor& Y) {
2022 1 const auto& v = detail::dispatch(W_int8, scales, X, Y);
2023
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
 1 if (!v.matmul_int8w_fp16) detail::throw_not_implemented("matmul_int8w_fp16", W_int8.device);
2024 ✗ detail::adopt_output(Y, W_int8.device);
2025 ✗ v.matmul_int8w_fp16(W_int8, scales, X, Y);
2026 ✗ }
2027
2028 ✗ void conv2d_int8w_fp16_forward(const Tensor& X,
2029 const Tensor& W_int8, const Tensor& scales,
2030 const Tensor* bias,
2031 int N, int C_in, int H, int W,
2032 int C_out, int kH, int kW,
2033 int stride_h, int stride_w,
2034 int pad_h, int pad_w,
2035 int dil_h, int dil_w, int groups,
2036 Tensor& Y) {
2037 ✗ const auto& v = detail::dispatch_with_opts(X, W_int8, {&scales, bias, &Y});
2038 ✗ if (!v.conv2d_int8w_fp16_forward)
2039 ✗ detail::throw_not_implemented("conv2d_int8w_fp16_forward", X.device);
2040 ✗ detail::adopt_output(Y, X.device);
2041 ✗ v.conv2d_int8w_fp16_forward(X, W_int8, scales, bias,
2042 ✗ N, C_in, H, W, C_out, kH, kW,
2043 ✗ stride_h, stride_w, pad_h, pad_w,
2044 ✗ dil_h, dil_w, groups, Y);
2045 ✗ }
2046
2047 ✗ void linear_forward_batched_int8w_fp16(const Tensor& W_int8,
2048 const Tensor& scales,
2049 const Tensor* bias,
2050 const Tensor& X_BD, Tensor& Y_BD) {
2051 ✗ const auto& v = detail::dispatch_with_opts(W_int8, scales, {bias, &X_BD, &Y_BD});
2052 ✗ if (!v.linear_forward_batched_int8w_fp16)
2053 ✗ detail::throw_not_implemented("linear_forward_batched_int8w_fp16", W_int8.device);
2054 ✗ detail::adopt_output(Y_BD, W_int8.device);
2055 ✗ v.linear_forward_batched_int8w_fp16(W_int8, scales, bias, X_BD, Y_BD);
2056 ✗ }
2057
2058 2 void dequant_q4k_to_fp16(const Tensor& W_q4k, Tensor& W_fp16) {
2059 2 const auto& v = detail::dispatch(W_q4k, W_fp16);
2060
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 1 time.
 2 if (!v.dequant_q4k_to_fp16)
2061 1 detail::throw_not_implemented("dequant_q4k_to_fp16", W_q4k.device);
2062 1 detail::adopt_output(W_fp16, W_q4k.device);
2063 1 v.dequant_q4k_to_fp16(W_q4k, W_fp16);
2064 1 }
2065
2066 34 void linear_forward_q4k_fp16(const Tensor& W_q4k, const Tensor* bias,
2067 const Tensor& x, Tensor& y) {
2068 34 const auto& v = detail::dispatch_with_opts(W_q4k, {bias, &x, &y});
2069
1/2
✓ Branch 0 taken 34 times.
✗ Branch 1 not taken.
 34 if (!v.linear_forward_q4k_fp16)
2070 ✗ detail::throw_not_implemented("linear_forward_q4k_fp16", W_q4k.device);
2071 34 detail::adopt_output(y, W_q4k.device);
2072 34 v.linear_forward_q4k_fp16(W_q4k, bias, x, y);
2073 34 }
2074
2075 3 void linear_forward_batched_q4k_fp16(const Tensor& W_q4k, const Tensor* bias,
2076 const Tensor& X_BD, Tensor& Y_BD) {
2077 3 const auto& v = detail::dispatch_with_opts(W_q4k, {bias, &X_BD, &Y_BD});
2078
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.linear_forward_batched_q4k_fp16)
2079 ✗ detail::throw_not_implemented("linear_forward_batched_q4k_fp16", W_q4k.device);
2080 3 detail::adopt_output(Y_BD, W_q4k.device);
2081 3 v.linear_forward_batched_q4k_fp16(W_q4k, bias, X_BD, Y_BD);
2082 3 }
2083
2084 4 void dequant_q8_0_to_fp16(const Tensor& W_q8, Tensor& W_fp16) {
2085 4 const auto& v = detail::dispatch(W_q8, W_fp16);
2086
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!v.dequant_q8_0_to_fp16)
2087 ✗ detail::throw_not_implemented("dequant_q8_0_to_fp16", W_q8.device);
2088 4 detail::adopt_output(W_fp16, W_q8.device);
2089 4 v.dequant_q8_0_to_fp16(W_q8, W_fp16);
2090 4 }
2091
2092 34 void linear_forward_q8_0_fp16(const Tensor& W_q8, const Tensor* bias,
2093 const Tensor& x, Tensor& y) {
2094 34 const auto& v = detail::dispatch_with_opts(W_q8, {bias, &x, &y});
2095
1/2
✓ Branch 0 taken 34 times.
✗ Branch 1 not taken.
 34 if (!v.linear_forward_q8_0_fp16)
2096 ✗ detail::throw_not_implemented("linear_forward_q8_0_fp16", W_q8.device);
2097 34 detail::adopt_output(y, W_q8.device);
2098 34 v.linear_forward_q8_0_fp16(W_q8, bias, x, y);
2099 34 }
2100
2101 3 void linear_forward_batched_q8_0_fp16(const Tensor& W_q8, const Tensor* bias,
2102 const Tensor& X_BD, Tensor& Y_BD) {
2103 3 const auto& v = detail::dispatch_with_opts(W_q8, {bias, &X_BD, &Y_BD});
2104
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.linear_forward_batched_q8_0_fp16)
2105 ✗ detail::throw_not_implemented("linear_forward_batched_q8_0_fp16", W_q8.device);
2106 3 detail::adopt_output(Y_BD, W_q8.device);
2107 3 v.linear_forward_batched_q8_0_fp16(W_q8, bias, X_BD, Y_BD);
2108 3 }
2109
2110 1 void dequant_q6k_to_fp16(const Tensor& W_q6k, Tensor& W_fp16) {
2111 1 const auto& v = detail::dispatch(W_q6k, W_fp16);
2112
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 if (!v.dequant_q6k_to_fp16)
2113 ✗ detail::throw_not_implemented("dequant_q6k_to_fp16", W_q6k.device);
2114 1 detail::adopt_output(W_fp16, W_q6k.device);
2115 1 v.dequant_q6k_to_fp16(W_q6k, W_fp16);
2116 1 }
2117
2118 34 void linear_forward_q6k_fp16(const Tensor& W_q6k, const Tensor* bias,
2119 const Tensor& x, Tensor& y) {
2120 34 const auto& v = detail::dispatch_with_opts(W_q6k, {bias, &x, &y});
2121
1/2
✓ Branch 0 taken 34 times.
✗ Branch 1 not taken.
 34 if (!v.linear_forward_q6k_fp16)
2122 ✗ detail::throw_not_implemented("linear_forward_q6k_fp16", W_q6k.device);
2123 34 detail::adopt_output(y, W_q6k.device);
2124 34 v.linear_forward_q6k_fp16(W_q6k, bias, x, y);
2125 34 }
2126
2127 3 void linear_forward_batched_q6k_fp16(const Tensor& W_q6k, const Tensor* bias,
2128 const Tensor& X_BD, Tensor& Y_BD) {
2129 3 const auto& v = detail::dispatch_with_opts(W_q6k, {bias, &X_BD, &Y_BD});
2130
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.linear_forward_batched_q6k_fp16)
2131 ✗ detail::throw_not_implemented("linear_forward_batched_q6k_fp16", W_q6k.device);
2132 3 detail::adopt_output(Y_BD, W_q6k.device);
2133 3 v.linear_forward_batched_q6k_fp16(W_q6k, bias, X_BD, Y_BD);
2134 3 }
2135
2136 ✗ void flash_attention_project_kv_int8w_fp16(const Tensor& ctx,
2137 const Tensor& Wk_int8,
2138 const Tensor& sk,
2139 const Tensor* bk,
2140 const Tensor& Wv_int8,
2141 const Tensor& sv,
2142 const Tensor* bv,
2143 Tensor& K_out, Tensor& V_out) {
2144 ✗ const auto& v = detail::dispatch_with_opts(
2145 ✗ ctx, Wk_int8, {&sk, bk, &Wv_int8, &sv, bv, &K_out, &V_out});
2146 ✗ if (!v.flash_attention_project_kv_int8w_fp16)
2147 ✗ detail::throw_not_implemented("flash_attention_project_kv_int8w_fp16", ctx.device);
2148 ✗ detail::adopt_output(K_out, ctx.device);
2149 ✗ detail::adopt_output(V_out, ctx.device);
2150 ✗ v.flash_attention_project_kv_int8w_fp16(ctx, Wk_int8, sk, bk,
2151 ✗ Wv_int8, sv, bv, K_out, V_out);
2152 ✗ }
2153
2154 ✗ void flash_attention_q_with_kv_cached_int8w_fp16(const Tensor& X,
2155 const Tensor& K, const Tensor& V,
2156 const Tensor& Wq_int8,
2157 const Tensor& sq,
2158 const Tensor* bq,
2159 const Tensor& Wo_int8,
2160 const Tensor& so,
2161 const Tensor* bo,
2162 const float* d_mask,
2163 int num_heads, bool causal,
2164 Tensor& O) {
2165 ✗ const auto& v = detail::dispatch_with_opts(
2166 ✗ X, K, {&V, &Wq_int8, &sq, bq, &Wo_int8, &so, bo, &O});
2167 ✗ if (!v.flash_attention_q_with_kv_cached_int8w_fp16)
2168 ✗ detail::throw_not_implemented("flash_attention_q_with_kv_cached_int8w_fp16", X.device);
2169 ✗ detail::adopt_output(O, X.device);
2170 ✗ v.flash_attention_q_with_kv_cached_int8w_fp16(X, K, V, Wq_int8, sq, bq,
2171 ✗ Wo_int8, so, bo,
2172 ✗ d_mask, num_heads, causal, O);
2173 ✗ }
2174
2175 ✗ void flash_attention_qkvo_int8w_fp16(const Tensor& X, const Tensor* Ctx,
2176 const Tensor& Wq_int8, const Tensor& sq, const Tensor* bq,
2177 const Tensor& Wk_int8, const Tensor& sk, const Tensor* bk,
2178 const Tensor& Wv_int8, const Tensor& sv, const Tensor* bv,
2179 const Tensor& Wo_int8, const Tensor& so, const Tensor* bo,
2180 const float* d_mask, int num_heads, bool causal,
2181 Tensor& O) {
2182 ✗ const auto& v = detail::dispatch_with_opts(
2183 ✗ X, Wq_int8, {Ctx, &sq, bq, &Wk_int8, &sk, bk, &Wv_int8, &sv, bv,
2184 ✗ &Wo_int8, &so, bo, &O});
2185 ✗ if (!v.flash_attention_qkvo_int8w_fp16)
2186 ✗ detail::throw_not_implemented("flash_attention_qkvo_int8w_fp16", X.device);
2187 ✗ detail::adopt_output(O, X.device);
2188 ✗ v.flash_attention_qkvo_int8w_fp16(X, Ctx, Wq_int8, sq, bq,
2189 ✗ Wk_int8, sk, bk, Wv_int8, sv, bv,
2190 ✗ Wo_int8, so, bo,
2191 ✗ d_mask, num_heads, causal, O);
2192 ✗ }
2193
2194 // ─── Spectral / FFT core (brosoundml) ──────────────────────────────────────
2195
2196 54 void complex_mul(const Tensor& a, const Tensor& b, Tensor& y) {
2197 54 const auto& v = detail::dispatch(a, b, y);
2198
1/2
✓ Branch 0 taken 54 times.
✗ Branch 1 not taken.
 54 if (!v.complex_mul) detail::throw_not_implemented("complex_mul", a.device);
2199 54 detail::adopt_output(y, a.device);
2200 54 v.complex_mul(a, b, y);
2201 54 }
2202 5 void complex_mul_backward(const Tensor& a, const Tensor& b, const Tensor& dY,
2203 Tensor& dA, Tensor& dB) {
2204 5 const auto& v = detail::dispatch(a, b, dY, dA, dB);
2205
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 if (!v.complex_mul_backward)
2206 ✗ detail::throw_not_implemented("complex_mul_backward", a.device);
2207 5 detail::adopt_output(dA, a.device);
2208 5 detail::adopt_output(dB, a.device);
2209 5 v.complex_mul_backward(a, b, dY, dA, dB);
2210 5 }
2211 36 void complex_abs(const Tensor& z, Tensor& y) {
2212 36 const auto& v = detail::dispatch(z, y);
2213
1/2
✓ Branch 0 taken 36 times.
✗ Branch 1 not taken.
 36 if (!v.complex_abs) detail::throw_not_implemented("complex_abs", z.device);
2214 36 detail::adopt_output(y, z.device);
2215 36 v.complex_abs(z, y);
2216 36 }
2217 3 void complex_abs_backward(const Tensor& z, const Tensor& dY, Tensor& dZ) {
2218 3 const auto& v = detail::dispatch(z, dY, dZ);
2219
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.complex_abs_backward)
2220 ✗ detail::throw_not_implemented("complex_abs_backward", z.device);
2221 3 detail::adopt_output(dZ, z.device);
2222 3 v.complex_abs_backward(z, dY, dZ);
2223 3 }
2224 3 void complex_angle(const Tensor& z, Tensor& y) {
2225 3 const auto& v = detail::dispatch(z, y);
2226
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.complex_angle) detail::throw_not_implemented("complex_angle", z.device);
2227 3 detail::adopt_output(y, z.device);
2228 3 v.complex_angle(z, y);
2229 3 }
2230 3 void complex_from_polar(const Tensor& mag, const Tensor& phase, Tensor& y) {
2231 3 const auto& v = detail::dispatch(mag, phase, y);
2232
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.complex_from_polar)
2233 ✗ detail::throw_not_implemented("complex_from_polar", mag.device);
2234 3 detail::adopt_output(y, mag.device);
2235 3 v.complex_from_polar(mag, phase, y);
2236 3 }
2237 339 void fft(const Tensor& x, Tensor& y) {
2238 339 const auto& v = detail::dispatch(x, y);
2239
1/2
✓ Branch 0 taken 339 times.
✗ Branch 1 not taken.
 339 if (!v.fft) detail::throw_not_implemented("fft", x.device);
2240 339 detail::adopt_output(y, x.device);
2241 339 v.fft(x, y);
2242 339 }
2243 23 void ifft(const Tensor& x, Tensor& y) {
2244 23 const auto& v = detail::dispatch(x, y);
2245
1/2
✓ Branch 0 taken 23 times.
✗ Branch 1 not taken.
 23 if (!v.ifft) detail::throw_not_implemented("ifft", x.device);
2246 23 detail::adopt_output(y, x.device);
2247 23 v.ifft(x, y);
2248 23 }
2249 335 void rfft(const Tensor& x, Tensor& y) {
2250 335 const auto& v = detail::dispatch(x, y);
2251
1/2
✓ Branch 0 taken 335 times.
✗ Branch 1 not taken.
 335 if (!v.rfft) detail::throw_not_implemented("rfft", x.device);
2252 335 detail::adopt_output(y, x.device);
2253 335 v.rfft(x, y);
2254 335 }
2255 367 void irfft(const Tensor& x, int L, Tensor& y) {
2256 367 const auto& v = detail::dispatch(x, y);
2257
1/2
✓ Branch 0 taken 367 times.
✗ Branch 1 not taken.
 367 if (!v.irfft) detail::throw_not_implemented("irfft", x.device);
2258 367 detail::adopt_output(y, x.device);
2259 367 v.irfft(x, L, y);
2260 367 }
2261 9 void rfft_backward(const Tensor& dY, int L, Tensor& dX) {
2262 9 const auto& v = detail::dispatch(dY, dX);
2263
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.rfft_backward) detail::throw_not_implemented("rfft_backward", dY.device);
2264 9 detail::adopt_output(dX, dY.device);
2265 9 v.rfft_backward(dY, L, dX);
2266 9 }
2267 9 void irfft_backward(const Tensor& dY, Tensor& dX) {
2268 9 const auto& v = detail::dispatch(dY, dX);
2269
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.irfft_backward)
2270 ✗ detail::throw_not_implemented("irfft_backward", dY.device);
2271 9 detail::adopt_output(dX, dY.device);
2272 9 v.irfft_backward(dY, dX);
2273 9 }
2274
2275 1290 void stft(const Tensor& signal, const Tensor& window,
2276 int N, int n_fft, int hop_length, int win_length,
2277 bool center, bool normalized, Tensor& spec) {
2278 1290 const auto& v = detail::dispatch(signal, window, spec);
2279
1/2
✓ Branch 0 taken 1290 times.
✗ Branch 1 not taken.
 1290 if (!v.stft) detail::throw_not_implemented("stft", signal.device);
2280 1290 detail::adopt_output(spec, signal.device);
2281 2580 v.stft(signal, window, N, n_fft, hop_length, win_length,
2282 1290 center, normalized, spec);
2283 1290 }
2284 14 void stft_backward(const Tensor& dSpec, const Tensor& window,
2285 int N, int signal_len, int n_fft, int hop_length,
2286 int win_length, bool center, bool normalized,
2287 Tensor& dSignal) {
2288 14 const auto& v = detail::dispatch(dSpec, window, dSignal);
2289
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.stft_backward)
2290 ✗ detail::throw_not_implemented("stft_backward", dSpec.device);
2291 14 detail::adopt_output(dSignal, dSpec.device);
2292 28 v.stft_backward(dSpec, window, N, signal_len, n_fft, hop_length,
2293 14 win_length, center, normalized, dSignal);
2294 14 }
2295 5114 void istft(const Tensor& spec, const Tensor& window,
2296 int N, int signal_len, int n_fft, int hop_length, int win_length,
2297 bool center, bool normalized, Tensor& signal) {
2298 5114 const auto& v = detail::dispatch(spec, window, signal);
2299
1/2
✓ Branch 0 taken 5114 times.
✗ Branch 1 not taken.
 5114 if (!v.istft) detail::throw_not_implemented("istft", spec.device);
2300 5114 detail::adopt_output(signal, spec.device);
2301 10228 v.istft(spec, window, N, signal_len, n_fft, hop_length, win_length,
2302 5114 center, normalized, signal);
2303 5114 }
2304 14 void istft_backward(const Tensor& dSignal, const Tensor& window,
2305 int N, int signal_len, int n_fft, int hop_length,
2306 int win_length, bool center, bool normalized,
2307 Tensor& dSpec) {
2308 14 const auto& v = detail::dispatch(dSignal, window, dSpec);
2309
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 if (!v.istft_backward)
2310 ✗ detail::throw_not_implemented("istft_backward", dSignal.device);
2311 14 detail::adopt_output(dSpec, dSignal.device);
2312 28 v.istft_backward(dSignal, window, N, signal_len, n_fft, hop_length,
2313 14 win_length, center, normalized, dSpec);
2314 14 }
2315
2316 // ─── 1D convolution family (brosoundml) ────────────────────────────────────
2317 //
2318 // conv1d / conv1d backward / conv1d_int8w_fp16 / causal_conv1d are header-only
2319 // inline wrappers in ops.h (they forward to the conv2d ops); only the three
2320 // genuinely new ops below have dispatcher wrappers.
2321
2322 144 void conv_transpose1d_forward(const Tensor& X, const Tensor& Wt,
2323 const Tensor* bias,
2324 int N, int C_in, int L, int C_out, int kL,
2325 int stride, int padding, int output_padding,
2326 int dilation, int groups, Tensor& Y) {
2327 144 const auto& v = detail::dispatch_with_opts(X, Wt, {bias, &Y});
2328
1/2
✓ Branch 0 taken 144 times.
✗ Branch 1 not taken.
 144 if (!v.conv_transpose1d_forward)
2329 ✗ detail::throw_not_implemented("conv_transpose1d_forward", X.device);
2330 144 detail::adopt_output(Y, X.device);
2331 288 v.conv_transpose1d_forward(X, Wt, bias, N, C_in, L, C_out, kL, stride,
2332 144 padding, output_padding, dilation, groups, Y);
2333 144 }
2334
2335 12 void conv_transpose1d_backward_input(const Tensor& Wt, const Tensor& dY,
2336 int N, int C_in, int L, int C_out, int kL,
2337 int stride, int padding,
2338 int output_padding, int dilation,
2339 int groups, Tensor& dX) {
2340 12 const auto& v = detail::dispatch(Wt, dY, dX);
2341
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.conv_transpose1d_backward_input)
2342 ✗ detail::throw_not_implemented("conv_transpose1d_backward_input", Wt.device);
2343 12 detail::adopt_output(dX, Wt.device);
2344 24 v.conv_transpose1d_backward_input(Wt, dY, N, C_in, L, C_out, kL, stride,
2345 12 padding, output_padding, dilation,
2346 12 groups, dX);
2347 12 }
2348
2349 10 void conv_transpose1d_backward_weight(const Tensor& X, const Tensor& dY,
2350 int N, int C_in, int L, int C_out, int kL,
2351 int stride, int padding,
2352 int output_padding, int dilation,
2353 int groups, Tensor& dWt) {
2354 10 const auto& v = detail::dispatch(X, dY, dWt);
2355
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.conv_transpose1d_backward_weight)
2356 ✗ detail::throw_not_implemented("conv_transpose1d_backward_weight", X.device);
2357 10 detail::adopt_output(dWt, X.device);
2358 20 v.conv_transpose1d_backward_weight(X, dY, N, C_in, L, C_out, kL, stride,
2359 10 padding, output_padding, dilation,
2360 10 groups, dWt);
2361 10 }
2362
2363 8 void conv_transpose1d_backward_bias(const Tensor& dY, int N, int C_out,
2364 int L_out, Tensor& dB) {
2365 8 const auto& v = detail::dispatch(dY, dB);
2366
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.conv_transpose1d_backward_bias)
2367 ✗ detail::throw_not_implemented("conv_transpose1d_backward_bias", dY.device);
2368 8 detail::adopt_output(dB, dY.device);
2369 8 v.conv_transpose1d_backward_bias(dY, N, C_out, L_out, dB);
2370 8 }
2371
2372 26 void causal_conv1d_update(const Tensor& X, const Tensor& Wt, const Tensor* bias,
2373 int N, int C, int L_step, int kL, int dilation,
2374 Tensor& state, Tensor& Y) {
2375 26 const auto& v = detail::dispatch_with_opts(X, Wt, {bias, &state, &Y});
2376
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 26 if (!v.causal_conv1d_update)
2377 ✗ detail::throw_not_implemented("causal_conv1d_update", X.device);
2378 26 detail::adopt_output(state, X.device);
2379 26 detail::adopt_output(Y, X.device);
2380 26 v.causal_conv1d_update(X, Wt, bias, N, C, L_step, kL, dilation, state, Y);
2381 26 }
2382
2383 157 void pad1d_forward(const Tensor& X, int N, int C, int L,
2384 int pad_left, int pad_right, int mode, Tensor& Y) {
2385 157 const auto& v = detail::dispatch(X, Y);
2386
1/2
✓ Branch 0 taken 157 times.
✗ Branch 1 not taken.
 157 if (!v.pad1d_forward)
2387 ✗ detail::throw_not_implemented("pad1d_forward", X.device);
2388 157 detail::adopt_output(Y, X.device);
2389 157 v.pad1d_forward(X, N, C, L, pad_left, pad_right, mode, Y);
2390 157 }
2391
2392 9 void pad1d_backward(const Tensor& dY, int N, int C, int L,
2393 int pad_left, int pad_right, int mode, Tensor& dX) {
2394 9 const auto& v = detail::dispatch(dY, dX);
2395
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.pad1d_backward)
2396 ✗ detail::throw_not_implemented("pad1d_backward", dY.device);
2397 9 detail::adopt_output(dX, dY.device);
2398 9 v.pad1d_backward(dY, N, C, L, pad_left, pad_right, mode, dX);
2399 9 }
2400
2401 // ─── Vocoder / codec activations (brosoundml) ──────────────────────────────
2402
2403 124 void snake_forward(const Tensor& X, const Tensor& alpha, const Tensor* beta,
2404 int N, int C, int L, Tensor& Y) {
2405 124 const auto& v = detail::dispatch_with_opts(X, alpha, {beta, &Y});
2406
1/2
✓ Branch 0 taken 124 times.
✗ Branch 1 not taken.
 124 if (!v.snake_forward)
2407 ✗ detail::throw_not_implemented("snake_forward", X.device);
2408 124 detail::adopt_output(Y, X.device);
2409 124 v.snake_forward(X, alpha, beta, N, C, L, Y);
2410 124 }
2411
2412 11 void snake_backward(const Tensor& X, const Tensor& alpha, const Tensor* beta,
2413 const Tensor& dY, int N, int C, int L,
2414 Tensor& dX, Tensor& dAlpha, Tensor* dBeta) {
2415 11 const auto& v = detail::dispatch_with_opts(
2416 11 X, alpha, {beta, &dY, &dX, &dAlpha, dBeta});
2417
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.snake_backward)
2418 ✗ detail::throw_not_implemented("snake_backward", X.device);
2419 11 detail::adopt_output(dX, X.device);
2420 11 detail::adopt_output(dAlpha, X.device);
2421
2/2
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 8 times.
 11 if (dBeta) detail::adopt_output(*dBeta, X.device);
2422 11 v.snake_backward(X, alpha, beta, dY, N, C, L, dX, dAlpha, dBeta);
2423 11 }
2424
2425 137 void elu_forward(const Tensor& x, float alpha, Tensor& y) {
2426 137 const auto& v = detail::dispatch(x, y);
2427
1/2
✓ Branch 0 taken 137 times.
✗ Branch 1 not taken.
 137 if (!v.elu_forward) detail::throw_not_implemented("elu_forward", x.device);
2428 137 detail::adopt_output(y, x.device);
2429 137 v.elu_forward(x, alpha, y);
2430 137 }
2431
2432 9 void elu_backward(const Tensor& x, const Tensor& dY, float alpha, Tensor& dX) {
2433 9 const auto& v = detail::dispatch(x, dY, dX);
2434
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.elu_backward) detail::throw_not_implemented("elu_backward", x.device);
2435 9 detail::adopt_output(dX, x.device);
2436 9 v.elu_backward(x, dY, alpha, dX);
2437 9 }
2438
2439 94 void leaky_relu_forward(const Tensor& x, float negative_slope, Tensor& y) {
2440 94 const auto& v = detail::dispatch(x, y);
2441
1/2
✓ Branch 0 taken 94 times.
✗ Branch 1 not taken.
 94 if (!v.leaky_relu_forward)
2442 ✗ detail::throw_not_implemented("leaky_relu_forward", x.device);
2443 94 detail::adopt_output(y, x.device);
2444 94 v.leaky_relu_forward(x, negative_slope, y);
2445 94 }
2446
2447 6 void leaky_relu_backward(const Tensor& x, const Tensor& dY,
2448 float negative_slope, Tensor& dX) {
2449 6 const auto& v = detail::dispatch(x, dY, dX);
2450
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.leaky_relu_backward)
2451 ✗ detail::throw_not_implemented("leaky_relu_backward", x.device);
2452 6 detail::adopt_output(dX, x.device);
2453 6 v.leaky_relu_backward(x, dY, negative_slope, dX);
2454 6 }
2455
2456 // ─── Codec quantization (brosoundml) ───────────────────────────────────────
2457
2458 9 void vq_encode_forward(const Tensor& x, const Tensor& codebook,
2459 Tensor& indices, Tensor& quantized) {
2460 9 const auto& v = detail::dispatch(x, codebook, indices, quantized);
2461
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.vq_encode_forward)
2462 ✗ detail::throw_not_implemented("vq_encode_forward", x.device);
2463 9 detail::adopt_output(indices, x.device);
2464 9 detail::adopt_output(quantized, x.device);
2465 9 v.vq_encode_forward(x, codebook, indices, quantized);
2466 9 }
2467
2468 6 void vq_encode_backward(const Tensor& dQuantized, Tensor& dX) {
2469 6 const auto& v = detail::dispatch(dQuantized, dX);
2470
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.vq_encode_backward)
2471 ✗ detail::throw_not_implemented("vq_encode_backward", dQuantized.device);
2472 6 detail::adopt_output(dX, dQuantized.device);
2473 6 v.vq_encode_backward(dQuantized, dX);
2474 6 }
2475
2476 8 void fsq_quantize_forward(const Tensor& x, const Tensor& levels,
2477 Tensor& quantized, Tensor& packed_indices) {
2478 8 const auto& v = detail::dispatch(x, levels, quantized, packed_indices);
2479
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.fsq_quantize_forward)
2480 ✗ detail::throw_not_implemented("fsq_quantize_forward", x.device);
2481 8 detail::adopt_output(quantized, x.device);
2482 8 detail::adopt_output(packed_indices, x.device);
2483 8 v.fsq_quantize_forward(x, levels, quantized, packed_indices);
2484 8 }
2485
2486 6 void fsq_quantize_backward(const Tensor& dQuantized, Tensor& dX) {
2487 6 const auto& v = detail::dispatch(dQuantized, dX);
2488
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.fsq_quantize_backward)
2489 ✗ detail::throw_not_implemented("fsq_quantize_backward", dQuantized.device);
2490 6 detail::adopt_output(dX, dQuantized.device);
2491 6 v.fsq_quantize_backward(dQuantized, dX);
2492 6 }
2493
2494 // ─── 1D resampling (brosoundml) ────────────────────────────────────────────
2495
2496 152 void resample1d_forward(const Tensor& X, int N, int C, int L_in, int L_out,
2497 int mode, Tensor& Y) {
2498 152 const auto& v = detail::dispatch(X, Y);
2499
1/2
✓ Branch 0 taken 152 times.
✗ Branch 1 not taken.
 152 if (!v.resample1d_forward)
2500 ✗ detail::throw_not_implemented("resample1d_forward", X.device);
2501 152 detail::adopt_output(Y, X.device);
2502 152 v.resample1d_forward(X, N, C, L_in, L_out, mode, Y);
2503 152 }
2504
2505 16 void resample1d_backward(const Tensor& dY, int N, int C, int L_in, int L_out,
2506 int mode, Tensor& dX) {
2507 16 const auto& v = detail::dispatch(dY, dX);
2508
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.resample1d_backward)
2509 ✗ detail::throw_not_implemented("resample1d_backward", dY.device);
2510 16 detail::adopt_output(dX, dY.device);
2511 16 v.resample1d_backward(dY, N, C, L_in, L_out, mode, dX);
2512 16 }
2513
2514 // ─── log / exp / round elementwise (brosoundml) ────────────────────────────
2515
2516 52 void log_forward(const Tensor& x, Tensor& y) {
2517 52 const auto& v = detail::dispatch(x, y);
2518
1/2
✓ Branch 0 taken 52 times.
✗ Branch 1 not taken.
 52 if (!v.log_forward) detail::throw_not_implemented("log_forward", x.device);
2519 52 detail::adopt_output(y, x.device);
2520 52 v.log_forward(x, y);
2521 52 }
2522
2523 5 void log_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
2524 5 const auto& v = detail::dispatch(x, dY, dX);
2525
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 if (!v.log_backward) detail::throw_not_implemented("log_backward", x.device);
2526 5 detail::adopt_output(dX, x.device);
2527 5 v.log_backward(x, dY, dX);
2528 5 }
2529
2530 48 void exp_forward(const Tensor& x, Tensor& y) {
2531 48 const auto& v = detail::dispatch(x, y);
2532
1/2
✓ Branch 0 taken 48 times.
✗ Branch 1 not taken.
 48 if (!v.exp_forward) detail::throw_not_implemented("exp_forward", x.device);
2533 48 detail::adopt_output(y, x.device);
2534 48 v.exp_forward(x, y);
2535 48 }
2536
2537 7 void exp_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
2538 7 const auto& v = detail::dispatch(x, dY, dX);
2539
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (!v.exp_backward) detail::throw_not_implemented("exp_backward", x.device);
2540 7 detail::adopt_output(dX, x.device);
2541 7 v.exp_backward(x, dY, dX);
2542 7 }
2543
2544 7 void round_forward(const Tensor& x, Tensor& y) {
2545 7 const auto& v = detail::dispatch(x, y);
2546
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (!v.round_forward) detail::throw_not_implemented("round_forward", x.device);
2547 7 detail::adopt_output(y, x.device);
2548 7 v.round_forward(x, y);
2549 7 }
2550
2551 3 void round_backward(const Tensor& dY, Tensor& dX) {
2552 3 const auto& v = detail::dispatch(dY, dX);
2553
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.round_backward) detail::throw_not_implemented("round_backward", dY.device);
2554 3 detail::adopt_output(dX, dY.device);
2555 3 v.round_backward(dY, dX);
2556 3 }
2557
2558 // ─── Autoregressive logit sampling (brosoundml CHUNK 7, family F) ───────────
2559
2560 50266 void sample_logits(const Tensor& logits, float temperature, int top_k,
2561 float top_p, uint64_t key, uint64_t counter,
2562 Tensor& indices) {
2563 50266 const auto& v = detail::dispatch(logits, indices);
2564
1/2
✓ Branch 0 taken 50266 times.
✗ Branch 1 not taken.
 50266 if (!v.sample_logits) detail::throw_not_implemented("sample_logits", logits.device);
2565 50266 detail::adopt_output(indices, logits.device);
2566 50266 v.sample_logits(logits, temperature, top_k, top_p, key, counter, indices);
2567 50266 }
2568
2569 13 void sample_logits_into(const Tensor& logits, float temperature, int top_k,
2570 float top_p, uint64_t key, Tensor& counter,
2571 Tensor& scratch, Tensor& indices) {
2572 13 const auto& v = detail::dispatch(logits, indices);
2573
1/2
✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
 13 if (!v.sample_logits_into)
2574 ✗ detail::throw_not_implemented("sample_logits_into", logits.device);
2575 13 detail::adopt_output(indices, logits.device);
2576 26 v.sample_logits_into(logits, temperature, top_k, top_p, key, counter,
2577 13 scratch, indices);
2578 13 }
2579
2580 // ─── L2 norm + Gated Delta Rule (linear-attention text path) ───────────────
2581
2582 73 void l2_norm_forward(const Tensor& X, int head_dim, int num_heads,
2583 float eps, Tensor& Y) {
2584 73 const auto& v = detail::dispatch(X, Y);
2585
1/2
✓ Branch 0 taken 73 times.
✗ Branch 1 not taken.
 73 if (!v.l2_norm_forward) detail::throw_not_implemented("l2_norm_forward", X.device);
2586 73 detail::adopt_output(Y, X.device);
2587 73 v.l2_norm_forward(X, head_dim, num_heads, eps, Y);
2588 73 }
2589
2590 11 void l2_norm_backward(const Tensor& X, int head_dim, int num_heads,
2591 float eps, const Tensor& dY, Tensor& dX) {
2592 11 const auto& v = detail::dispatch(X, dY, dX);
2593
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.l2_norm_backward) detail::throw_not_implemented("l2_norm_backward", X.device);
2594 11 detail::adopt_output(dX, X.device);
2595 11 v.l2_norm_backward(X, head_dim, num_heads, eps, dY, dX);
2596 11 }
2597
2598 11 void gated_delta_rule_chunked(const Tensor& Q, const Tensor& K, const Tensor& V,
2599 const Tensor& a_raw, const Tensor& beta,
2600 const Tensor& log_A,
2601 int num_heads, int d_k, int d_v,
2602 Tensor& state, Tensor& O) {
2603 11 const auto& v = detail::dispatch(Q, K, V, a_raw, beta, log_A, state, O);
2604
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.gated_delta_rule_chunked)
2605 ✗ detail::throw_not_implemented("gated_delta_rule_chunked", Q.device);
2606 11 detail::adopt_output(state, Q.device);
2607 11 detail::adopt_output(O, Q.device);
2608 22 v.gated_delta_rule_chunked(Q, K, V, a_raw, beta, log_A,
2609 11 num_heads, d_k, d_v, state, O);
2610 11 }
2611
2612 12 void gated_delta_rule_step(const Tensor& Q, const Tensor& K, const Tensor& V,
2613 const Tensor& a_raw, const Tensor& beta,
2614 const Tensor& log_A,
2615 int num_heads, int d_k, int d_v,
2616 Tensor& state, Tensor& O) {
2617 12 const auto& v = detail::dispatch(Q, K, V, a_raw, beta, log_A, state, O);
2618
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.gated_delta_rule_step)
2619 ✗ detail::throw_not_implemented("gated_delta_rule_step", Q.device);
2620 12 detail::adopt_output(state, Q.device);
2621 12 detail::adopt_output(O, Q.device);
2622 24 v.gated_delta_rule_step(Q, K, V, a_raw, beta, log_A,
2623 12 num_heads, d_k, d_v, state, O);
2624 12 }
2625
2626 // ─── BatchNorm ─────────────────────────────────────────────────────────────
2627
2628 9 void batch_norm_forward(const Tensor& X,
2629 const Tensor& gamma, const Tensor& beta,
2630 Tensor& running_mean, Tensor& running_var,
2631 int N, int C, int H, int W,
2632 float eps, float momentum,
2633 Tensor& Y,
2634 Tensor& saved_mean, Tensor& saved_rstd) {
2635 18 const auto& v = detail::dispatch(X, gamma, beta, running_mean, running_var,
2636 9 Y, saved_mean, saved_rstd);
2637
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.batch_norm_forward)
2638 ✗ detail::throw_not_implemented("batch_norm_forward", X.device);
2639 9 detail::adopt_output(running_mean, X.device);
2640 9 detail::adopt_output(running_var, X.device);
2641 9 detail::adopt_output(Y, X.device);
2642 9 detail::adopt_output(saved_mean, X.device);
2643 9 detail::adopt_output(saved_rstd, X.device);
2644 18 v.batch_norm_forward(X, gamma, beta, running_mean, running_var,
2645 9 N, C, H, W, eps, momentum,
2646 9 Y, saved_mean, saved_rstd);
2647 9 }
2648
2649 11 void batch_norm_inference(const Tensor& X,
2650 const Tensor& gamma, const Tensor& beta,
2651 const Tensor& running_mean,
2652 const Tensor& running_var,
2653 int N, int C, int H, int W,
2654 float eps,
2655 Tensor& Y) {
2656 11 const auto& v = detail::dispatch(X, gamma, beta, running_mean, running_var, Y);
2657
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.batch_norm_inference)
2658 ✗ detail::throw_not_implemented("batch_norm_inference", X.device);
2659 11 detail::adopt_output(Y, X.device);
2660 22 v.batch_norm_inference(X, gamma, beta, running_mean, running_var,
2661 11 N, C, H, W, eps, Y);
2662 11 }
2663
2664 5 void batch_norm_backward(const Tensor& X,
2665 const Tensor& gamma,
2666 const Tensor& saved_mean,
2667 const Tensor& saved_rstd,
2668 const Tensor& dY,
2669 int N, int C, int H, int W,
2670 Tensor& dX,
2671 Tensor& dGamma, Tensor& dBeta) {
2672 10 const auto& v = detail::dispatch(X, gamma, saved_mean, saved_rstd, dY,
2673 5 dX, dGamma, dBeta);
2674
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 5 if (!v.batch_norm_backward)
2675 ✗ detail::throw_not_implemented("batch_norm_backward", X.device);
2676 5 detail::adopt_output(dX, X.device);
2677 5 detail::adopt_output(dGamma, X.device);
2678 5 detail::adopt_output(dBeta, X.device);
2679 10 v.batch_norm_backward(X, gamma, saved_mean, saved_rstd, dY,
2680 5 N, C, H, W, dX, dGamma, dBeta);
2681 5 }
2682
2683 // ─── Image preprocessing helpers ───────────────────────────────────────────
2684
2685 6 void image_normalize(const Tensor& X,
2686 const Tensor& mean, const Tensor& std_,
2687 int N, int C, int H, int W,
2688 Tensor& Y) {
2689 6 const auto& v = detail::dispatch(X, mean, std_, Y);
2690
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.image_normalize)
2691 ✗ detail::throw_not_implemented("image_normalize", X.device);
2692 6 detail::adopt_output(Y, X.device);
2693 6 v.image_normalize(X, mean, std_, N, C, H, W, Y);
2694 6 }
2695
2696 9 void image_u8_to_f32_nhwc_to_nchw(const uint8_t* src,
2697 int N, int H, int W, int C,
2698 float scale, float bias,
2699 Tensor& Y) {
2700 // No tensor inputs to dispatch on — Y is the only one. adopt_output
2701 // pins an uncommitted Y to the default device before lookup.
2702 9 const auto& v = detail::dispatch(Y);
2703
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.image_u8_to_f32_nhwc_to_nchw)
2704 ✗ detail::throw_not_implemented("image_u8_to_f32_nhwc_to_nchw", Y.device);
2705 9 detail::adopt_output(Y, Y.device);
2706 9 v.image_u8_to_f32_nhwc_to_nchw(src, N, H, W, C, scale, bias, Y);
2707 9 }
2708
2709 // ─── Counter-based noise generation (Philox 4x32-10) ───────────────────────
2710
2711 6 void randn(uint64_t key, uint64_t counter, Tensor& Y) {
2712 6 const auto& v = detail::dispatch(Y);
2713
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!v.randn) detail::throw_not_implemented("randn", Y.device);
2714 6 detail::adopt_output(Y, Y.device);
2715 6 v.randn(key, counter, Y);
2716 6 }
2717
2718 4 void rand_uniform(uint64_t key, uint64_t counter, Tensor& Y) {
2719 4 const auto& v = detail::dispatch(Y);
2720
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!v.rand_uniform) detail::throw_not_implemented("rand_uniform", Y.device);
2721 4 detail::adopt_output(Y, Y.device);
2722 4 v.rand_uniform(key, counter, Y);
2723 4 }
2724
2725 4 void rand_bernoulli(float p, uint64_t key, uint64_t counter, Tensor& Y) {
2726 4 const auto& v = detail::dispatch(Y);
2727
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!v.rand_bernoulli) detail::throw_not_implemented("rand_bernoulli", Y.device);
2728 4 detail::adopt_output(Y, Y.device);
2729 4 v.rand_bernoulli(p, key, counter, Y);
2730 4 }
2731
2732 2 void randn_truncated(float lo, float hi, uint64_t key, uint64_t counter,
2733 Tensor& Y) {
2734 2 const auto& v = detail::dispatch(Y);
2735
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!v.randn_truncated) detail::throw_not_implemented("randn_truncated", Y.device);
2736 2 detail::adopt_output(Y, Y.device);
2737 2 v.randn_truncated(lo, hi, key, counter, Y);
2738 2 }
2739
2740 // ─── StyleGAN3 synthesis-input primitives — sin/cos/rsqrt + pixel_norm ──────
2741
2742 32 void sin_forward(const Tensor& x, Tensor& y) {
2743 32 const auto& v = detail::dispatch(x, y);
2744
1/2
✓ Branch 0 taken 32 times.
✗ Branch 1 not taken.
 32 if (!v.sin_forward) detail::throw_not_implemented("sin_forward", x.device);
2745 32 detail::adopt_output(y, x.device);
2746 32 v.sin_forward(x, y);
2747 32 }
2748
2749 3 void sin_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
2750 3 const auto& v = detail::dispatch(x, dY, dX);
2751
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.sin_backward) detail::throw_not_implemented("sin_backward", x.device);
2752 3 detail::adopt_output(dX, x.device);
2753 3 v.sin_backward(x, dY, dX);
2754 3 }
2755
2756 32 void cos_forward(const Tensor& x, Tensor& y) {
2757 32 const auto& v = detail::dispatch(x, y);
2758
1/2
✓ Branch 0 taken 32 times.
✗ Branch 1 not taken.
 32 if (!v.cos_forward) detail::throw_not_implemented("cos_forward", x.device);
2759 32 detail::adopt_output(y, x.device);
2760 32 v.cos_forward(x, y);
2761 32 }
2762
2763 3 void cos_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
2764 3 const auto& v = detail::dispatch(x, dY, dX);
2765
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!v.cos_backward) detail::throw_not_implemented("cos_backward", x.device);
2766 3 detail::adopt_output(dX, x.device);
2767 3 v.cos_backward(x, dY, dX);
2768 3 }
2769
2770 32 void rsqrt_forward(const Tensor& x, Tensor& y) {
2771 32 const auto& v = detail::dispatch(x, y);
2772
1/2
✓ Branch 0 taken 32 times.
✗ Branch 1 not taken.
 32 if (!v.rsqrt_forward) detail::throw_not_implemented("rsqrt_forward", x.device);
2773 32 detail::adopt_output(y, x.device);
2774 32 v.rsqrt_forward(x, y);
2775 32 }
2776
2777 7 void rsqrt_backward(const Tensor& y, const Tensor& dY, Tensor& dX) {
2778 7 const auto& v = detail::dispatch(y, dY, dX);
2779
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (!v.rsqrt_backward) detail::throw_not_implemented("rsqrt_backward", y.device);
2780 7 detail::adopt_output(dX, y.device);
2781 7 v.rsqrt_backward(y, dY, dX);
2782 7 }
2783
2784 55 void pixel_norm_forward(const Tensor& X, float eps, Tensor& Y) {
2785 55 const auto& v = detail::dispatch(X, Y);
2786
1/2
✓ Branch 0 taken 55 times.
✗ Branch 1 not taken.
 55 if (!v.pixel_norm_forward) detail::throw_not_implemented("pixel_norm_forward", X.device);
2787 55 detail::adopt_output(Y, X.device);
2788 55 v.pixel_norm_forward(X, eps, Y);
2789 55 }
2790
2791 7 void pixel_norm_backward(const Tensor& X, const Tensor& dY, float eps, Tensor& dX) {
2792 7 const auto& v = detail::dispatch(X, dY, dX);
2793
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (!v.pixel_norm_backward) detail::throw_not_implemented("pixel_norm_backward", X.device);
2794 7 detail::adopt_output(dX, X.device);
2795 7 v.pixel_norm_backward(X, dY, eps, dX);
2796 7 }
2797
2798 // ─── StyleGAN3 bias_act ─────────────────────────────────────────────────────
2799
2800 42 void bias_act_forward(const Tensor& X, const Tensor* b,
2801 int N, int C, int HW, int act, float alpha,
2802 float gain, float clamp, Tensor& Y) {
2803 42 const auto& v = detail::dispatch_with_opts(X, {b, &Y});
2804
1/2
✓ Branch 0 taken 42 times.
✗ Branch 1 not taken.
 42 if (!v.bias_act_forward) detail::throw_not_implemented("bias_act_forward", X.device);
2805 42 detail::adopt_output(Y, X.device);
2806 42 v.bias_act_forward(X, b, N, C, HW, act, alpha, gain, clamp, Y);
2807 42 }
2808
2809 42 void bias_act_backward(const Tensor& dY, const Tensor& X, const Tensor* b,
2810 int N, int C, int HW, int act, float alpha,
2811 float gain, float clamp, Tensor& dX, Tensor* dB) {
2812 42 const auto& v = detail::dispatch_with_opts(dY, X, {b, dB, &dX});
2813
1/2
✓ Branch 0 taken 42 times.
✗ Branch 1 not taken.
 42 if (!v.bias_act_backward) detail::throw_not_implemented("bias_act_backward", dY.device);
2814 42 detail::adopt_output(dX, dY.device);
2815
2/2
✓ Branch 0 taken 20 times.
✓ Branch 1 taken 22 times.
 42 if (dB) detail::adopt_output(*dB, dY.device);
2816 42 v.bias_act_backward(dY, X, b, N, C, HW, act, alpha, gain, clamp, dX, dB);
2817 42 }
2818
2819 // ─── StyleGAN3 upfirdn2d ────────────────────────────────────────────────────
2820
2821 25 void upfirdn2d_forward(const Tensor& X, const Tensor& f,
2822 int N, int C, int H, int Wd, int fH, int fW,
2823 int up_x, int up_y, int down_x, int down_y,
2824 int pad_x0, int pad_x1, int pad_y0, int pad_y1,
2825 bool flip_filter, float gain, Tensor& Y) {
2826 25 const auto& v = detail::dispatch(X, f, Y);
2827
1/2
✓ Branch 0 taken 25 times.
✗ Branch 1 not taken.
 25 if (!v.upfirdn2d_forward) detail::throw_not_implemented("upfirdn2d_forward", X.device);
2828 25 detail::adopt_output(Y, X.device);
2829 50 v.upfirdn2d_forward(X, f, N, C, H, Wd, fH, fW, up_x, up_y, down_x, down_y,
2830 25 pad_x0, pad_x1, pad_y0, pad_y1, flip_filter, gain, Y);
2831 25 }
2832
2833 21 void upfirdn2d_backward(const Tensor& dY, const Tensor& f,
2834 int N, int C, int H, int Wd, int fH, int fW,
2835 int up_x, int up_y, int down_x, int down_y,
2836 int pad_x0, int pad_x1, int pad_y0, int pad_y1,
2837 bool flip_filter, float gain, Tensor& dX) {
2838 21 const auto& v = detail::dispatch(dY, f, dX);
2839
1/2
✓ Branch 0 taken 21 times.
✗ Branch 1 not taken.
 21 if (!v.upfirdn2d_backward) detail::throw_not_implemented("upfirdn2d_backward", dY.device);
2840 21 detail::adopt_output(dX, dY.device);
2841 42 v.upfirdn2d_backward(dY, f, N, C, H, Wd, fH, fW, up_x, up_y, down_x, down_y,
2842 21 pad_x0, pad_x1, pad_y0, pad_y1, flip_filter, gain, dX);
2843 21 }
2844
2845 // ─── StyleGAN3 modulated_conv2d ─────────────────────────────────────────────
2846
2847 12 void modulated_conv2d_forward(const Tensor& X, const Tensor& W, const Tensor& s,
2848 int N, int C_in, int H, int Wd,
2849 int C_out, int kH, int kW,
2850 int pad_h, int pad_w,
2851 bool demodulate, float eps,
2852 Tensor& dcoef, Tensor& Y) {
2853 12 const auto& v = detail::dispatch(X, W, s, dcoef, Y);
2854
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.modulated_conv2d_forward)
2855 ✗ detail::throw_not_implemented("modulated_conv2d_forward", X.device);
2856 12 detail::adopt_output(dcoef, X.device);
2857 12 detail::adopt_output(Y, X.device);
2858 24 v.modulated_conv2d_forward(X, W, s, N, C_in, H, Wd, C_out, kH, kW,
2859 12 pad_h, pad_w, demodulate, eps, dcoef, Y);
2860 12 }
2861
2862 16 void modulated_conv2d_backward(const Tensor& X, const Tensor& W, const Tensor& s,
2863 const Tensor& dcoef, const Tensor& dY,
2864 int N, int C_in, int H, int Wd,
2865 int C_out, int kH, int kW,
2866 int pad_h, int pad_w, bool demodulate, float eps,
2867 Tensor& dX, Tensor& dW, Tensor& ds) {
2868 16 const auto& v = detail::dispatch(X, W, s, dcoef, dY, dX, dW, ds);
2869
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.modulated_conv2d_backward)
2870 ✗ detail::throw_not_implemented("modulated_conv2d_backward", X.device);
2871 16 detail::adopt_output(dX, X.device);
2872 16 detail::adopt_output(dW, X.device);
2873 16 detail::adopt_output(ds, X.device);
2874 32 v.modulated_conv2d_backward(X, W, s, dcoef, dY, N, C_in, H, Wd, C_out, kH, kW,
2875 16 pad_h, pad_w, demodulate, eps, dX, dW, ds);
2876 16 }
2877
2878 } // namespace brotensor
2879