GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 97.2% 351 / 0 / 361
Functions: 100.0% 29 / 0 / 29
Branches: 38.0% 429 / 0 / 1130
src/api/native_tensor_attention.cpp
Line Branch Exec Source
1 // The attention family of bro.tensor, ported from the QuickJS binding
2 // (tensor_bindings.cpp: attention/mha; tensor_bindings_attention.cpp: self /
3 // cross / bias attention, the flash-attention family, resblock). Every
4 // nullable tensor (mask|null, the projection biases, Ctx|null, the optional
5 // gradient outputs) crosses as `dynamic`; js/tensor_ext.js has already
6 // checked each one is null or a GpuTensor, and unpacks the options-object
7 // forms (flashAttentionQkvoBackward(opts), resblockForward(opts),
8 // resblockBackward(opts)) into these positional calls.
9 //
10 // Each body checks its operands against the shapes X (and Ctx) imply before
11 // the op runs (api_internal.h "input-size contract"): the weights are (D,D)
12 // or (D,D_ctx), the caches and upstream gradients match the forward, the
13 // accumulated weight gradients match their weights, and a mask covers every
14 // key.
15
16 #include "native_tensor_decl.h"
17 #include "api_internal.h"
18
19 using namespace brotensor::api;
20 using brotensor::Tensor;
21
22 namespace {
23
24 388 Tensor* T(void* p) { return toTensor(p); }
25
26 // Square (D,D) projections Wq/Wk/Wv/Wo in X's dtype.
27 24 bool needSquareProj(const char* L, const Tensor* X, void* Wq, void* Wk, void* Wv, void* Wo) {
28 24 const int64_t DD = elems({X->cols, X->cols});
29 24 return needOperands(L, X, {{"Wq", T(Wq), DD}, {"Wk", T(Wk), DD}, {"Wv", T(Wv), DD}, {"Wo", T(Wo), DD}});
30 }
31
32 // Cross projections: Wq/Wo (D,D), Wk/Wv (D,D_ctx).
33 12 bool needCrossProj(const char* L, const Tensor* X, const Tensor* Ctx, void* Wq, void* Wk, void* Wv, void* Wo) {
34 12 const int64_t DD = elems({X->cols, X->cols});
35 12 const int64_t DC = elems({X->cols, Ctx->cols});
36
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 12 times.
 12 if (!needSameDtype(L, "Ctx", Ctx, X)) return false;
37 12 return needOperands(L, X, {{"Wq", T(Wq), DD}, {"Wk", T(Wk), DC}, {"Wv", T(Wv), DC}, {"Wo", T(Wo), DD}});
38 12 }
39
40 // The multi-head backward caches of a self-attention forward over X (K,D).
41 4 bool needMhaCaches(const char* L, const Tensor* X, int32_t numHeads,
42 void* dO, void* Qh, void* Kh, void* Vh, void* Attnh, void* Yconcat) {
43 4 const int64_t KD = elems({X->rows, X->cols});
44 12 return needOperands(L, X, {{"dO", T(dO), KD}, {"Qh", T(Qh), KD}, {"Kh", T(Kh), KD}, {"Vh", T(Vh), KD},
45 4 {"Attnh", T(Attnh), elems({numHeads, X->rows, X->rows})},
46 4 {"Yconcat", T(Yconcat), KD}});
47 }
48
49 // Accumulated (D,D) weight gradients.
50 6 bool needSquareGrads(const char* L, const Tensor* X, void* dWq, void* dWk, void* dWv, void* dWo) {
51 6 const int64_t DD = elems({X->cols, X->cols});
52 6 return needOperands(L, X, {{"dWq", T(dWq), DD}, {"dWk", T(dWk), DD}, {"dWv", T(dWv), DD}, {"dWo", T(dWo), DD}});
53 }
54
55 // Pre-projected flash attention: Q (Lq,D), K/V (Lk,D), numHeads | D, a
56 // causal run needs Lq == Lk.
57 8 bool needFlashQkv(const char* L, const Tensor* Q, const Tensor* K, const Tensor* V, int32_t numHeads, bool causal) {
58
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (!needHeads(L, Q->cols, numHeads)) return false;
59
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 (K->cols != Q->cols || V->cols != Q->cols || V->rows != K->rows) {
60 ✗ setError(std::string(L) + ": Q is (Lq,D); K and V must both be (Lk,D)");
61 ✗ return false;
62 }
63
3/4
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 6 times.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 8 if (causal && Q->rows != K->rows) {
64 ✗ setError(std::string(L) + ": causal attention needs Lq == Lk");
65 ✗ return false;
66 }
67
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 return needSameDtype(L, "K", K, Q) && needSameDtype(L, "V", V, Q);
68 8 }
69
70 // The optional (D,1) projection biases of the qkvo family, in X's dtype.
71 4 bool needBiases(const char* L, const Tensor* X, std::initializer_list<Operand> biases) {
72 4 return needOperands(L, X, biases);
73 }
74
75 } // namespace
76
77 extern "C" {
78
79 6 void bro_tensor_attentionForward(void* X, void* Wq, void* Wk, void* Wv, void* Wo, uint64_t mask_bits, void* Q, void* K, void* V, void* Attn, void* Y_pre_Wo, void* O) {
80 6 const char* L = "attentionForward";
81
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
 6 if (!need(L, {X, Wq, Wk, Wv, Wo, Q, K, V, Attn, Y_pre_Wo, O})) return;
82 6 const Tensor* x = T(X);
83
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 6 times.
 6 if (!needSquareProj(L, x, Wq, Wk, Wv, Wo) || !needMask(L, tensorFromValue(mask_bits), x->rows)) return;
84 BROTENSOR_API_TRY
85
6/12
✓ 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.
✓ Branch 8 taken 6 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 6 times.
✗ Branch 11 not taken.
 12 brotensor::attention_forward(*toTensor(X), *toTensor(Wq), *toTensor(Wk), *toTensor(Wv), *toTensor(Wo),
86
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 maskPtr(mask_bits),
87
5/10
✓ 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.
✓ Branch 8 taken 6 times.
✗ Branch 9 not taken.
 6 *toTensor(Q), *toTensor(K), *toTensor(V), *toTensor(Attn), *toTensor(Y_pre_Wo),
88
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 *toTensor(O));
89
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(L)
90 6 }
91
92 2 void bro_tensor_attentionBackward(void* dO, void* X, void* Q, void* K, void* V, void* Attn, void* Y_pre_Wo, void* Wq, void* Wk, void* Wv, void* Wo, uint64_t mask_bits, void* dX, void* dWq, void* dWk, void* dWv, void* dWo) {
93 2 const char* L = "attentionBackward";
94
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {dO, X, Q, K, V, Attn, Y_pre_Wo, Wq, Wk, Wv, Wo, dX, dWq, dWk, dWv, dWo})) return;
95 2 const Tensor* x = T(X);
96 2 const int64_t ND = elems({x->rows, x->cols});
97
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!needSquareProj(L, x, Wq, Wk, Wv, Wo) || !needSquareGrads(L, x, dWq, dWk, dWv, dWo)) return;
98
3/6
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 2 times.
 6 if (!needOperands(L, x, {{"dO", T(dO), ND}, {"Q", T(Q), ND}, {"K", T(K), ND}, {"V", T(V), ND},
99 4 {"Attn", T(Attn), elems({x->rows, x->rows})}, {"Y_pre_Wo", T(Y_pre_Wo), ND}})) return;
100
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needMask(L, tensorFromValue(mask_bits), x->rows)) return;
101 BROTENSOR_API_TRY
102
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 brotensor::attention_backward(*toTensor(dO), *toTensor(X), *toTensor(Q), *toTensor(K), *toTensor(V),
103
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(Attn), *toTensor(Y_pre_Wo),
104
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(Wq), *toTensor(Wk), *toTensor(Wv), *toTensor(Wo),
105
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 maskPtr(mask_bits),
106
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 *toTensor(dX), *toTensor(dWq), *toTensor(dWk), *toTensor(dWv), *toTensor(dWo));
107
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)
108 2 }
109
110 2 void bro_tensor_mhaForward(void* X, void* Wq, void* Wk, void* Wv, void* Wo, uint64_t mask_bits, int32_t numHeads, void* Qh, void* Kh, void* Vh, void* Attnh, void* Yconcat, void* O) {
111 2 const char* L = "mhaForward";
112
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {X, Wq, Wk, Wv, Wo, Qh, Kh, Vh, Attnh, Yconcat, O})) return;
113 2 const Tensor* x = T(X);
114
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!needHeads(L, x->cols, numHeads) || !needSquareProj(L, x, Wq, Wk, Wv, Wo)) return;
115
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needMask(L, tensorFromValue(mask_bits), x->rows)) return;
116 BROTENSOR_API_TRY
117
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 brotensor::mha_forward(*toTensor(X), *toTensor(Wq), *toTensor(Wk), *toTensor(Wv), *toTensor(Wo),
118
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 maskPtr(mask_bits), numHeads,
119
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 *toTensor(Qh), *toTensor(Kh), *toTensor(Vh), *toTensor(Attnh), *toTensor(Yconcat),
120
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 *toTensor(O));
121
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)
122 2 }
123
124 2 void bro_tensor_mhaBackward(void* dO, void* X, void* Qh, void* Kh, void* Vh, void* Attnh, void* Yconcat, void* Wq, void* Wk, void* Wv, void* Wo, uint64_t mask_bits, int32_t numHeads, void* dX, void* dWq, void* dWk, void* dWv, void* dWo) {
125 2 const char* L = "mhaBackward";
126
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {dO, X, Qh, Kh, Vh, Attnh, Yconcat, Wq, Wk, Wv, Wo, dX, dWq, dWk, dWv, dWo})) return;
127 2 const Tensor* x = T(X);
128
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!needHeads(L, x->cols, numHeads) || !needSquareProj(L, x, Wq, Wk, Wv, Wo)) return;
129
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!needSquareGrads(L, x, dWq, dWk, dWv, dWo) || !needMhaCaches(L, x, numHeads, dO, Qh, Kh, Vh, Attnh, Yconcat)) return;
130
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needMask(L, tensorFromValue(mask_bits), x->rows)) return;
131 BROTENSOR_API_TRY
132
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 brotensor::mha_backward(*toTensor(dO), *toTensor(X), *toTensor(Qh), *toTensor(Kh), *toTensor(Vh),
133
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(Attnh), *toTensor(Yconcat),
134
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(Wq), *toTensor(Wk), *toTensor(Wv), *toTensor(Wo),
135
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 maskPtr(mask_bits), numHeads,
136
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 *toTensor(dX), *toTensor(dWq), *toTensor(dWk), *toTensor(dWv), *toTensor(dWo));
137
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)
138 2 }
139
140 8 void bro_tensor_selfAttentionForward(void* X, void* Wq, void* Wk, void* Wv, void* Wo, uint64_t mask_bits, int32_t numHeads, void* O) {
141 8 const char* L = "selfAttentionForward";
142
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (!need(L, {X, Wq, Wk, Wv, Wo, O})) return;
143 8 const Tensor* x = T(X);
144
4/4
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 4 times.
 8 if (!needHeads(L, x->cols, numHeads) || !needSquareProj(L, x, Wq, Wk, Wv, Wo)) return;
145
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 2 times.
 4 if (!needMask(L, tensorFromValue(mask_bits), x->rows)) return;
146 BROTENSOR_API_TRY
147
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 brotensor::self_attention_forward(*toTensor(X), *toTensor(Wq), *toTensor(Wk), *toTensor(Wv), *toTensor(Wo),
148
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 maskPtr(mask_bits), numHeads, *toTensor(O));
149
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)
150 8 }
151
152 2 void bro_tensor_selfAttentionForwardTrain(void* X, void* Wq, void* Wk, void* Wv, void* Wo, uint64_t mask_bits, int32_t numHeads, void* Qh, void* Kh, void* Vh, void* Attnh, void* Yconcat, void* O) {
153 2 const char* L = "selfAttentionForwardTrain";
154
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {X, Wq, Wk, Wv, Wo, Qh, Kh, Vh, Attnh, Yconcat, O})) return;
155 2 const Tensor* x = T(X);
156
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!needHeads(L, x->cols, numHeads) || !needSquareProj(L, x, Wq, Wk, Wv, Wo)) return;
157
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needMask(L, tensorFromValue(mask_bits), x->rows)) return;
158 BROTENSOR_API_TRY
159
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 brotensor::self_attention_forward_train(*toTensor(X), *toTensor(Wq), *toTensor(Wk), *toTensor(Wv), *toTensor(Wo),
160
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 maskPtr(mask_bits), numHeads,
161
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(Qh), *toTensor(Kh), *toTensor(Vh), *toTensor(Attnh),
162
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(Yconcat), *toTensor(O));
163
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)
164 2 }
165
166 2 void bro_tensor_selfAttentionBackward(void* dO, void* X, void* Qh, void* Kh, void* Vh, void* Attnh, void* Yconcat, void* Wq, void* Wk, void* Wv, void* Wo, uint64_t mask_bits, int32_t numHeads, void* dX, void* dWq, void* dWk, void* dWv, void* dWo) {
167 2 const char* L = "selfAttentionBackward";
168
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {dO, X, Qh, Kh, Vh, Attnh, Yconcat, Wq, Wk, Wv, Wo, dX, dWq, dWk, dWv, dWo})) return;
169 2 const Tensor* x = T(X);
170
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!needHeads(L, x->cols, numHeads) || !needSquareProj(L, x, Wq, Wk, Wv, Wo)) return;
171
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!needSquareGrads(L, x, dWq, dWk, dWv, dWo) || !needMhaCaches(L, x, numHeads, dO, Qh, Kh, Vh, Attnh, Yconcat)) return;
172
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needMask(L, tensorFromValue(mask_bits), x->rows)) return;
173 BROTENSOR_API_TRY
174
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 brotensor::self_attention_backward(*toTensor(dO), *toTensor(X), *toTensor(Qh), *toTensor(Kh), *toTensor(Vh),
175
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(Attnh), *toTensor(Yconcat),
176
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(Wq), *toTensor(Wk), *toTensor(Wv), *toTensor(Wo),
177
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 maskPtr(mask_bits), numHeads,
178
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 *toTensor(dX), *toTensor(dWq), *toTensor(dWk), *toTensor(dWv), *toTensor(dWo));
179
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)
180 2 }
181
182 2 void bro_tensor_selfAttentionBiasForward(void* X, void* Wq, void* Wk, void* Wv, void* Wo, uint64_t mask_bits, uint64_t attnBias_bits, int32_t numHeads, double scale, void* O) {
183 2 const char* L = "selfAttentionBiasForward";
184
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {X, Wq, Wk, Wv, Wo, O})) return;
185 2 const Tensor* x = T(X);
186
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!needHeads(L, x->cols, numHeads) || !needSquareProj(L, x, Wq, Wk, Wv, Wo)) return;
187
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needMask(L, tensorFromValue(mask_bits), x->rows)) return;
188 // attn_bias: (numHeads*L, L) FP32.
189
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!needMask(L, tensorFromValue(attnBias_bits), elems({numHeads, x->rows, x->rows}), "attnBias")) return;
190 BROTENSOR_API_TRY
191
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 brotensor::self_attention_bias_forward(*toTensor(X), *toTensor(Wq), *toTensor(Wk), *toTensor(Wv), *toTensor(Wo),
192 nullptr, nullptr, nullptr, nullptr,
193
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 maskPtr(mask_bits), tensorFromValue(attnBias_bits),
194
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 numHeads, static_cast<float>(scale), *toTensor(O));
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(L)
196 2 }
197
198 2 void bro_tensor_crossAttentionForward(void* X, void* Ctx, void* Wq, void* Wk, void* Wv, void* Wo, uint64_t mask_bits, int32_t numHeads, void* O) {
199 2 const char* L = "crossAttentionForward";
200
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {X, Ctx, Wq, Wk, Wv, Wo, O})) return;
201 2 const Tensor* x = T(X);
202 2 const Tensor* ctx = T(Ctx);
203
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!needHeads(L, x->cols, numHeads) || !needCrossProj(L, x, ctx, Wq, Wk, Wv, Wo)) return;
204
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needMask(L, tensorFromValue(mask_bits), ctx->rows)) return;
205 BROTENSOR_API_TRY
206
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 brotensor::cross_attention_forward(*toTensor(X), *toTensor(Ctx), *toTensor(Wq), *toTensor(Wk), *toTensor(Wv),
207
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(Wo), maskPtr(mask_bits), numHeads, *toTensor(O));
208
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)
209 2 }
210
211 2 void bro_tensor_crossAttentionForwardWithAttn(void* X, void* Ctx, void* Wq, void* Wk, void* Wv, void* Wo, uint64_t mask_bits, uint64_t attnLogitBias_bits, int32_t numHeads, void* O, void* AttnAvg) {
212 2 const char* L = "crossAttentionForwardWithAttn";
213
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {X, Ctx, Wq, Wk, Wv, Wo, O, AttnAvg})) return;
214 2 const Tensor* x = T(X);
215 2 const Tensor* ctx = T(Ctx);
216
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!needHeads(L, x->cols, numHeads) || !needCrossProj(L, x, ctx, Wq, Wk, Wv, Wo)) return;
217
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needMask(L, tensorFromValue(mask_bits), ctx->rows)) return;
218 // attn_logit_bias: (Lq, Lk) FP32.
219
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!needMask(L, tensorFromValue(attnLogitBias_bits), elems({x->rows, ctx->rows}), "attnLogitBias")) return;
220 BROTENSOR_API_TRY
221
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.
 4 brotensor::cross_attention_forward_with_attn(*toTensor(X), *toTensor(Ctx), *toTensor(Wq), *toTensor(Wk),
222
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(Wv), *toTensor(Wo),
223
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 maskPtr(mask_bits), tensorFromValue(attnLogitBias_bits),
224
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 numHeads, *toTensor(O), *toTensor(AttnAvg));
225
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)
226 2 }
227
228 2 void bro_tensor_crossAttentionForwardTrain(void* X, void* Ctx, void* Wq, void* Wk, void* Wv, void* Wo, uint64_t mask_bits, int32_t numHeads, void* Qh, void* Kh, void* Vh, void* Attnh, void* Yconcat, void* O) {
229 2 const char* L = "crossAttentionForwardTrain";
230
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {X, Ctx, Wq, Wk, Wv, Wo, Qh, Kh, Vh, Attnh, Yconcat, O})) return;
231 2 const Tensor* x = T(X);
232 2 const Tensor* ctx = T(Ctx);
233
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!needHeads(L, x->cols, numHeads) || !needCrossProj(L, x, ctx, Wq, Wk, Wv, Wo)) return;
234
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needMask(L, tensorFromValue(mask_bits), ctx->rows)) return;
235 BROTENSOR_API_TRY
236
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.
 4 brotensor::cross_attention_forward_train(*toTensor(X), *toTensor(Ctx), *toTensor(Wq), *toTensor(Wk),
237
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(Wv), *toTensor(Wo), maskPtr(mask_bits), numHeads,
238
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(Qh), *toTensor(Kh), *toTensor(Vh), *toTensor(Attnh),
239
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(Yconcat), *toTensor(O));
240
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)
241 2 }
242
243 2 void bro_tensor_crossAttentionBackward(void* dO, void* X, void* Ctx, void* Qh, void* Kh, void* Vh, void* Attnh, void* Yconcat, void* Wq, void* Wk, void* Wv, void* Wo, uint64_t mask_bits, int32_t numHeads, void* dX, void* dCtx, void* dWq, void* dWk, void* dWv, void* dWo) {
244 2 const char* L = "crossAttentionBackward";
245
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {dO, X, Ctx, Qh, Kh, Vh, Attnh, Yconcat, Wq, Wk, Wv, Wo, dX, dCtx, dWq, dWk, dWv, dWo})) return;
246 2 const Tensor* x = T(X);
247 2 const Tensor* ctx = T(Ctx);
248
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!needHeads(L, x->cols, numHeads) || !needCrossProj(L, x, ctx, Wq, Wk, Wv, Wo)) return;
249 2 const int64_t QD = elems({x->rows, x->cols});
250 2 const int64_t KD = elems({ctx->rows, x->cols});
251 2 const int64_t DD = elems({x->cols, x->cols});
252 2 const int64_t DC = elems({x->cols, ctx->cols});
253
7/14
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 2 times.
✗ Branch 6 not taken.
✓ Branch 7 taken 2 times.
✗ Branch 8 not taken.
✓ Branch 9 taken 2 times.
✗ Branch 10 not taken.
✓ Branch 11 taken 2 times.
✗ Branch 12 not taken.
✓ Branch 13 taken 2 times.
 14 if (!needOperands(L, x, {{"dO", T(dO), QD}, {"Qh", T(Qh), QD}, {"Kh", T(Kh), KD}, {"Vh", T(Vh), KD},
254 4 {"Attnh", T(Attnh), elems({numHeads, x->rows, ctx->rows})}, {"Yconcat", T(Yconcat), QD},
255 8 {"dWq", T(dWq), DD}, {"dWk", T(dWk), DC}, {"dWv", T(dWv), DC}, {"dWo", T(dWo), DD}})) return;
256
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needMask(L, tensorFromValue(mask_bits), ctx->rows)) return;
257 BROTENSOR_API_TRY
258
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 brotensor::cross_attention_backward(*toTensor(dO), *toTensor(X), *toTensor(Ctx), *toTensor(Qh), *toTensor(Kh),
259
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(Vh), *toTensor(Attnh), *toTensor(Yconcat),
260
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(Wq), *toTensor(Wk), *toTensor(Wv), *toTensor(Wo),
261
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 maskPtr(mask_bits), numHeads,
262
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(dX), *toTensor(dCtx),
263
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(dWq), *toTensor(dWk), *toTensor(dWv), *toTensor(dWo));
264
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)
265 2 }
266
267 4 void bro_tensor_flashAttentionForward(void* Q, void* K, void* V, uint64_t mask_bits, int32_t numHeads, bool causal, void* O) {
268 4 const char* L = "flashAttentionForward";
269
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 if (!need(L, {Q, K, V, O})) return;
270
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!needFlashQkv(L, T(Q), T(K), T(V), numHeads, causal)) return;
271
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (!needMask(L, tensorFromValue(mask_bits), T(K)->rows)) return;
272 BROTENSOR_API_TRY
273
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 brotensor::flash_attention_forward(*toTensor(Q), *toTensor(K), *toTensor(V), maskPtr(mask_bits),
274
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 numHeads, causal, *toTensor(O));
275
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(L)
276 4 }
277
278 2 void bro_tensor_flashAttentionWindowedForward(void* Q, void* K, void* V, uint64_t mask_bits, int32_t numHeads, int32_t window, void* O) {
279 2 const char* L = "flashAttentionWindowedForward";
280
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {Q, K, V, O})) return;
281
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needFlashQkv(L, T(Q), T(K), T(V), numHeads, false)) return;
282
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 2 if (T(K)->rows < T(Q)->rows) { setError("flashAttentionWindowedForward: needs Lk >= Lq"); return; }
283
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needMask(L, tensorFromValue(mask_bits), T(K)->rows)) return;
284 BROTENSOR_API_TRY
285
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.
 4 brotensor::flash_attention_windowed_forward(*toTensor(Q), *toTensor(K), *toTensor(V), maskPtr(mask_bits),
286
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 numHeads, window, *toTensor(O));
287
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)
288 2 }
289
290 2 void bro_tensor_flashAttentionBackward(void* Q, void* K, void* V, void* O, void* dO, uint64_t mask_bits, int32_t numHeads, bool causal, void* dQ, void* dK, void* dV) {
291 2 const char* L = "flashAttentionBackward";
292
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {Q, K, V, O, dO, dQ, dK, dV})) return;
293
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needFlashQkv(L, T(Q), T(K), T(V), numHeads, causal)) return;
294
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needPair(L, "dO", T(dO), T(Q))) return;
295
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needMask(L, tensorFromValue(mask_bits), T(K)->rows)) return;
296 BROTENSOR_API_TRY
297
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 brotensor::flash_attention_backward(*toTensor(Q), *toTensor(K), *toTensor(V), *toTensor(O), *toTensor(dO),
298
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 maskPtr(mask_bits), numHeads, causal,
299
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(dQ), *toTensor(dK), *toTensor(dV));
300
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)
301 2 }
302
303 2 void bro_tensor_flashAttentionQkvoForward(void* X, uint64_t Ctx_bits, void* Wq, uint64_t bq_bits, void* Wk, uint64_t bk_bits, void* Wv, uint64_t bv_bits, void* Wo, uint64_t bo_bits, uint64_t mask_bits, int32_t numHeads, bool causal, void* O) {
304 2 const char* L = "flashAttentionQkvoForward";
305
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {X, Wq, Wk, Wv, Wo, O})) return;
306 2 const Tensor* x = T(X);
307 2 const Tensor* ctx = tensorFromValue(Ctx_bits);
308
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 const Tensor* kv = ctx ? ctx : x;
309
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!needHeads(L, x->cols, numHeads) || !needCrossProj(L, x, kv, Wq, Wk, Wv, Wo)) return;
310
3/6
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 2 times.
 6 if (!needBiases(L, x, {{"bq", tensorFromValue(bq_bits), x->cols}, {"bk", tensorFromValue(bk_bits), x->cols},
311 4 {"bv", tensorFromValue(bv_bits), x->cols}, {"bo", tensorFromValue(bo_bits), x->cols}})) return;
312
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needMask(L, tensorFromValue(mask_bits), kv->rows)) return;
313 BROTENSOR_API_TRY
314
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.
 4 brotensor::flash_attention_qkvo_forward(*toTensor(X), tensorFromValue(Ctx_bits),
315
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(Wq), tensorFromValue(bq_bits),
316
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(Wk), tensorFromValue(bk_bits),
317
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(Wv), tensorFromValue(bv_bits),
318
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(Wo), tensorFromValue(bo_bits),
319
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 maskPtr(mask_bits), numHeads, causal, *toTensor(O));
320
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)
321 2 }
322
323 2 void bro_tensor_flashAttentionQkvoBackward(void* X, uint64_t Ctx_bits, void* Wq, uint64_t bq_bits, void* Wk, uint64_t bk_bits, void* Wv, uint64_t bv_bits, void* Wo, uint64_t bo_bits, uint64_t mask_bits, int32_t numHeads, bool causal, void* dO, void* dX, uint64_t dCtx_bits, void* dWq, uint64_t dbq_bits, void* dWk, uint64_t dbk_bits, void* dWv, uint64_t dbv_bits, void* dWo, uint64_t dbo_bits) {
324 2 const char* L = "flashAttentionQkvoBackward";
325
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {X, Wq, Wk, Wv, Wo, dO, dX, dWq, dWk, dWv, dWo})) return;
326 2 const Tensor* x = T(X);
327 2 const Tensor* ctx = tensorFromValue(Ctx_bits);
328
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 const Tensor* kv = ctx ? ctx : x;
329
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (!needHeads(L, x->cols, numHeads) || !needCrossProj(L, x, kv, Wq, Wk, Wv, Wo)) return;
330 2 const int64_t D = x->cols;
331 2 const int64_t DD = elems({D, D});
332 2 const int64_t DC = elems({D, kv->cols});
333
7/14
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 2 times.
✗ Branch 6 not taken.
✓ Branch 7 taken 2 times.
✗ Branch 8 not taken.
✓ Branch 9 taken 2 times.
✗ Branch 10 not taken.
✓ Branch 11 taken 2 times.
✗ Branch 12 not taken.
✓ Branch 13 taken 2 times.
 14 if (!needBiases(L, x, {{"bq", tensorFromValue(bq_bits), D}, {"bk", tensorFromValue(bk_bits), D},
334 4 {"bv", tensorFromValue(bv_bits), D}, {"bo", tensorFromValue(bo_bits), D},
335 4 {"dbq", tensorFromValue(dbq_bits), D}, {"dbk", tensorFromValue(dbk_bits), D},
336 4 {"dbv", tensorFromValue(dbv_bits), D}, {"dbo", tensorFromValue(dbo_bits), D}})) return;
337
5/10
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 2 times.
✗ Branch 6 not taken.
✓ Branch 7 taken 2 times.
✗ Branch 8 not taken.
✓ Branch 9 taken 2 times.
 10 if (!needOperands(L, x, {{"dO", T(dO), elems({x->rows, D})},
338 8 {"dWq", T(dWq), DD}, {"dWk", T(dWk), DC}, {"dWv", T(dWv), DC}, {"dWo", T(dWo), DD}})) return;
339
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needMask(L, tensorFromValue(mask_bits), kv->rows)) return;
340 BROTENSOR_API_TRY
341
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.
 4 brotensor::flash_attention_qkvo_backward(*toTensor(X), tensorFromValue(Ctx_bits),
342
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(Wq), tensorFromValue(bq_bits),
343
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(Wk), tensorFromValue(bk_bits),
344
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(Wv), tensorFromValue(bv_bits),
345
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(Wo), tensorFromValue(bo_bits),
346
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 maskPtr(mask_bits), numHeads, causal,
347
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 *toTensor(dO),
348
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(dX), tensorFromValue(dCtx_bits),
349
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(dWq), tensorFromValue(dbq_bits),
350
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(dWk), tensorFromValue(dbk_bits),
351
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(dWv), tensorFromValue(dbv_bits),
352
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(dWo), tensorFromValue(dbo_bits));
353
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)
354 2 }
355
356 2 void bro_tensor_flashAttentionProjectKv(void* ctx, void* Wk, uint64_t bk_bits, void* Wv, uint64_t bv_bits, void* K_out, void* V_out) {
357 2 const char* L = "flashAttentionProjectKv";
358
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {ctx, Wk, Wv, K_out, V_out})) return;
359 // Wk/Wv are (D, D_ctx); D is Wk's row count.
360 2 const Tensor* c = T(ctx);
361 2 const Tensor* wk = T(Wk);
362 2 const int64_t D = wk->rows;
363
3/6
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 2 times.
 6 if (!needOperands(L, c, {{"Wk", wk, elems({D, c->cols})}, {"Wv", T(Wv), elems({D, c->cols})},
364 4 {"bk", tensorFromValue(bk_bits), D}, {"bv", tensorFromValue(bv_bits), D}})) return;
365
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 2 if (wk->cols != c->cols) { setError("flashAttentionProjectKv: Wk must be (D, ctx.cols)"); return; }
366 BROTENSOR_API_TRY
367
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::flash_attention_project_kv(*toTensor(ctx), *toTensor(Wk), tensorFromValue(bk_bits),
368
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(Wv), tensorFromValue(bv_bits),
369
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(K_out), *toTensor(V_out));
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(L)
371 2 }
372
373 2 void bro_tensor_flashAttentionQWithKvCachedForward(void* X, void* K, void* V, void* Wq, uint64_t bq_bits, void* Wo, uint64_t bo_bits, uint64_t mask_bits, int32_t numHeads, bool causal, void* O) {
374 2 const char* L = "flashAttentionQWithKvCachedForward";
375
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {X, K, V, Wq, Wo, O})) return;
376 2 const Tensor* x = T(X);
377 2 const int64_t D = x->cols;
378
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needHeads(L, D, numHeads)) return;
379
3/6
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 2 times.
 2 if (T(K)->cols != D || T(V)->cols != D || T(V)->rows != T(K)->rows) {
380 ✗ setError("flashAttentionQWithKvCachedForward: K and V must both be (Lk, D)");
381 ✗ return;
382 }
383
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 2 if (causal && x->rows != T(K)->rows) { setError("flashAttentionQWithKvCachedForward: causal attention needs Lq == Lk"); return; }
384
3/6
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 2 times.
 6 if (!needOperands(L, x, {{"K", T(K), 0}, {"V", T(V), 0}, {"Wq", T(Wq), elems({D, D})}, {"Wo", T(Wo), elems({D, D})},
385 4 {"bq", tensorFromValue(bq_bits), D}, {"bo", tensorFromValue(bo_bits), D}})) return;
386
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!needMask(L, tensorFromValue(mask_bits), T(K)->rows)) return;
387 BROTENSOR_API_TRY
388
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::flash_attention_q_with_kv_cached_forward(*toTensor(X), *toTensor(K), *toTensor(V),
389
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(Wq), tensorFromValue(bq_bits),
390
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(Wo), tensorFromValue(bo_bits),
391
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 maskPtr(mask_bits), numHeads, causal, *toTensor(O));
392
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)
393 2 }
394
395 // The fused SD ResBlock: 3x3 convs W1 (C_out, C_in*9) and W2 (C_out, C_out*9),
396 // per-channel GN params, an optional 1x1 skip.
397 4 static bool needResblock(const char* L, const Tensor* x, void* gamma1, void* beta1, void* W1, uint64_t b1_bits,
398 uint64_t t_emb_shift_bits, void* gamma2, void* beta2, void* W2, uint64_t b2_bits,
399 uint64_t Wskip_bits, uint64_t bskip_bits,
400 int32_t N, int32_t C_in, int32_t C_out, int32_t H, int32_t W, int32_t numGroups) {
401
2/4
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
 4 if (!needPositive(L, {numGroups}) || !needNonNegative(L, {N, C_in, C_out, H, W})) return false;
402
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
 4 if (C_in % numGroups != 0 || C_out % numGroups != 0) {
403 ✗ setError(std::string(L) + ": numGroups must divide C_in and C_out");
404 ✗ return false;
405 }
406
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 4 if (C_in != C_out && !tensorFromValue(Wskip_bits)) {
407 ✗ setError(std::string(L) + ": C_in != C_out needs a Wskip projection");
408 ✗ return false;
409 }
410 48 return needOperands(L, x, {{"X", x, elems({N, C_in, H, W})},
411 8 {"gamma1", T(gamma1), C_in}, {"beta1", T(beta1), C_in},
412 8 {"W1", T(W1), elems({C_out, C_in, 9})}, {"b1", tensorFromValue(b1_bits), C_out},
413 4 {"t_emb_shift", tensorFromValue(t_emb_shift_bits), C_out},
414 8 {"gamma2", T(gamma2), C_out}, {"beta2", T(beta2), C_out},
415 8 {"W2", T(W2), elems({C_out, C_out, 9})}, {"b2", tensorFromValue(b2_bits), C_out},
416 4 {"Wskip", tensorFromValue(Wskip_bits), elems({C_out, C_in})},
417 4 {"bskip", tensorFromValue(bskip_bits), C_out}});
418 4 }
419
420 2 void bro_tensor_resblockForward(void* X, void* gamma1, void* beta1, void* W1, uint64_t b1_bits, uint64_t t_emb_shift_bits, void* gamma2, void* beta2, void* W2, uint64_t b2_bits, uint64_t Wskip_bits, uint64_t bskip_bits, int32_t N, int32_t C_in, int32_t C_out, int32_t H, int32_t W, int32_t numGroups, double eps, void* Y) {
421 2 const char* L = "resblockForward";
422
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {X, gamma1, beta1, W1, gamma2, beta2, W2, Y})) return;
423
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 4 if (!needResblock(L, T(X), gamma1, beta1, W1, b1_bits, t_emb_shift_bits, gamma2, beta2, W2, b2_bits,
424 2 Wskip_bits, bskip_bits, N, C_in, C_out, H, W, numGroups)) return;
425 BROTENSOR_API_TRY
426
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::resblock_forward(*toTensor(X), *toTensor(gamma1), *toTensor(beta1),
427
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(W1), tensorFromValue(b1_bits), tensorFromValue(t_emb_shift_bits),
428
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(gamma2), *toTensor(beta2),
429
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(W2), tensorFromValue(b2_bits),
430
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 tensorFromValue(Wskip_bits), tensorFromValue(bskip_bits),
431
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 N, C_in, C_out, H, W, numGroups, static_cast<float>(eps), *toTensor(Y));
432
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)
433 2 }
434
435 2 void bro_tensor_resblockBackward(void* X, void* gamma1, void* beta1, void* W1, uint64_t b1_bits, uint64_t t_emb_shift_bits, void* gamma2, void* beta2, void* W2, uint64_t b2_bits, uint64_t Wskip_bits, uint64_t bskip_bits, int32_t N, int32_t C_in, int32_t C_out, int32_t H, int32_t W, int32_t numGroups, double eps, void* dY, void* dX, void* dGamma1, void* dBeta1, void* dW1, uint64_t db1_bits, uint64_t dt_emb_shift_bits, void* dGamma2, void* dBeta2, void* dW2, uint64_t db2_bits, uint64_t dWskip_bits, uint64_t dbskip_bits) {
436 2 const char* L = "resblockBackward";
437
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (!need(L, {X, gamma1, beta1, W1, gamma2, beta2, W2, dY, dX, dGamma1, dBeta1, dW1, dGamma2, dBeta2, dW2})) return;
438 2 const Tensor* x = T(X);
439
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 4 if (!needResblock(L, x, gamma1, beta1, W1, b1_bits, t_emb_shift_bits, gamma2, beta2, W2, b2_bits,
440 2 Wskip_bits, bskip_bits, N, C_in, C_out, H, W, numGroups)) return;
441 // The gradients mirror their parameters.
442
12/24
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 2 times.
✗ Branch 6 not taken.
✓ Branch 7 taken 2 times.
✗ Branch 8 not taken.
✓ Branch 9 taken 2 times.
✗ Branch 10 not taken.
✓ Branch 11 taken 2 times.
✗ Branch 12 not taken.
✓ Branch 13 taken 2 times.
✗ Branch 14 not taken.
✓ Branch 15 taken 2 times.
✗ Branch 16 not taken.
✓ Branch 17 taken 2 times.
✗ Branch 18 not taken.
✓ Branch 19 taken 2 times.
✗ Branch 20 not taken.
✓ Branch 21 taken 2 times.
✗ Branch 22 not taken.
✓ Branch 23 taken 2 times.
 24 if (!needOperands(L, x, {{"dY", T(dY), elems({N, C_out, H, W})},
443 4 {"dGamma1", T(dGamma1), C_in}, {"dBeta1", T(dBeta1), C_in},
444 4 {"dW1", T(dW1), elems({C_out, C_in, 9})}, {"db1", tensorFromValue(db1_bits), C_out},
445 2 {"dt_emb_shift", tensorFromValue(dt_emb_shift_bits), C_out},
446 4 {"dGamma2", T(dGamma2), C_out}, {"dBeta2", T(dBeta2), C_out},
447 4 {"dW2", T(dW2), elems({C_out, C_out, 9})}, {"db2", tensorFromValue(db2_bits), C_out},
448 2 {"dWskip", tensorFromValue(dWskip_bits), elems({C_out, C_in})},
449 2 {"dbskip", tensorFromValue(dbskip_bits), C_out}})) return;
450 BROTENSOR_API_TRY
451
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::resblock_backward(*toTensor(X), *toTensor(gamma1), *toTensor(beta1),
452
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(W1), tensorFromValue(b1_bits), tensorFromValue(t_emb_shift_bits),
453
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(gamma2), *toTensor(beta2),
454
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(W2), tensorFromValue(b2_bits),
455
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 tensorFromValue(Wskip_bits), tensorFromValue(bskip_bits),
456 2 N, C_in, C_out, H, W, numGroups, static_cast<float>(eps),
457
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(dY), *toTensor(dX),
458
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(dGamma1), *toTensor(dBeta1),
459
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(dW1), tensorFromValue(db1_bits), tensorFromValue(dt_emb_shift_bits),
460
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(dGamma2), *toTensor(dBeta2),
461
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 *toTensor(dW2), tensorFromValue(db2_bits),
462
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 tensorFromValue(dWskip_bits), tensorFromValue(dbskip_bits));
463
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)
464 2 }
465
466 } // extern "C"
467