GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 47.5% 248 / 0 / 522
Functions: 61.2% 30 / 0 / 49
Branches: 17.8% 228 / 0 / 1279
src/api/native_tensor_nn.cpp
Line Branch Exec Source
1 #include "native_tensor_decl.h"
2 #include "api_internal.h"
3 #include <vector>
4
5 using namespace brotensor::api;
6 using brotensor::Tensor;
7 using brotensor::Dtype;
8
9 // Operand checks come first in every body (api_internal.h "input-size
10 // contract"); accumulated outputs ("caller zeros") are checked like inputs.
11
12 namespace {
13
14 // X is (L, numHeads*headDim) with an even, positive headDim.
15 8 bool needRopeInput(const char* L, const Tensor* xt, int32_t headDim, int32_t numHeads) {
16
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (!needPositive(L, {headDim, numHeads})) return false;
17
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 6 times.
 8 if (headDim % 2 != 0) {
18
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 setError(std::string(L) + ": headDim must be even");
19 2 return false;
20 }
21
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 4 times.
 6 if (static_cast<int64_t>(xt->cols) != static_cast<int64_t>(numHeads) * headDim) {
22
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.
 4 setError(std::string(L) + ": X has " + std::to_string(xt->cols) + " columns, numHeads*headDim is " +
23
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 std::to_string(static_cast<int64_t>(numHeads) * headDim));
24 2 return false;
25 }
26 4 return true;
27 8 }
28
29 // cos/sin tables: (L, headDim/2) FP32.
30 4 bool needRopeTables(const char* L, const Tensor* xt, const Tensor* ct, const Tensor* st, int32_t headDim) {
31 4 const int64_t n = elems({xt->rows, headDim / 2});
32
3/4
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 4 if (!needElems(L, "cosTbl", ct, n) || !needElems(L, "sinTbl", st, n)) return false;
33
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (ct->dtype != Dtype::FP32 || st->dtype != Dtype::FP32) {
34 ✗ setError(std::string(L) + ": cosTbl and sinTbl must be FP32");
35 ✗ return false;
36 }
37 2 return true;
38 4 }
39
40 // Decode attention: Q (L_q, H*hd), caches (L_max, kvH*hd), kvH | H.
41 ✗ bool needDecodeShapes(const char* L, const Tensor* qt, const Tensor* kt, const Tensor* vt,
42 int32_t numHeads, int32_t kvHeads) {
43 ✗ if (!needPositive(L, {numHeads, kvHeads})) return false;
44 ✗ if (numHeads % kvHeads != 0 || qt->cols % numHeads != 0) {
45 ✗ setError(std::string(L) + ": numKvHeads must divide numHeads, and numHeads Q.cols");
46 ✗ return false;
47 }
48 ✗ const int64_t headDim = qt->cols / numHeads;
49 ✗ if (kt->cols != kvHeads * headDim || vt->cols != kvHeads * headDim || vt->rows != kt->rows) {
50 ✗ setError(std::string(L) + ": K_cache/V_cache must both be (L_max, numKvHeads*headDim) with headDim = Q.cols/numHeads");
51 ✗ return false;
52 }
53 ✗ return needSameDtype(L, "K_cache", kt, qt) && needSameDtype(L, "V_cache", vt, qt);
54 ✗ }
55
56 // Embedding indices: B values in [0, V), read back and range-checked (an
57 // index out of range is an out-of-bounds read of the table on the device),
58 // then handed to the op as an INT32 buffer on the table's device. The index
59 // tensor is INT32 storage or FP32 values.
60 20 bool embeddingIndices(const char* L, const Tensor* it, int32_t B, int32_t V, brotensor::Device dev,
61 Tensor& deviceIdx) {
62
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 18 times.
 20 if (!needElems(L, "indices", it, B)) return false;
63 18 std::vector<int32_t> idx(static_cast<size_t>(B));
64
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 auto host = it->to(brotensor::Device::CPU);
65
2/2
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 6 times.
 18 if (it->dtype == Dtype::INT32) {
66
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 std::memcpy(idx.data(), host.host_raw(), idx.size() * sizeof(int32_t));
67
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 18 } else if (it->dtype == Dtype::FP32) {
68
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 const float* f = host.host_f32();
69
2/2
✓ Branch 0 taken 10 times.
✓ Branch 1 taken 6 times.
 16 for (int32_t i = 0; i < B; ++i) idx[i] = static_cast<int32_t>(f[i]);
70 6 } else {
71 ✗ setError(std::string(L) + ": indices must be INT32 storage or FP32 values");
72 ✗ return false;
73 }
74
2/2
✓ Branch 0 taken 14 times.
✓ Branch 1 taken 24 times.
 38 for (int32_t i = 0; i < B; ++i) {
75
4/4
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 20 times.
 24 if (idx[i] < 0 || idx[i] >= V) {
76
9/18
✓ 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.
 8 setError(std::string(L) + ": index " + std::to_string(idx[i]) + " at " + std::to_string(i) +
77
3/6
✓ 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.
 4 " is outside the table's " + std::to_string(V) + " rows");
78 4 return false;
79 }
80 20 }
81
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 deviceIdx = Tensor::from_raw_bytes_on(dev, idx.data(), B, 1, Dtype::INT32, idx.size() * sizeof(int32_t));
82 14 return true;
83 20 }
84
85 } // namespace
86
87 extern "C" {
88
89 ✗ void bro_tensor_layernormBackward(void* dY, void* xhat, void* gamma, double rstd, void* dX, void* dGamma, void* dBeta) {
90 ✗ const char* L = "layernormBackward";
91 ✗ auto* dyt = toTensor(dY);
92 ✗ auto* xht = toTensor(xhat);
93 ✗ auto* gt = toTensor(gamma);
94 ✗ auto* dxt = toTensor(dX);
95 ✗ auto* dgt = toTensor(dGamma);
96 ✗ auto* dbt = toTensor(dBeta);
97 ✗ if (!need(L, {dyt, xht, gt, dxt, dgt, dbt})) return;
98 ✗ if (!needPair(L, "xhat", xht, dyt) || !needPair(L, "gamma", gt, dyt)) return;
99 ✗ if (!needElems(L, "dGamma", dgt, dyt->size()) || !needElems(L, "dBeta", dbt, dyt->size())) return;
100 ✗ if (!needSameDtype(L, "dGamma", dgt, dyt) || !needSameDtype(L, "dBeta", dbt, dyt)) return;
101 BROTENSOR_API_TRY
102 ✗ brotensor::layernorm_backward(*dyt, *xht, *gt, static_cast<float>(rstd), *dxt, *dgt, *dbt);
103 ✗ BROTENSOR_API_CATCH(L)
104 ✗ }
105
106 ✗ void bro_tensor_layernormForwardInferenceBatched(void* X_RD, void* gamma, void* beta, void* Y_RD, double eps) {
107 ✗ const char* L = "layernormForwardInferenceBatched";
108 ✗ auto* xt = toTensor(X_RD);
109 ✗ auto* gt = toTensor(gamma);
110 ✗ auto* bt = toTensor(beta);
111 ✗ auto* yt = toTensor(Y_RD);
112 ✗ if (!need(L, {xt, gt, bt, yt})) return;
113 ✗ if (!needElems(L, "gamma", gt, xt->cols) || !needElems(L, "beta", bt, xt->cols)) return;
114 ✗ if (!needSameDtype(L, "gamma", gt, xt) || !needSameDtype(L, "beta", bt, xt)) return;
115 BROTENSOR_API_TRY
116 ✗ brotensor::layernorm_forward_inference_batched(*xt, *gt, *bt, *yt, static_cast<float>(eps));
117 ✗ BROTENSOR_API_CATCH(L)
118 ✗ }
119
120 ✗ void bro_tensor_layernormForwardInferenceBatchedFp16(void* X_RD, void* gamma, void* beta, void* Y_RD, double eps) {
121 ✗ const char* L = "layernormForwardInferenceBatchedFp16";
122 ✗ auto* xt = toTensor(X_RD);
123 ✗ auto* gt = toTensor(gamma);
124 ✗ auto* bt = toTensor(beta);
125 ✗ auto* yt = toTensor(Y_RD);
126 ✗ if (!need(L, {xt, gt, bt, yt})) return;
127 ✗ if (!needElems(L, "gamma", gt, xt->cols) || !needElems(L, "beta", bt, xt->cols)) return;
128 ✗ if (!needSameDtype(L, "gamma", gt, xt) || !needSameDtype(L, "beta", bt, xt)) return;
129 BROTENSOR_API_TRY
130 ✗ brotensor::layernorm_forward_inference_batched_fp16(*xt, *gt, *bt, *yt, static_cast<float>(eps));
131 ✗ BROTENSOR_API_CATCH(L)
132 ✗ }
133
134 2 void bro_tensor_rmsNormForward(void* X, void* gamma, double eps, void* Y) {
135 2 const char* L = "rmsNormForward";
136 2 auto* xt = toTensor(X);
137 2 auto* gt = toTensor(gamma);
138 2 auto* yt = toTensor(Y);
139
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {xt, gt, yt})) return;
140 // gamma matches X or is FP32 against 16-bit activations.
141
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!needElems(L, "gamma", gt, xt->cols)) return;
142 ✗ if (gt->dtype != Dtype::FP32 && !needSameDtype(L, "gamma", gt, xt)) return;
143 BROTENSOR_API_TRY
144 ✗ brotensor::rms_norm_forward(*xt, *gt, static_cast<float>(eps), *yt);
145 ✗ BROTENSOR_API_CATCH(L)
146 2 }
147
148 ✗ void bro_tensor_rmsNormBackward(void* X, void* gamma, void* dY, double eps, void* dX, void* dGamma) {
149 ✗ const char* L = "rmsNormBackward";
150 ✗ auto* xt = toTensor(X);
151 ✗ auto* gt = toTensor(gamma);
152 ✗ auto* dyt = toTensor(dY);
153 ✗ auto* dxt = toTensor(dX);
154 ✗ auto* dgt = toTensor(dGamma);
155 ✗ if (!need(L, {xt, gt, dyt, dxt, dgt})) return;
156 ✗ if (!needElems(L, "gamma", gt, xt->cols) || !needSameDtype(L, "gamma", gt, xt)) return;
157 ✗ if (!needPair(L, "dY", dyt, xt)) return;
158 ✗ if (!needElems(L, "dGamma", dgt, xt->cols) || !needSameDtype(L, "dGamma", dgt, xt)) return;
159 BROTENSOR_API_TRY
160 ✗ brotensor::rms_norm_backward(*xt, *gt, *dyt, static_cast<float>(eps), *dxt, *dgt);
161 ✗ BROTENSOR_API_CATCH(L)
162 ✗ }
163
164 2 void bro_tensor_groupNormForward(void* X, void* gamma, void* beta, int32_t N, int32_t C, int32_t H, int32_t W, int32_t numGroups, double eps, void* Y) {
165 2 const char* L = "groupNormForward";
166 2 auto* xt = toTensor(X);
167 2 auto* gt = toTensor(gamma);
168 2 auto* bt = toTensor(beta);
169 2 auto* yt = toTensor(Y);
170
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {xt, gt, bt, yt})) return;
171
2/4
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!needPositive(L, {numGroups}) || !needNonNegative(L, {N, C, H, W})) return;
172
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 if (C % numGroups != 0) { setError("groupNormForward: numGroups must divide C"); return; }
173 ✗ if (!needElems(L, "X", xt, elems({N, C, H, W}))) return;
174 ✗ if (!needElems(L, "gamma", gt, C) || !needElems(L, "beta", bt, C)) return;
175 ✗ if (!needSameDtype(L, "gamma", gt, xt) || !needSameDtype(L, "beta", bt, xt)) return;
176 BROTENSOR_API_TRY
177 ✗ brotensor::group_norm_forward(*xt, *gt, *bt, N, C, H, W, numGroups, static_cast<float>(eps), *yt);
178 ✗ BROTENSOR_API_CATCH(L)
179 2 }
180
181 ✗ void bro_tensor_groupNormBackward(void* X, void* gamma, void* dY, int32_t N, int32_t C, int32_t H, int32_t W, int32_t numGroups, double eps, void* dX, void* dGamma, void* dBeta) {
182 ✗ const char* L = "groupNormBackward";
183 ✗ auto* xt = toTensor(X);
184 ✗ auto* gt = toTensor(gamma);
185 ✗ auto* dyt = toTensor(dY);
186 ✗ auto* dxt = toTensor(dX);
187 ✗ auto* dgt = toTensor(dGamma);
188 ✗ auto* dbt = toTensor(dBeta);
189 ✗ if (!need(L, {xt, gt, dyt, dxt, dgt, dbt})) return;
190 ✗ if (!needPositive(L, {numGroups}) || !needNonNegative(L, {N, C, H, W})) return;
191 ✗ if (C % numGroups != 0) { setError("groupNormBackward: numGroups must divide C"); return; }
192 ✗ const int64_t n = elems({N, C, H, W});
193 ✗ if (!needElems(L, "X", xt, n) || !needElems(L, "dY", dyt, n) || !needSameDtype(L, "dY", dyt, xt)) return;
194 ✗ for (const Tensor* t : {gt, dgt, dbt}) {
195 ✗ if (!needElems(L, "gamma/dGamma/dBeta", t, C) || !needSameDtype(L, "gamma/dGamma/dBeta", t, xt)) return;
196 }
197 BROTENSOR_API_TRY
198 ✗ brotensor::group_norm_backward(*xt, *gt, *dyt, N, C, H, W, numGroups, static_cast<float>(eps), *dxt, *dgt, *dbt);
199 ✗ BROTENSOR_API_CATCH(L)
200 ✗ }
201
202 4 void bro_tensor_ropeForward(void* X, int32_t headDim, int32_t numHeads, int32_t seqOffset, double thetaBase, void* Y) {
203 4 const char* L = "ropeForward";
204 4 auto* xt = toTensor(X);
205 4 auto* yt = toTensor(Y);
206
2/4
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 if (!need(L, {xt, yt}) || !needRopeInput(L, xt, headDim, numHeads)) return;
207 BROTENSOR_API_TRY
208 ✗ brotensor::rope_forward(*xt, headDim, numHeads, seqOffset, static_cast<float>(thetaBase), *yt);
209 ✗ BROTENSOR_API_CATCH(L)
210 4 }
211
212 ✗ void bro_tensor_ropeBackward(void* dY, int32_t headDim, int32_t numHeads, int32_t seqOffset, double thetaBase, void* dX) {
213 ✗ const char* L = "ropeBackward";
214 ✗ auto* dyt = toTensor(dY);
215 ✗ auto* dxt = toTensor(dX);
216 ✗ if (!need(L, {dyt, dxt}) || !needRopeInput(L, dyt, headDim, numHeads)) return;
217 BROTENSOR_API_TRY
218 ✗ brotensor::rope_backward(*dyt, headDim, numHeads, seqOffset, static_cast<float>(thetaBase), *dxt);
219 ✗ BROTENSOR_API_CATCH(L)
220 ✗ }
221
222 4 void bro_tensor_ropeApply(void* X, void* cosTbl, void* sinTbl, int32_t headDim, int32_t numHeads, void* Y) {
223 4 const char* L = "ropeApply";
224 4 auto* xt = toTensor(X);
225 4 auto* ct = toTensor(cosTbl);
226 4 auto* st = toTensor(sinTbl);
227 4 auto* yt = toTensor(Y);
228
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (!need(L, {xt, ct, st, yt})) return;
229
3/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 2 times.
 4 if (!needRopeInput(L, xt, headDim, numHeads) || !needRopeTables(L, xt, ct, st, headDim)) return;
230 BROTENSOR_API_TRY
231
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 brotensor::rope_apply(*xt, *ct, *st, headDim, numHeads, *yt);
232
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(L)
233 4 }
234
235 ✗ void bro_tensor_ropeApplyBackward(void* dY, void* cosTbl, void* sinTbl, int32_t headDim, int32_t numHeads, void* dX) {
236 ✗ const char* L = "ropeApplyBackward";
237 ✗ auto* dyt = toTensor(dY);
238 ✗ auto* ct = toTensor(cosTbl);
239 ✗ auto* st = toTensor(sinTbl);
240 ✗ auto* dxt = toTensor(dX);
241 ✗ if (!need(L, {dyt, ct, st, dxt})) return;
242 ✗ if (!needRopeInput(L, dyt, headDim, numHeads) || !needRopeTables(L, dyt, ct, st, headDim)) return;
243 BROTENSOR_API_TRY
244 ✗ brotensor::rope_apply_backward(*dyt, *ct, *st, headDim, numHeads, *dxt);
245 ✗ BROTENSOR_API_CATCH(L)
246 ✗ }
247
248 ✗ void bro_tensor_modulate(void* X, void* scale, void* shift, void* Y) {
249 ✗ const char* L = "modulate";
250 ✗ auto* xt = toTensor(X);
251 ✗ auto* st = toTensor(scale);
252 ✗ auto* sht = toTensor(shift);
253 ✗ auto* yt = toTensor(Y);
254 ✗ if (!need(L, {xt, st, sht, yt})) return;
255 ✗ if (!needElems(L, "scale", st, xt->cols) || !needElems(L, "shift", sht, xt->cols)) return;
256 ✗ if (!needSameDtype(L, "scale", st, xt) || !needSameDtype(L, "shift", sht, xt)) return;
257 BROTENSOR_API_TRY
258 ✗ brotensor::modulate(*xt, *st, *sht, *yt);
259 ✗ BROTENSOR_API_CATCH(L)
260 ✗ }
261
262 ✗ void bro_tensor_broadcastMul(void* X, void* v, void* Y) {
263 ✗ const char* L = "broadcastMul";
264 ✗ auto* xt = toTensor(X);
265 ✗ auto* vt = toTensor(v);
266 ✗ auto* yt = toTensor(Y);
267 ✗ if (!need(L, {xt, vt, yt})) return;
268 ✗ if (!needElems(L, "v", vt, xt->cols) || !needSameDtype(L, "v", vt, xt)) return;
269 BROTENSOR_API_TRY
270 ✗ brotensor::broadcast_mul(*xt, *vt, *yt);
271 ✗ BROTENSOR_API_CATCH(L)
272 ✗ }
273
274 ✗ void bro_tensor_attentionTokenMoments(void* Attn, int32_t h_lat, int32_t w_lat, void* mass, void* centroid) {
275 ✗ const char* L = "attentionTokenMoments";
276 ✗ auto* at = toTensor(Attn);
277 ✗ auto* mt = toTensor(mass);
278 ✗ auto* ct = toTensor(centroid);
279 ✗ if (!need(L, {at, mt, ct}) || !needPositive(L, {h_lat, w_lat})) return;
280 ✗ if (static_cast<int64_t>(at->rows) != static_cast<int64_t>(h_lat) * w_lat) {
281 ✗ setError("attentionTokenMoments: Attn.rows must equal h_lat*w_lat");
282 ✗ return;
283 }
284 BROTENSOR_API_TRY
285 ✗ brotensor::attention_token_moments(*at, h_lat, w_lat, *mt, *ct);
286 ✗ BROTENSOR_API_CATCH(L)
287 ✗ }
288
289 // mask[k] = x[offset + k*stride] > 0.5: the last read is offset+(K-1)*stride.
290 ✗ void bro_tensor_buildSlotMask(void* x, int32_t offset, int32_t K, int32_t stride, void* mask) {
291 ✗ const char* L = "buildSlotMask";
292 ✗ auto* xt = toTensor(x);
293 ✗ auto* mt = toTensor(mask);
294 ✗ if (!need(L, {xt, mt}) || !needNonNegative(L, {offset, K, stride})) return;
295 ✗ if (K > 0 && !needElems(L, "x", xt, static_cast<int64_t>(offset) + static_cast<int64_t>(K - 1) * stride + 1)) return;
296 ✗ if (xt->dtype != Dtype::FP32) { setError("buildSlotMask: x must be FP32"); return; }
297 BROTENSOR_API_TRY
298 ✗ brotensor::build_slot_mask(*xt, offset, K, stride, *mt);
299 ✗ BROTENSOR_API_CATCH(L)
300 ✗ }
301
302 ✗ void bro_tensor_buildCausalMaskRow(int32_t L_, int32_t q, void* mask) {
303 ✗ const char* L = "buildCausalMaskRow";
304 ✗ auto* mt = toTensor(mask);
305 ✗ if (!need(L, {mt}) || !needNonNegative(L, {L_})) return;
306 BROTENSOR_API_TRY
307 ✗ brotensor::build_causal_mask_row(L_, q, *mt);
308 ✗ BROTENSOR_API_CATCH(L)
309 ✗ }
310
311 ✗ void bro_tensor_flashAttentionDecode(void* Q, void* K_cache, void* V_cache, int32_t validLen, int32_t numHeads, void* O, bool numKvHeads_given, int32_t numKvHeads, double attnSoftcap, int32_t window) {
312 ✗ const char* L = "flashAttentionDecode";
313 ✗ auto* qt = toTensor(Q);
314 ✗ auto* kt = toTensor(K_cache);
315 ✗ auto* vt = toTensor(V_cache);
316 ✗ auto* ot = toTensor(O);
317 ✗ if (!need(L, {qt, kt, vt, ot})) return;
318 ✗ const int kvHeads = (numKvHeads_given && numKvHeads > 0) ? numKvHeads : numHeads;
319 ✗ if (!needDecodeShapes(L, qt, kt, vt, numHeads, kvHeads)) return;
320 ✗ if (validLen < qt->rows || validLen > kt->rows) {
321 ✗ setError("flashAttentionDecode: validLen must lie in [Q.rows, K_cache.rows]");
322 ✗ return;
323 }
324 BROTENSOR_API_TRY
325 ✗ brotensor::flash_attention_decode(*qt, *kt, *vt, validLen, numHeads, kvHeads, *ot, static_cast<float>(attnSoftcap), window);
326 ✗ BROTENSOR_API_CATCH(L)
327 ✗ }
328
329 ✗ void bro_tensor_flashAttentionDecodeMasked(void* Q, void* K_cache, void* V_cache, void* dMask, int32_t numHeads, void* O, bool numKvHeads_given, int32_t numKvHeads, double attnSoftcap, int32_t window) {
330 ✗ const char* L = "flashAttentionDecodeMasked";
331 ✗ auto* qt = toTensor(Q);
332 ✗ auto* kt = toTensor(K_cache);
333 ✗ auto* vt = toTensor(V_cache);
334 ✗ auto* mt = toTensor(dMask);
335 ✗ auto* ot = toTensor(O);
336 // The op reads a length-L_max mask and must not be handed null.
337 ✗ if (!need(L, {qt, kt, vt, mt, ot})) return;
338 ✗ const int kvHeads = (numKvHeads_given && numKvHeads > 0) ? numKvHeads : numHeads;
339 ✗ if (!needDecodeShapes(L, qt, kt, vt, numHeads, kvHeads)) return;
340 ✗ if (qt->rows != 1) { setError("flashAttentionDecodeMasked: Q must be a single row"); return; }
341 ✗ if (!needMask(L, mt, kt->rows)) return;
342 BROTENSOR_API_TRY
343 ✗ brotensor::flash_attention_decode_masked(*qt, *kt, *vt, static_cast<const float*>(mt->data), numHeads, kvHeads,
344 ✗ *ot, static_cast<float>(attnSoftcap), window);
345 ✗ BROTENSOR_API_CATCH(L)
346 ✗ }
347
348 ✗ void bro_tensor_kvCacheAppend(void* K_new, void* V_new, int32_t curLen, void* K_cache, void* V_cache) {
349 ✗ const char* L = "kvCacheAppend";
350 ✗ auto* knt = toTensor(K_new);
351 ✗ auto* vnt = toTensor(V_new);
352 ✗ auto* kct = toTensor(K_cache);
353 ✗ auto* vct = toTensor(V_cache);
354 ✗ if (!need(L, {knt, vnt, kct, vct}) || !needNonNegative(L, {curLen})) return;
355 ✗ if (vnt->rows != knt->rows || vnt->cols != knt->cols || kct->cols != knt->cols || vct->cols != knt->cols) {
356 ✗ setError("kvCacheAppend: K_new/V_new/K_cache/V_cache must share a column count, K_new and V_new a shape");
357 ✗ return;
358 }
359 ✗ const int64_t end = static_cast<int64_t>(curLen) + knt->rows;
360 ✗ if (end > kct->rows || end > vct->rows) {
361 ✗ setError("kvCacheAppend: curLen + K_new.rows exceeds the cache's rows");
362 ✗ return;
363 }
364 ✗ for (const Tensor* t : {vnt, kct, vct}) {
365 ✗ if (!needSameDtype(L, "every operand", t, knt)) return;
366 }
367 BROTENSOR_API_TRY
368 ✗ brotensor::kv_cache_append(*knt, *vnt, curLen, *kct, *vct);
369 ✗ BROTENSOR_API_CATCH(L)
370 ✗ }
371
372 // The indices are range-checked on the host and reach the op as a device
373 // INT32 buffer. (The FP32-index path used to hand a host pointer to the
374 // device kernel.)
375 20 void bro_tensor_embeddingLookupForward(void* table, void* idxAsInt32, int32_t B, void* out) {
376 20 const char* L = "embeddingLookupForward";
377 20 auto* tt = toTensor(table);
378 20 auto* it = toTensor(idxAsInt32);
379 20 auto* ot = toTensor(out);
380
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 20 times.
 20 if (!need(L, {tt, it, ot})) return;
381
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 20 if (B <= 0) { setError("embeddingLookupForward: B must be positive"); return; }
382 BROTENSOR_API_TRY
383 20 Tensor idx;
384
3/4
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 14 times.
✓ Branch 3 taken 6 times.
 20 if (!embeddingIndices(L, it, B, tt->rows, tt->device, idx)) return;
385
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 brotensor::embedding_lookup_forward(*tt, static_cast<const int32_t*>(idx.data), B, *ot);
386
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 6 times.
✓ Branch 12 taken 14 times.
✗ Branch 13 not taken.
✗ Branch 14 not taken.
 20 BROTENSOR_API_CATCH(L)
387 20 }
388
389 ✗ void bro_tensor_embeddingLookupBackward(void* dOut, void* idxAsInt32, int32_t B, void* dTable) {
390 ✗ const char* L = "embeddingLookupBackward";
391 ✗ auto* dot = toTensor(dOut);
392 ✗ auto* it = toTensor(idxAsInt32);
393 ✗ auto* dtt = toTensor(dTable);
394 ✗ if (!need(L, {dot, it, dtt})) return;
395 ✗ if (B <= 0) { setError("embeddingLookupBackward: B must be positive"); return; }
396 ✗ if (!needElems(L, "dOut", dot, elems({B, dtt->cols})) || !needSameDtype(L, "dOut", dot, dtt)) return;
397 BROTENSOR_API_TRY
398 ✗ Tensor idx;
399 ✗ if (!embeddingIndices(L, it, B, dtt->rows, dtt->device, idx)) return;
400 ✗ brotensor::embedding_lookup_backward(*dot, static_cast<const int32_t*>(idx.data), B, *dtt);
401 ✗ BROTENSOR_API_CATCH(L)
402 ✗ }
403
404 2 void bro_tensor_sgdStep(void* param, void* grad, void* velocity, double lr, double momentum) {
405 2 const char* L = "sgdStep";
406 2 auto* pt = toTensor(param);
407 2 auto* gt = toTensor(grad);
408 2 auto* vt = toTensor(velocity);
409
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {pt, gt, vt})) return;
410
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 if (!needPair(L, "grad", gt, pt) || !needPair(L, "velocity", vt, pt)) return;
411 // The kernels read every operand as float.
412 ✗ if (pt->dtype != Dtype::FP32) { setError("sgdStep: tensors must be FP32"); return; }
413 BROTENSOR_API_TRY
414 ✗ brotensor::sgd_step(*pt, *gt, *vt, static_cast<float>(lr), static_cast<float>(momentum));
415 ✗ BROTENSOR_API_CATCH(L)
416 2 }
417
418 2 void bro_tensor_adamStep(void* param, void* grad, void* m, void* v, double lr, double beta1, double beta2, double eps, int32_t step) {
419 2 const char* L = "adamStep";
420 2 auto* pt = toTensor(param);
421 2 auto* gt = toTensor(grad);
422 2 auto* mt = toTensor(m);
423 2 auto* vt = toTensor(v);
424
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {pt, gt, mt, vt})) return;
425
2/6
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 2 if (!needPair(L, "grad", gt, pt) || !needPair(L, "m", mt, pt) || !needPair(L, "v", vt, pt)) return;
426 ✗ if (step < 1) { setError("adamStep: step is the 1-based bias-correction counter"); return; }
427 ✗ if (pt->dtype != Dtype::FP32) { setError("adamStep: tensors must be FP32"); return; }
428 BROTENSOR_API_TRY
429 ✗ brotensor::adam_step(*pt, *gt, *mt, *vt, static_cast<float>(lr), static_cast<float>(beta1), static_cast<float>(beta2), static_cast<float>(eps), step);
430 ✗ BROTENSOR_API_CATCH(L)
431 2 }
432
433 // ---- restored from the QuickJS binding (tensor_bindings.cpp) ---------------
434
435 110 void bro_tensor_cast(void* src, void* dst, const char* outDtype) {
436
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 110 times.
 110 if (!need("cast", {src, dst})) return;
437 BROTENSOR_API_TRY
438
4/8
✓ Branch 0 taken 110 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 110 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 110 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 110 times.
✗ Branch 7 not taken.
 110 brotensor::cast(*toTensor(src), *toTensor(dst), parseDtype(outDtype));
439
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.
 110 BROTENSOR_API_CATCH("cast")
440 110 }
441
442 // softmaxForward(logits, probs, mask|null): the old third argument.
443 6 void bro_tensor_softmaxForwardMasked(void* logits, void* probs, uint64_t mask_bits) {
444
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (!need("softmaxForward", {logits, probs})) return;
445
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 if (!needMask("softmaxForward", tensorFromValue(mask_bits), toTensor(logits)->size())) return;
446 BROTENSOR_API_TRY
447
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 taken 6 times.
✗ Branch 7 not taken.
 6 brotensor::softmax_forward(*toTensor(logits), *toTensor(probs), maskPtr(mask_bits));
448
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.
 6 BROTENSOR_API_CATCH("softmaxForward")
449 6 }
450
451 // layernormForward(x, gamma, beta, y, xhat, eps) -> {mean, rstd}: the two
452 // scalar caches come back as a length-2 f64[] the wrapper unpacks.
453 2 void bro_tensor_layernormForward(void* x, void* gamma, void* beta, void* y, void* xhat, double eps, bronze_native_buffer* out) {
454 static thread_local double stats[2];
455
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!out) return;
456 2 out->data = nullptr;
457 2 out->length = 0;
458 2 out->release = nullptr;
459 2 out->ctx = nullptr;
460
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need("layernormForward", {x, gamma, beta, y, xhat})) return;
461
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!needPair("layernormForward", "gamma", toTensor(gamma), toTensor(x)) ||
462 2 !needPair("layernormForward", "beta", toTensor(beta), toTensor(x))) return;
463 BROTENSOR_API_TRY
464 2 float mean = 0.f, rstd = 0.f;
465
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.
 2 brotensor::layernorm_forward(*toTensor(x), *toTensor(gamma), *toTensor(beta), *toTensor(y), *toTensor(xhat),
466 2 mean, rstd, static_cast<float>(eps));
467 2 stats[0] = mean;
468 2 stats[1] = rstd;
469 2 out->data = stats;
470 2 out->length = 2;
471
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("layernormForward")
472 2 }
473
474 4 void bro_tensor_maskedMeanPoolForward(void* X, uint64_t mask_bits, void* y) {
475
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (!need("maskedMeanPoolForward", {X, y})) return;
476
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!needMask("maskedMeanPoolForward", tensorFromValue(mask_bits), toTensor(X)->rows)) return;
477 BROTENSOR_API_TRY
478
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 taken 4 times.
✗ Branch 7 not taken.
 4 brotensor::masked_mean_pool_forward(*toTensor(X), maskPtr(mask_bits), *toTensor(y));
