GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 87.9% 261 / 0 / 297
Functions: 100.0% 17 / 0 / 17
Branches: 34.9% 286 / 0 / 819
src/api/native_tensor_extra_nn.cpp
Line Branch Exec Source
1 // native_tensor_extra_nn.cpp — the dense / recurrent / attention half of the
2 // never-bound ops: linearForwardBatchedEx, linearForwardBatchedFp16Act,
3 // matmulAbt, the LSTM training pair, GQA flash attention, the packed-QKV
4 // encoder attention pair, the Transformer-XL relative-position bias and the
5 // two RoPE variants (per-head tables, in-place over a packed QKV buffer).
6 //
7 // Same mechanism as native_tensor_extra.cpp: every operand is checked against
8 // the dims it is handed before the op runs, and index operands whose values
9 // address memory (sequence bounds, RoPE positions) are read and range-checked
10 // on the host first.
11
12 #include "native_tensor_extra_decl.h"
13 #include "native_tensor_extra_util.h"
14 #include "api_internal.h"
15 #include "native_register.h"
16
17 #include <algorithm>
18 #include <cmath>
19 #include <string>
20 #include <vector>
21
22 using namespace brotensor::api;
23 using namespace brotensor::api::extra;
24
25 namespace {
26
27 // A LinearActivation value (brotensor/ops/linear.h).
28 20 bool needAct(const char* label, int32_t act) {
29 20 return needRange(label, "act", act, brotensor::kLinearActNone, brotensor::kLinearActQuickGelu);
30 }
31
32 // A JS number carrying a `long long` element stride: whole and non-negative.
33 24 bool strideArg(const char* label, const char* name, double v, int64_t& out) {
34
3/6
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 24 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 24 times.
 24 if (!(v >= 0.0) || v != std::floor(v) || v > 9007199254740991.0) {
35 ✗ setError(std::string(label) + ": " + name + " must be a non-negative whole number");
36 ✗ return false;
37 }
38 24 out = static_cast<int64_t>(v);
39 24 return true;
40 24 }
41
42 // Elements a batch of `batch` slices of `slice` elements, `stride` apart, spans.
43 20 int64_t spanOf(int64_t batch, int64_t stride, int64_t slice) {
44
2/4
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 20 times.
 20 if (batch <= 0 || slice <= 0) return 0;
45 20 const int64_t head = prod({batch - 1, stride});
46
2/4
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 20 times.
 20 if (head == INT64_MAX || head > INT64_MAX - slice) return INT64_MAX;
47 20 return head + slice;
48 20 }
49
50 } // namespace
51
52 extern "C" {
53
54 // ---- batched linear with a fused epilogue ------------------------------------
55
56 8 void bro_tensor_linearForwardBatchedEx(void* W, uint64_t bias_bits, void* X, int32_t act, int32_t epilogue,
57 uint64_t workspace_bits, void* Y) {
58 8 constexpr const char* lb = "linearForwardBatchedEx";
59
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (!need(lb, {W, X, Y})) return;
60 8 auto* w = toTensor(W);
61 8 auto* x = toTensor(X);
62 8 auto* y = toTensor(Y);
63 8 auto* b = tensorFromValue(bias_bits);
64 8 auto* ws = tensorFromValue(workspace_bits);
65
3/6
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 8 times.
 8 if (!needFloat(lb, "W", w) || !needSameDtype(lb, "X", x, w) || !needAct(lb, act)) return;
66
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 8 if (w->empty()) { setError(std::string(lb) + ": W is empty"); return; }
67
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 6 times.
 8 if (x->cols != w->cols) {
68
7/14
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 2 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 2 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 2 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 2 times.
✗ Branch 13 not taken.
 2 setError(std::string(lb) + ": X has " + std::to_string(x->cols) + " columns, W has " + std::to_string(w->cols));
69 2 return;
70 }
71
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 2 times.
 6 if (!optElems(lb, "bias", b, w->rows, w)) return;
72 4 const int32_t base = epilogue & ~static_cast<int32_t>(brotensor::kLinearEpiFastAccum);
73
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 2 times.
 4 if (base == brotensor::kLinearEpiAccumulate) {
74
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needSameDtype(lb, "Y", y, w)) return;
75
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (y->rows != x->rows || y->cols != w->rows) {
76 ✗ setError(std::string(lb) + ": the accumulate epilogue needs Y already (" + std::to_string(x->rows) +
77 ✗ ", " + std::to_string(w->rows) + ")");
78 ✗ return;
79 }
80
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 4 } else if (base == brotensor::kLinearEpiGeglu) {
81 ✗ if (act != 0 || (w->rows & 1)) {
82 ✗ setError(std::string(lb) + ": the GeGLU epilogue needs act 0 and an even W row count");
83 ✗ return;
84 }
85
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 } else if (base != brotensor::kLinearEpiStore) {
86 ✗ setError(std::string(lb) + ": unknown epilogue " + std::to_string(epilogue));
87 ✗ return;
88 }
89 BROTENSOR_API_TRY
90
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 brotensor::linear_forward_batched_ex(*w, b, *x, act, epilogue, ws, *y);
91
0/12
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
 4 BROTENSOR_API_CATCH(lb)
92 8 }
93
94 4 void bro_tensor_linearForwardBatchedFp16Act(void* W, uint64_t bias_bits, void* X, int32_t act, void* Y) {
95 4 constexpr const char* lb = "linearForwardBatchedFp16Act";
96
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (!need(lb, {W, X, Y})) return;
97 4 auto* w = toTensor(W);
98 4 auto* x = toTensor(X);
99 4 auto* b = tensorFromValue(bias_bits);
100
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 4 if (w->dtype != brotensor::Dtype::FP16 && w->dtype != brotensor::Dtype::BF16) {
101 ✗ setError(std::string(lb) + ": W is " + dtypeToString(w->dtype) + ", expected fp16 or bf16");
102 ✗ return;
103 }
104
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
 4 if (!needSameDtype(lb, "X", x, w) || !needAct(lb, act)) return;
105
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (x->cols != w->cols) {
106 ✗ setError(std::string(lb) + ": X has " + std::to_string(x->cols) + " columns, W has " + std::to_string(w->cols));
107 ✗ return;
108 }
109
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 2 times.
 4 if (!optElems(lb, "bias", b, w->rows, w)) return;
110 BROTENSOR_API_TRY
111
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 brotensor::linear_forward_batched_fp16_act(*w, b, *x, act, *toTensor(Y));
112
0/12
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
 2 BROTENSOR_API_CATCH(lb)
113 4 }
114
115 // ---- batched A @ B^T ---------------------------------------------------------
116
117 8 void bro_tensor_matmulAbt(void* A, void* B, void* C, int32_t batch, int32_t M, int32_t N, int32_t K,
118 double strideA, double strideB, double strideC, uint64_t bias_bits, int32_t act) {
119 8 constexpr const char* lb = "matmulAbt";
120
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (!need(lb, {A, B, C})) return;
121 8 auto* a = toTensor(A);
122 8 auto* b = toTensor(B);
123 8 auto* c = toTensor(C);
124 8 auto* bias = tensorFromValue(bias_bits);
125
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 8 if (a->dtype != brotensor::Dtype::FP16 && a->dtype != brotensor::Dtype::BF16) {
126 ✗ setError(std::string(lb) + ": A is " + dtypeToString(a->dtype) + ", expected fp16 or bf16");
127 ✗ return;
128 }
129
3/6
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 8 times.
 8 if (!needSameDtype(lb, "B", b, a) || !needSameDtype(lb, "C", c, a) || !needAct(lb, act)) return;
130
4/8
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 8 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 8 times.
 8 if (batch < 1 || M < 0 || N < 0 || K < 0) {
131 ✗ setError(std::string(lb) + ": batch must be >= 1 and M, N, K >= 0");
132 ✗ return;
133 }
134 8 int64_t sA = 0, sB = 0, sC = 0;
135
3/6
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 8 times.
 8 if (!strideArg(lb, "strideA", strideA, sA) || !strideArg(lb, "strideB", strideB, sB) ||
136 8 !strideArg(lb, "strideC", strideC, sC)) return;
137 // C is not resized by the op: it must already hold every slice it writes.
138
4/4
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 4 times.
✓ Branch 3 taken 4 times.
 14 if (!needElems(lb, "A", a, spanOf(batch, sA, prod({M, K}))) ||
139
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 !needElems(lb, "B", b, spanOf(batch, sB, prod({N, K}))) ||
140
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 4 times.
 6 !needElems(lb, "C", c, spanOf(batch, sC, prod({M, N}))) ||
141 8 !optElems(lb, "bias", bias, N, a)) return;
142 BROTENSOR_API_TRY
143
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 brotensor::matmul_abt(*a, *b, *c, batch, M, N, K, sA, sB, sC, bias, act);
144
0/12
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
 4 BROTENSOR_API_CATCH(lb)
145 8 }
146
147 } // extern "C"
148
149 // ---- LSTM (training forward + BPTT) ------------------------------------------
150
151 namespace {
152
153 // The shape the LSTM pair shares: FP32 X / W_ih / W_hh, H and I from the
154 // weights, X (T*B, I). Answers H and I.
155 12 bool lstmShapes(const char* lb, brotensor::Tensor* X, brotensor::Tensor* Wih, brotensor::Tensor* Whh,
156 int32_t T, int32_t B, int64_t& H, int64_t& I) {
157 12 constexpr auto F = brotensor::Dtype::FP32;
158
3/6
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 12 times.
 12 if (!needDtype(lb, "X", X, F) || !needDtype(lb, "W_ih", Wih, F) || !needDtype(lb, "W_hh", Whh, F)) return false;
159
2/8
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 12 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
 12 if (T < 1 || B < 1) { setError(std::string(lb) + ": T and B must be >= 1"); return false; }
160
2/8
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 12 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
 12 if (Wih->rows < 4 || Wih->rows % 4 != 0) { setError(std::string(lb) + ": W_ih must have 4*H rows"); return false; }
161 12 H = Wih->rows / 4;
162 12 I = Wih->cols;
163
2/4
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 12 times.
 12 if (Whh->rows != Wih->rows || Whh->cols != H) {
164 ✗ setError(std::string(lb) + ": W_hh must be (4H, H) = (" + std::to_string(4 * H) + ", " + std::to_string(H) + ")");
165 ✗ return false;
166 }
167
3/4
✗ Branch 0 not taken.
✓ Branch 1 taken 12 times.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 10 times.
 12 if (X->cols != I || !needElems(lb, "X", X, prod({T, B, I}))) {
168
1/10
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
 2 if (X->cols != I) setError(std::string(lb) + ": X must be (T*B, I) with I = W_ih.cols = " + std::to_string(I));
169 2 return false;
170 }
171 10 return true;
172 12 }
173
174 } // namespace
175
176 extern "C" {
177
178 6 void bro_tensor_lstmForwardTrain(void* X, void* W_ih, void* W_hh, uint64_t b_ih_bits, uint64_t b_hh_bits,
179 uint64_t h0_bits, uint64_t c0_bits, int32_t T, int32_t B,
180 void* Y, void* gates, void* C, uint64_t hT_bits, uint64_t cT_bits) {
181 6 constexpr const char* lb = "lstmForwardTrain";
182
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (!need(lb, {X, W_ih, W_hh, Y, gates, C})) return;
183 6 int64_t H = 0, I = 0;
184
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 2 times.
 6 if (!lstmShapes(lb, toTensor(X), toTensor(W_ih), toTensor(W_hh), T, B, H, I)) return;
185 4 auto* bih = tensorFromValue(b_ih_bits);
186 4 auto* bhh = tensorFromValue(b_hh_bits);
187 4 auto* h0 = tensorFromValue(h0_bits);
188 4 auto* c0 = tensorFromValue(c0_bits);
189
5/6
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 2 times.
 6 if (!optF32Elems(lb, "b_ih", bih, 4 * H) || !optF32Elems(lb, "b_hh", bhh, 4 * H) ||
190
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 4 !optF32Elems(lb, "h0", h0, prod({B, H})) || !optF32Elems(lb, "c0", c0, prod({B, H}))) return;
191 BROTENSOR_API_TRY
192
4/8
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 2 times.
✗ Branch 7 not taken.
 4 brotensor::lstm_forward_train(*toTensor(X), *toTensor(W_ih), *toTensor(W_hh), bih, bhh, h0, c0, T, B,
193
3/6
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
 2 *toTensor(Y), *toTensor(gates), *toTensor(C),
194
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 tensorFromValue(hT_bits), tensorFromValue(cT_bits));
195
0/12
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
 2 BROTENSOR_API_CATCH(lb)
196 6 }
197
198 6 void bro_tensor_lstmBackward(void* X, void* W_ih, void* W_hh, uint64_t h0_bits, uint64_t c0_bits,
199 void* Y, void* gates, void* C, void* dY, int32_t T, int32_t B,
200 void* dX, void* dW_ih, void* dW_hh, uint64_t db_ih_bits, uint64_t db_hh_bits,
201 uint64_t dh0_bits, uint64_t dc0_bits) {
202 6 constexpr const char* lb = "lstmBackward";
203
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (!need(lb, {X, W_ih, W_hh, Y, gates, C, dY, dX, dW_ih, dW_hh})) return;
204 6 int64_t H = 0, I = 0;
205
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!lstmShapes(lb, toTensor(X), toTensor(W_ih), toTensor(W_hh), T, B, H, I)) return;
206 6 constexpr auto F = brotensor::Dtype::FP32;
207 6 const int64_t TB = prod({T, B});
208 6 auto* h0 = tensorFromValue(h0_bits);
209 6 auto* c0 = tensorFromValue(c0_bits);
210 6 auto* dbih = tensorFromValue(db_ih_bits);
211 6 auto* dbhh = tensorFromValue(db_hh_bits);
212
4/6
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 4 times.
✓ Branch 5 taken 2 times.
 12 if (!optF32Elems(lb, "h0", h0, prod({B, H})) || !optF32Elems(lb, "c0", c0, prod({B, H})) ||
213
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
 6 !needDtype(lb, "Y", toTensor(Y), F) || !needElems(lb, "Y", toTensor(Y), prod({TB, H})) ||
214
3/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✓ Branch 3 taken 2 times.
 6 !needDtype(lb, "gates", toTensor(gates), F) || !needElems(lb, "gates", toTensor(gates), prod({TB, 4 * H})) ||
215
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 !needDtype(lb, "C", toTensor(C), F) || !needElems(lb, "C", toTensor(C), prod({TB, H})) ||
216
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 !needDtype(lb, "dY", toTensor(dY), F) || !needElems(lb, "dY", toTensor(dY), prod({TB, H})) ||
217 // The parameter gradients accumulate, so they are read: check them too.
218
3/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 2 times.
 4 !needDtype(lb, "dW_ih", toTensor(dW_ih), F) || !needElems(lb, "dW_ih", toTensor(dW_ih), prod({4 * H, I})) ||
219
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 !needDtype(lb, "dW_hh", toTensor(dW_hh), F) || !needElems(lb, "dW_hh", toTensor(dW_hh), prod({4 * H, H})) ||
220
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 6 !optF32Elems(lb, "db_ih", dbih, 4 * H) || !optF32Elems(lb, "db_hh", dbhh, 4 * H)) return;
221 BROTENSOR_API_TRY
222
4/8
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 2 times.
✗ Branch 7 not taken.
 4 brotensor::lstm_backward(*toTensor(X), *toTensor(W_ih), *toTensor(W_hh), h0, c0,
223
4/8
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 2 times.
✗ Branch 7 not taken.
 2 *toTensor(Y), *toTensor(gates), *toTensor(C), *toTensor(dY), T, B,
224
3/6
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
 2 *toTensor(dX), *toTensor(dW_ih), *toTensor(dW_hh), dbih, dbhh,
225
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 tensorFromValue(dh0_bits), tensorFromValue(dc0_bits));
226
0/12
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
 2 BROTENSOR_API_CATCH(lb)
227 6 }
228
229 // ---- GQA flash attention -------------------------------------------------------
230
231 8 void bro_tensor_flashAttentionGqaForward(void* Q, void* K, void* V, uint64_t mask_bits, int32_t numQHeads,
232 int32_t numKvHeads, bool causal, void* O) {
233 8 constexpr const char* lb = "flashAttentionGqaForward";
234
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (!need(lb, {Q, K, V, O})) return;
235 8 auto* q = toTensor(Q);
236 8 auto* k = toTensor(K);
237 8 auto* v = toTensor(V);
238
3/6
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 8 times.
 8 if (!needFloat(lb, "Q", q) || !needSameDtype(lb, "K", k, q) || !needSameDtype(lb, "V", v, q)) return;
239
3/6
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 8 times.
 8 if (numQHeads < 1 || numKvHeads < 1 || numQHeads % numKvHeads != 0) {
240 ✗ setError(std::string(lb) + ": numKvHeads must divide numQHeads (both >= 1)");
241 ✗ return;
242 }
243
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 8 if (q->cols % numQHeads != 0) { setError(std::string(lb) + ": Q.cols must be numQHeads * headDim"); return; }
244 8 const int64_t hd = q->cols / numQHeads;
245
4/6
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 6 times.
 8 if (k->cols != numKvHeads * hd || v->cols != k->cols || v->rows != k->rows) {
246
4/8
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 2 times.
✗ Branch 7 not taken.
 2 setError(std::string(lb) + ": K and V must both be (Lk, numKvHeads*headDim) with headDim " + std::to_string(hd));
247 2 return;
248 }
249
2/6
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 6 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 6 if (k->rows < q->rows || (causal && k->rows != q->rows)) {
250 ✗ setError(std::string(lb) + ": needs Lk >= Lq (and Lk == Lq when causal)");
251 ✗ return;
252 }
253
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 2 times.
 6 if (!optF32Elems(lb, "mask", tensorFromValue(mask_bits), k->rows)) return;
254 BROTENSOR_API_TRY
255
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 brotensor::flash_attention_gqa_forward(*q, *k, *v, maskPtr(mask_bits), numQHeads, numKvHeads, causal,
256
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 *toTensor(O));
257
0/12
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
 4 BROTENSOR_API_CATCH(lb)
258 8 }
259
260 } // extern "C"
261
262 // ---- packed-QKV encoder attention ----------------------------------------------
263
264 namespace {
265
266 // QKV (L, 3*numHeads*headDim) float storage, seqBounds (L, 2) whose row r is
267 // a sequence [b0, b1) inside [0, L) containing r. Answers the INT32 operand
268 // the op reads (seqBounds itself, or a converted copy in `scratch`).
269 8 const brotensor::Tensor* packedBounds(const char* lb, brotensor::Tensor* qkv, int32_t numHeads,
270 brotensor::Tensor* sb, brotensor::Tensor& scratch) {
271
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (!needFloat(lb, "QKV", qkv)) return nullptr;
272
3/6
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 8 times.
 8 if (numHeads < 1 || qkv->cols % (3 * static_cast<int64_t>(numHeads)) != 0 || qkv->cols == 0) {
273 ✗ setError(std::string(lb) + ": QKV.cols must be 3 * numHeads * headDim");
274 ✗ return nullptr;
275 }
276 8 const int L = qkv->rows;
277
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
 8 if (sb->rows != L || sb->cols != 2) {
278 ✗ setError(std::string(lb) + ": seqBounds must be (" + std::to_string(L) + ", 2)");
279 ✗ return nullptr;
280 }
281 8 std::vector<int32_t> vals;
282
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
 8 if (!hostInt32(lb, "seqBounds", sb, prod({L, 2}), vals)) return nullptr;
283
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 24 times.
 28 for (int r = 0; r < L; ++r) {
284 24 const int32_t b0 = vals[2 * static_cast<size_t>(r)], b1 = vals[2 * static_cast<size_t>(r) + 1];
285
6/8
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 24 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 22 times.
✓ Branch 5 taken 2 times.
✓ Branch 6 taken 20 times.
✓ Branch 7 taken 2 times.
 24 if (b0 < 0 || b0 > r || b1 <= r || b1 > L) {
286
10/20
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 4 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 4 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 4 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 4 times.
✗ Branch 11 not taken.
✓ Branch 12 taken 4 times.
✗ Branch 13 not taken.
✓ Branch 14 taken 4 times.
✗ Branch 15 not taken.
✓ Branch 16 taken 4 times.
✗ Branch 17 not taken.
✓ Branch 18 taken 4 times.
✗ Branch 19 not taken.
 12 setError(std::string(lb) + ": seqBounds row " + std::to_string(r) + " is [" + std::to_string(b0) + ", " +
287
5/10
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 4 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 4 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 4 times.
✗ Branch 9 not taken.
 8 std::to_string(b1) + "), which must lie in [0, " + std::to_string(L) + ") and contain the row");
288 4 return nullptr;
289 }
290 20 }
291
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 return &int32Operand(*sb, vals, scratch);
292 8 }
293
294 } // namespace
295
296 extern "C" {
297
298 6 void bro_tensor_flashAttentionPackedQkvForward(void* QKV, void* seqBounds, int32_t numHeads, int32_t window,
299 void* O) {
300 6 constexpr const char* lb = "flashAttentionPackedQkvForward";
301
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (!need(lb, {QKV, seqBounds, O})) return;
302 BROTENSOR_API_TRY
303 6 brotensor::Tensor scratch;
304
3/6
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
 6 const brotensor::Tensor* sb = packedBounds(lb, toTensor(QKV), numHeads, toTensor(seqBounds), scratch);
305
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 4 times.
 6 if (!sb) return;
306
3/6
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
 2 brotensor::flash_attention_packed_qkv_forward(*toTensor(QKV), *sb, numHeads, window, *toTensor(O));
307
2/15
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 4 times.
✓ Branch 12 taken 2 times.
✗ Branch 13 not taken.
✗ Branch 14 not taken.
 6 BROTENSOR_API_CATCH(lb)
308 6 }
309
310 4 void bro_tensor_flashAttentionPackedQkvBackward(void* QKV, void* dO, void* seqBounds, int32_t numHeads,
311 int32_t window, void* dQKV) {
312 4 constexpr const char* lb = "flashAttentionPackedQkvBackward";
313
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (!need(lb, {QKV, dO, seqBounds, dQKV})) return;
314 4 auto* qkv = toTensor(QKV);
315 4 auto* dout = toTensor(dO);
316
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!needSameDtype(lb, "dO", dout, qkv)) return;
317
3/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 2 times.
 4 if (dout->rows != qkv->rows || static_cast<int64_t>(dout->cols) * 3 != qkv->cols) {
318
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 setError(std::string(lb) + ": dO must be (L, numHeads*headDim), a third of QKV's width");
319 2 return;
320 }
321 BROTENSOR_API_TRY
322 2 brotensor::Tensor scratch;
323
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 const brotensor::Tensor* sb = packedBounds(lb, qkv, numHeads, toTensor(seqBounds), scratch);
324
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!sb) return;
325
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 brotensor::flash_attention_packed_qkv_backward(*qkv, *dout, *sb, numHeads, window, *toTensor(dQKV));
326
1/15
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✓ Branch 10 taken 2 times.
✗ Branch 11 not taken.
✗ Branch 12 not taken.
✗ Branch 13 not taken.
✗ Branch 14 not taken.
 2 BROTENSOR_API_CATCH(lb)