479
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("maskedMeanPoolForward")
480 4 }
481
482 2 void bro_tensor_maskedMeanPoolBackward(void* dY, uint64_t mask_bits, int32_t K, void* dX) {
483
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need("maskedMeanPoolBackward", {dY, dX})) return;
484
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!needNonNegative("maskedMeanPoolBackward", {K})) return;
485
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needMask("maskedMeanPoolBackward", tensorFromValue(mask_bits), K)) return;
486 BROTENSOR_API_TRY
487
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::masked_mean_pool_backward(*toTensor(dY), maskPtr(mask_bits), K, *toTensor(dX));
488
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("maskedMeanPoolBackward")
489 2 }
490
491 2 double bro_tensor_softmaxXentFused(void* logits, void* target, uint64_t mask_bits, void* probs, void* dLogits) {
492
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need("softmaxXentFused", {logits, target, probs, dLogits})) return 0.0;
493
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needPair("softmaxXentFused", "target", toTensor(target), toTensor(logits))) return 0.0;
494
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needMask("softmaxXentFused", tensorFromValue(mask_bits), toTensor(logits)->size())) return 0.0;
495 BROTENSOR_API_TRY
496
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 return brotensor::softmax_xent_fused(*toTensor(logits), *toTensor(target), maskPtr(mask_bits),
497
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(probs), *toTensor(dLogits));
498 ✗ BROTENSOR_API_CATCH("softmaxXentFused")
499 ✗ return 0.0;
500 2 }
501
502 // The head_offsets GpuTensor is a device INT32 buffer (cumulative, n_heads+1),
503 // the old binding's convention.
504 ✗ void bro_tensor_softmaxXentFusedBatched(void* logits_BL, void* target_BL, uint64_t mask_bits, void* headOffsets, int32_t nHeads, void* probs_BL, void* dLogits_BL, void* lossPerSample) {
505 ✗ const char* L = "softmaxXentFusedBatched";
506 ✗ if (!need(L, {logits_BL, target_BL, headOffsets, probs_BL, dLogits_BL, lossPerSample})) return;
507 ✗ const Tensor* lt = toTensor(logits_BL);
508 ✗ if (!needPair(L, "target_BL", toTensor(target_BL), lt)) return;
509 ✗ if (!needMask(L, tensorFromValue(mask_bits), lt->size())) return;
510 ✗ if (!needPositive(L, {nHeads})) return;
511 // The offsets are read on the device as slice bounds into every row: read
512 // them back and check them (n_heads+1 cumulative ints within a row).
513 ✗ const Tensor* ht = toTensor(headOffsets);
514 ✗ if (ht->dtype != Dtype::INT32 || !needElems(L, "headOffsets", ht, static_cast<int64_t>(nHeads) + 1)) {
515 ✗ if (ht->dtype != Dtype::INT32) setError("softmaxXentFusedBatched: headOffsets must be INT32 storage");
516 ✗ return;
517 }
518 BROTENSOR_API_TRY
519 ✗ auto host = ht->to(brotensor::Device::CPU);
520 ✗ const auto* off = static_cast<const int32_t*>(host.host_raw());
521 ✗ for (int32_t h = 0; h <= nHeads; ++h) {
522 ✗ if (off[h] < 0 || off[h] > lt->cols || (h > 0 && off[h] < off[h - 1])) {
523 ✗ setError("softmaxXentFusedBatched: headOffsets must be non-decreasing within [0, logits.cols]");
524 ✗ return;
525 }
526 ✗ }
527 ✗ brotensor::softmax_xent_fused_batched(*toTensor(logits_BL), *toTensor(target_BL), maskPtr(mask_bits),
528 ✗ static_cast<const int*>(ht->data), nHeads,
529 ✗ *toTensor(probs_BL), *toTensor(dLogits_BL), *toTensor(lossPerSample));
530 ✗ BROTENSOR_API_CATCH(L)
531 ✗ }
532
533 } // extern "C"
534
535 namespace {
536
537 6 std::vector<const brotensor::Tensor*> asConst(const std::vector<brotensor::Tensor*>& v) {
538 6 return std::vector<const brotensor::Tensor*>(v.begin(), v.end());
539 }
540
541 4 std::vector<int> asInts(const int32_t* p, uint32_t n) {
542
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 return p ? std::vector<int>(p, p + n) : std::vector<int>();
543 }
544
545 // The concat/split kernels copy every part with the FIRST part's element size.
546 12 bool needOneDtype(const char* L, const std::vector<brotensor::Tensor*>& parts, const Tensor* ref) {
547
2/2
✓ Branch 0 taken 24 times.
✓ Branch 1 taken 10 times.
 34 for (const Tensor* p : parts) {
548
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 22 times.
 24 if (!needSameDtype(L, "every part", p, ref)) return false;
549 }
550 10 return true;
551 12 }
552
553 // The RNG fills write Y as float.
554 20 bool needFp32Out(const char* L, const Tensor* y) {
555
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 20 times.
 20 if (y->dtype != Dtype::FP32) {
556 ✗ setError(std::string(L) + ": Y must be FP32");
557 ✗ return false;
558 }
559 20 return true;
560 20 }
561
562 } // namespace
563
564 extern "C" {
565
566 4 void bro_tensor_concatRows(uint64_t parts_bits, void* out) {
567 4 std::vector<brotensor::Tensor*> parts;
568
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 (!need("concatRows", {out}) || !readTensorArray(parts_bits, parts, "concatRows")) return;
569
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 (!parts.empty() && !needOneDtype("concatRows", parts, parts[0])) return;
570 BROTENSOR_API_TRY
571
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::concat_rows(asConst(parts), *toTensor(out));
572
0/14
✗ 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.
✗ Branch 12 not taken.
✗ Branch 13 not taken.
 2 BROTENSOR_API_CATCH("concatRows")
573 4 }
574
575 // split_rows copies consecutive runs of `in` into each part at the part's
576 // size: the parts must fit in `in` and share its dtype.
577 4 void bro_tensor_splitRows(void* in, uint64_t parts_bits) {
578 4 std::vector<brotensor::Tensor*> parts;
579
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 taken 4 times.
✗ Branch 7 not taken.
 4 if (!need("splitRows", {in}) || !readTensorArray(parts_bits, parts, "splitRows")) return;
580
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 const Tensor* it = toTensor(in);
581
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 if (!needOneDtype("splitRows", parts, it)) return;
582 4 int64_t total = 0;
583
3/4
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 8 times.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 12 for (const Tensor* p : parts) total += p->size();
584
3/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 2 times.
 4 if (total > static_cast<int64_t>(it->size())) {
585
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("splitRows: the parts hold " + std::to_string(total) + " elements, in has " + std::to_string(it->size()));
586 2 return;
587 }
588 BROTENSOR_API_TRY
589
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 brotensor::split_rows(*toTensor(in), parts);
590
0/14
✗ 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.
✗ Branch 12 not taken.
✗ Branch 13 not taken.
 2 BROTENSOR_API_CATCH("splitRows")
591 4 }
592
593 // Every part is (B, d_i) with one B.
594 4 void bro_tensor_concatBatchedRows(uint64_t parts_bits, void* out) {
595 4 std::vector<brotensor::Tensor*> parts;
596
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 (!need("concatBatchedRows", {out}) || !readTensorArray(parts_bits, parts, "concatBatchedRows")) return;
597
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!parts.empty()) {
598
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 if (!needOneDtype("concatBatchedRows", parts, parts[0])) return;
599
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 2 times.
 10 for (const Tensor* p : parts) {
600
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 6 times.
 8 if (p->rows != parts[0]->rows) {
601
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 setError("concatBatchedRows: every part must have the same row count");
602 2 return;
603 }
604 }
605 2 }
606 BROTENSOR_API_TRY
607
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::concat_batched_rows(asConst(parts), *toTensor(out));
608
0/14
✗ 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.
✗ Branch 12 not taken.
✗ Branch 13 not taken.
 2 BROTENSOR_API_CATCH("concatBatchedRows")
609 4 }
610
611 // C_per_part is an i32[] view INTO THE MOVING HEAP, and readTensorArray walks
612 // the parts array with allocating getProperty/getElement calls: copy the ints
613 // out first, or the second read sees a buffer the collector has moved (the
614 // gcstress run caught exactly that).
615 2 void bro_tensor_concatNchwChannels(uint64_t parts_bits, int32_t N, int32_t H, int32_t W, const int32_t* C_per_part, uint32_t C_len, void* out) {
616 2 const std::vector<int> channels = asInts(C_per_part, C_len);
617 2 std::vector<brotensor::Tensor*> parts;
618
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 not taken.
✓ Branch 7 taken 2 times.
 2 if (!need("concatNchwChannels", {out}) || !readTensorArray(parts_bits, parts, "concatNchwChannels")) return;
619 BROTENSOR_API_TRY
620
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::concat_nchw_channels(asConst(parts), N, H, W, channels, *toTensor(out));
621
0/14
✗ 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.
✗ Branch 12 not taken.
✗ Branch 13 not taken.
 2 BROTENSOR_API_CATCH("concatNchwChannels")
622 2 }
623
624 2 void bro_tensor_concatNchwChannelsBackward(void* dY, int32_t N, int32_t H, int32_t W, const int32_t* C_per_part, uint32_t C_len, uint64_t parts_bits) {
625 2 const std::vector<int> channels = asInts(C_per_part, C_len);
626 2 std::vector<brotensor::Tensor*> parts;
627
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 not taken.
✓ Branch 7 taken 2 times.
 2 if (!need("concatNchwChannelsBackward", {dY}) || !readTensorArray(parts_bits, parts, "concatNchwChannelsBackward")) return;
628 BROTENSOR_API_TRY
629
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 brotensor::concat_nchw_channels_backward(*toTensor(dY), N, H, W, channels, parts);
630
0/14
✗ 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.
✗ Branch 12 not taken.
✗ Branch 13 not taken.
 2 BROTENSOR_API_CATCH("concatNchwChannelsBackward")
631 2 }
632
633 // ---- counter-based RNG + init (key / counter / state: BigInt or Number) ----
634
635 2 void bro_tensor_randUniform(uint64_t key_bits, uint64_t counter_bits, void* Y) {
636
2/4
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!need("randUniform", {Y}) || !needFp32Out("randUniform", toTensor(Y))) return;
637 BROTENSOR_API_TRY
638
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 4 brotensor::rand_uniform(bronze::embed::toUint64(bronze::Value{key_bits}),
639
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 bronze::embed::toUint64(bronze::Value{counter_bits}), *toTensor(Y));
640
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("randUniform")
641 2 }
642
643 12 void bro_tensor_randn(uint64_t key_bits, uint64_t counter_bits, void* Y) {
644
2/4
✗ Branch 0 not taken.
✓ Branch 1 taken 12 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 12 times.
 12 if (!need("randn", {Y}) || !needFp32Out("randn", toTensor(Y))) return;
645 BROTENSOR_API_TRY
646
2/4
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
 24 brotensor::randn(bronze::embed::toUint64(bronze::Value{key_bits}),
647
2/4
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
 12 bronze::embed::toUint64(bronze::Value{counter_bits}), *toTensor(Y));
648
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.
 12 BROTENSOR_API_CATCH("randn")
649 12 }
650
651 4 void bro_tensor_randBernoulli(double p, uint64_t key_bits, uint64_t counter_bits, void* Y) {
652
2/4
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
 4 if (!need("randBernoulli", {Y}) || !needFp32Out("randBernoulli", toTensor(Y))) return;
653 BROTENSOR_API_TRY
654
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 8 brotensor::rand_bernoulli(static_cast<float>(p), bronze::embed::toUint64(bronze::Value{key_bits}),
655
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 bronze::embed::toUint64(bronze::Value{counter_bits}), *toTensor(Y));
656
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("randBernoulli")
657 4 }
658
659 2 void bro_tensor_randnTruncated(double lo, double hi, uint64_t key_bits, uint64_t counter_bits, void* Y) {
660
2/4
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!need("randnTruncated", {Y}) || !needFp32Out("randnTruncated", toTensor(Y))) return;
661 BROTENSOR_API_TRY
662
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 4 brotensor::randn_truncated(static_cast<float>(lo), static_cast<float>(hi),
663
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 bronze::embed::toUint64(bronze::Value{key_bits}),
664
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 bronze::embed::toUint64(bronze::Value{counter_bits}), *toTensor(Y));
665
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("randnTruncated")
666 2 }
667
668 // xavierInit(W, rngState) -> advanced state. The old binding returned a
669 // BigInt; the embed API mints no BigInt, so the state comes back as a Number
670 // (exact below 2^53 — thread it through successive inits as before).
671 ✗ double bro_tensor_xavierInit(void* W, uint64_t state_bits) {
672 ✗ uint64_t state = bronze::embed::toUint64(bronze::Value{state_bits});
673 ✗ if (!need("xavierInit", {W})) return static_cast<double>(state);
674 BROTENSOR_API_TRY
675 ✗ brotensor::xavier_init(*toTensor(W), state);
676 ✗ BROTENSOR_API_CATCH("xavierInit")
677 ✗ return static_cast<double>(state);
678 ✗ }
679
680 } // extern "C"
681
682 namespace brotensor::api {
683
684 16 bool readTensorArray(uint64_t bits, std::vector<brotensor::Tensor*>& out, const char* label) {
685 namespace ev = bronze::embed;
686 16 ev::Persistent arr(bronze::Value{bits});
687 16 out.clear();
688
3/6
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 16 times.
 16 if (!ev::isObject(arr.get())) {
689 ✗ setError(std::string(label) + ": expected an array of GpuTensors");
690 ✗ return false;
691 }
692
4/8
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 16 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 16 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
 16 const auto n = static_cast<uint32_t>(ev::toDouble(ev::getProperty(arr.get(), "length")));
693
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 out.reserve(n);
694
2/2
✓ Branch 0 taken 32 times.
✓ Branch 1 taken 16 times.
 48 for (uint32_t i = 0; i < n; ++i) {
695
3/6
✓ Branch 0 taken 32 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 32 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 32 times.
✗ Branch 5 not taken.
 32 void* d = ev::handleData(ev::getElement(arr.get(), i));
696
2/4
✓ Branch 0 taken 32 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 32 times.
✗ Branch 3 not taken.
 32 auto* t = d ? toTensor(d) : nullptr;
697
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 32 times.
 32 if (!t) {
698 ✗ setError(std::string(label) + ": element " + std::to_string(i) + " is not a GpuTensor");
699 ✗ return false;
700 }
701
1/2
✓ Branch 0 taken 32 times.
✗ Branch 1 not taken.
 32 out.push_back(t);
702 32 }
703 16 return true;
704 16 }
705
706 } // namespace brotensor::api
707