327 4 }
328
329 // ---- Transformer-XL relative-position bias -------------------------------------
330
331 4 void bro_tensor_relPosBiasXlForward(void* Qv, void* Pk, int32_t numHeads, int32_t headDim, void* Bias) {
332 4 constexpr const char* lb = "relPosBiasXlForward";
333
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (!need(lb, {Qv, Pk, Bias})) return;
334 4 auto* q = toTensor(Qv);
335 4 auto* pk = toTensor(Pk);
336 4 constexpr auto F = brotensor::Dtype::FP32;
337
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
 4 if (!needDtype(lb, "Qv", q, F) || !needDtype(lb, "Pk", pk, F)) return;
338
3/6
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 4 times.
 4 if (numHeads < 1 || headDim < 1 || q->cols != static_cast<int64_t>(numHeads) * headDim) {
339 ✗ setError(std::string(lb) + ": Qv must be (T, numHeads*headDim)");
340 ✗ return;
341 }
342 4 const int64_t T = q->rows;
343
1/6
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 4 if (T < 1) { setError(std::string(lb) + ": Qv is empty"); return; }
344
3/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 2 times.
 4 if (pk->cols != q->cols || pk->rows < 2 * T - 1) {
345
6/12
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 2 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 2 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 2 times.
✗ Branch 11 not taken.
 6 setError(std::string(lb) + ": Pk must be (2T-1, numHeads*headDim) = (" + std::to_string(2 * T - 1) + ", " +
346
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 4 std::to_string(q->cols) + ")");
347 2 return;
348 }
349
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!needExtent(lb, "numHeads*T*T", prod({numHeads, T, T}))) return;
350 BROTENSOR_API_TRY
351
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 brotensor::rel_pos_bias_xl_forward(*q, *pk, numHeads, headDim, *toTensor(Bias));
352
0/12
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
 2 BROTENSOR_API_CATCH(lb)
353 4 }
354
355 // ---- RoPE: per-head tables, packed QKV in place --------------------------------
356
357 4 void bro_tensor_ropeApplyPerhead(void* X, void* cosTbl, void* sinTbl, int32_t headDim, int32_t numHeads, void* Y) {
358 4 constexpr const char* lb = "ropeApplyPerhead";
359
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (!need(lb, {X, cosTbl, sinTbl, Y})) return;
360 4 auto* x = toTensor(X);
361
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!needFloat(lb, "X", x)) return;
362
4/8
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 4 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 4 times.
 4 if (headDim < 2 || (headDim & 1) || numHeads < 1 || x->cols != static_cast<int64_t>(numHeads) * headDim) {
363 ✗ setError(std::string(lb) + ": X must be (L, numHeads*headDim) with headDim even");
364 ✗ return;
365 }
366 4 const int64_t n = prod({x->rows, numHeads, headDim / 2});
367
3/4
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 4 if (!optF32Elems(lb, "cosTbl", toTensor(cosTbl), n) || !optF32Elems(lb, "sinTbl", toTensor(sinTbl), n)) return;
368 BROTENSOR_API_TRY
369
4/8
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 2 times.
✗ Branch 7 not taken.
 2 brotensor::rope_apply_perhead(*x, *toTensor(cosTbl), *toTensor(sinTbl), headDim, numHeads, *toTensor(Y));
370
0/12
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
 2 BROTENSOR_API_CATCH(lb)
371 4 }
372
373 6 void bro_tensor_ropeQkvPackedInplace(void* QKV, void* cosTbl, void* sinTbl, void* pos, int32_t numHeads,
374 int32_t headDim) {
375 6 constexpr const char* lb = "ropeQkvPackedInplace";
376
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (!need(lb, {QKV, cosTbl, sinTbl, pos})) return;
377 6 auto* qkv = toTensor(QKV);
378 6 auto* ct = toTensor(cosTbl);
379 6 auto* st = toTensor(sinTbl);
380 6 auto* ps = toTensor(pos);
381
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!needFloat(lb, "QKV", qkv)) return;
382
4/8
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 6 times.
 6 if (headDim < 2 || (headDim & 1) || numHeads < 1 ||
383 6 qkv->cols != 3 * static_cast<int64_t>(numHeads) * headDim) {
384 ✗ setError(std::string(lb) + ": QKV must be (L, 3*numHeads*headDim) with headDim even");
385 ✗ return;
386 }
387 6 constexpr auto F = brotensor::Dtype::FP32;
388
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 6 times.
 6 if (!needDtype(lb, "cosTbl", ct, F) || !needDtype(lb, "sinTbl", st, F)) return;
389
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 6 times.
 6 if (ct->cols != headDim / 2 || st->cols != headDim / 2) {
390 ✗ setError(std::string(lb) + ": cosTbl / sinTbl must be (P, headDim/2)");
391 ✗ return;
392 }
393
3/4
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
 6 if (ps->rows != qkv->rows || ps->cols != 1) {
394
5/10
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 2 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 2 times.
✗ Branch 9 not taken.
 2 setError(std::string(lb) + ": pos must be (" + std::to_string(qkv->rows) + ", 1)");
395 2 return;
396 }
397 BROTENSOR_API_TRY
398 4 std::vector<int32_t> vals;
399
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 if (!hostInt32(lb, "pos", ps, qkv->rows, vals)) return;
400
4/6
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 2 times.
 4 if (!needIndicesIn(lb, "pos", vals, qkv->rows, 0, std::min(ct->rows, st->rows))) return;
401 2 brotensor::Tensor scratch;
402
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 brotensor::rope_qkv_packed_inplace(*qkv, *ct, *st, int32Operand(*ps, vals, scratch), numHeads, headDim);
403
2/15
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✗ Branch 12 not taken.
✗ Branch 13 not taken.
✗ Branch 14 not taken.
 4 BROTENSOR_API_CATCH(lb)
404 6 }
405
406 } // extern "C"
407
408 namespace brotensor::api {
409
410 3 bool registerTensorNatives_extra_nn(std::string* error) {
411 using namespace brotensor::api::reg;
412 3 return
413 6 fn("__bro_native.tensor.linearForwardBatchedEx", p(&bro_tensor_linearForwardBatchedEx), "void",
414
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 6 {kTensorCls, kDyn, kTensorCls, "i32", "i32", kDyn, kTensorCls}, error) &&
415 6 fn("__bro_native.tensor.linearForwardBatchedFp16Act", p(&bro_tensor_linearForwardBatchedFp16Act), "void",
416
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 6 {kTensorCls, kDyn, kTensorCls, "i32", kTensorCls}, error) &&
417 6 fn("__bro_native.tensor.matmulAbt", p(&bro_tensor_matmulAbt), "void",
418 3 {kTensorCls, kTensorCls, kTensorCls, "i32", "i32", "i32", "i32", "f64", "f64", "f64", kDyn, "i32"},
419
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 6 error) &&
420 6 fn("__bro_native.tensor.lstmForwardTrain", p(&bro_tensor_lstmForwardTrain), "void",
421 3 {kTensorCls, kTensorCls, kTensorCls, kDyn, kDyn, kDyn, kDyn, "i32", "i32",
422
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 6 kTensorCls, kTensorCls, kTensorCls, kDyn, kDyn}, error) &&
423 6 fn("__bro_native.tensor.lstmBackward", p(&bro_tensor_lstmBackward), "void",
424 3 {kTensorCls, kTensorCls, kTensorCls, kDyn, kDyn, kTensorCls, kTensorCls, kTensorCls, kTensorCls,
425
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 6 "i32", "i32", kTensorCls, kTensorCls, kTensorCls, kDyn, kDyn, kDyn, kDyn}, error) &&
426 6 fn("__bro_native.tensor.flashAttentionGqaForward", p(&bro_tensor_flashAttentionGqaForward), "void",
427
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 6 {kTensorCls, kTensorCls, kTensorCls, kDyn, "i32", "i32", "bool", kTensorCls}, error) &&
428 6 fn("__bro_native.tensor.flashAttentionPackedQkvForward", p(&bro_tensor_flashAttentionPackedQkvForward),
429
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 6 "void", {kTensorCls, kTensorCls, "i32", "i32", kTensorCls}, error) &&
430 6 fn("__bro_native.tensor.flashAttentionPackedQkvBackward", p(&bro_tensor_flashAttentionPackedQkvBackward),
431
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 6 "void", {kTensorCls, kTensorCls, kTensorCls, "i32", "i32", kTensorCls}, error) &&
432 6 fn("__bro_native.tensor.relPosBiasXlForward", p(&bro_tensor_relPosBiasXlForward), "void",
433
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 3 times.
✗ Branch 3 not taken.
 6 {kTensorCls, kTensorCls, "i32", "i32", kTensorCls}, error) &&
434 6 fn("__bro_native.tensor.ropeApplyPerhead", p(&bro_tensor_ropeApplyPerhead), "void",
435
2/4
✗ Branch 0 not taken.
✓ Branch 1 taken 3 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 3 times.
 6 {kTensorCls, kTensorCls, kTensorCls, "i32", "i32", kTensorCls}, error) &&
436 6 fn("__bro_native.tensor.ropeQkvPackedInplace", p(&bro_tensor_ropeQkvPackedInplace), "void",
437 3 {kTensorCls, kTensorCls, kTensorCls, kTensorCls, "i32", "i32"}, error);
438 }
439
440 } // namespace brotensor::api
441