GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 93.9% 169 / 0 / 180
Functions: 100.0% 5 / 0 / 5
Branches: 54.3% 102 / 0 / 188
src/cpu/self_attention_decomposed_rel_pos.cpp
Line Branch Exec Source
1 // ─── CPU self-attention with decomposed 2D relative-position bias ──────────
2 //
3 // FP32 scalar host reference for the SAM / ViTDet image-encoder attention.
4 // Multi-head self-attention where the pre-softmax bias is the *decomposed*
5 // 2D relative-position term from Dosovitskiy-style ViTDet (segment_anything
6 // `add_decomposed_rel_pos`):
7 //
8 // r_q = Q_h[q] (projected, UNSCALED query)
9 // bias[h,q,k] = r_q . Rh[qh, kh] + r_q . Rw[qw, kw]
10 // S[h,q,k] = scale * (Q_h[q] . K_h[k]) + bias[h,q,k]
11 // O = concat_h( softmax_k(S[h]) @ V_h ) @ Wo
12 //
13 // where a token index t maps to grid coords (t / grid_w, t % grid_w) over a
14 // grid_h × grid_w patch grid (so L == grid_h*grid_w), and
15 // Rh[qh, kh] = rel_pos_h[(qh - kh) + (grid_h - 1)] (length head_dim)
16 // Rw[qw, kw] = rel_pos_w[(qw - kw) + (grid_w - 1)].
17 //
18 // This is the q*size == k*size case of segment_anything's get_rel_pos (no
19 // rel-pos interpolation): rel_pos_h has exactly 2*grid_h-1 rows, indexed by the
20 // signed query−key row offset shifted into [0, 2*grid_h-2]. Windowed blocks
21 // call this per window (grid_h == grid_w == window); global blocks call it once
22 // over the full 64×64 grid. The bias is data-dependent (it reads Q), which is
23 // why it can't be expressed through the static-bias self_attention_bias op.
24 //
25 // Unlike that static (num_heads*L, L) bias, the decomposed form is never
26 // materialised: per (head, query) we precompute length-grid_h and length-grid_w
27 // dot-product vectors and add bias[k] = relh[kh] + relw[kw] inside the score
28 // loop — O(L*(grid_h+grid_w)*head_dim) work, not O(L*L*head_dim).
29 //
30 // qkv and output projections carry optional biases (SAM uses both). Wq/Wk/Wv/Wo
31 // are (D, D); per-head split takes contiguous weight rows hh*dh .. hh*dh+dh.
32
33 #include <brotensor/tensor.h>
34
35 #include <cmath>
36 #include <stdexcept>
37 #include <string>
38 #include <vector>
39
40 namespace brotensor::detail::cpu {
41
42 namespace {
43
44 712 inline void check_fp32(const ::brotensor::Tensor& t, const char* name) {
45
1/2
✓ Branch 0 taken 712 times.
✗ Branch 1 not taken.
 712 if (t.dtype != Dtype::FP32) {
46 ✗ throw std::runtime_error(
47 ✗ std::string("self_attention_decomposed_rel_pos_forward: ") + name +
48 " must be FP32 (CPU backend is FP32-only)");
49 }
50 712 }
51
52 } // namespace
53
54 66 void self_attention_decomposed_rel_pos_forward(
55 const ::brotensor::Tensor& X,
56 const ::brotensor::Tensor& Wq, const ::brotensor::Tensor* bq,
57 const ::brotensor::Tensor& Wk, const ::brotensor::Tensor* bk,
58 const ::brotensor::Tensor& Wv, const ::brotensor::Tensor* bv,
59 const ::brotensor::Tensor& Wo, const ::brotensor::Tensor* bo,
60 const ::brotensor::Tensor& rel_pos_h,
61 const ::brotensor::Tensor& rel_pos_w,
62 int num_heads, int grid_h, int grid_w, float scale,
63 ::brotensor::Tensor& O) {
64 66 check_fp32(X, "X");
65 66 check_fp32(Wq, "Wq"); check_fp32(Wk, "Wk");
66 66 check_fp32(Wv, "Wv"); check_fp32(Wo, "Wo");
67 66 check_fp32(rel_pos_h, "rel_pos_h"); check_fp32(rel_pos_w, "rel_pos_w");
68
69 66 const int L = X.rows;
70 66 const int D = X.cols;
71
1/2
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
 66 if (num_heads <= 0 || D % num_heads != 0)
72 ✗ throw std::runtime_error(
73 "self_attention_decomposed_rel_pos_forward: num_heads must divide D");
74
1/2
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
 66 if (grid_h <= 0 || grid_w <= 0 || grid_h * grid_w != L)
75 ✗ throw std::runtime_error(
76 "self_attention_decomposed_rel_pos_forward: grid_h*grid_w must equal X.rows");
77
1/2
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
 132 if (Wq.rows != D || Wq.cols != D || Wk.rows != D || Wk.cols != D ||
78 66 Wv.rows != D || Wv.cols != D || Wo.rows != D || Wo.cols != D)
79 ✗ throw std::runtime_error(
80 "self_attention_decomposed_rel_pos_forward: Wq/Wk/Wv/Wo must be (D, D)");
81
82 66 const int H = num_heads;
83 66 const int dh = D / H;
84
1/2
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
 66 if (rel_pos_h.rows != 2 * grid_h - 1 || rel_pos_h.cols != dh)
85 ✗ throw std::runtime_error(
86 "self_attention_decomposed_rel_pos_forward: rel_pos_h must be (2*grid_h-1, head_dim)");
87
1/2
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
 66 if (rel_pos_w.rows != 2 * grid_w - 1 || rel_pos_w.cols != dh)
88 ✗ throw std::runtime_error(
89 "self_attention_decomposed_rel_pos_forward: rel_pos_w must be (2*grid_w-1, head_dim)");
90
91 330 auto bias_ptr = [&](const ::brotensor::Tensor* b, const char* name) -> const float* {
92
3/4
✓ Branch 0 taken 240 times.
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 240 times.
 264 if (!b || !b->data) return nullptr;
93 240 check_fp32(*b, name);
94
1/2
✓ Branch 0 taken 240 times.
✗ Branch 1 not taken.
 240 if (b->size() != D)
95 ✗ throw std::runtime_error(
96 ✗ std::string("self_attention_decomposed_rel_pos_forward: ") + name +
97 " must have D entries");
98 240 return b->host_f32();
99 264 };
100 66 const float* bqp = bias_ptr(bq, "bq");
101 66 const float* bkp = bias_ptr(bk, "bk");
102 66 const float* bvp = bias_ptr(bv, "bv");
103 66 const float* bop = bias_ptr(bo, "bo");
104
105
4/6
✓ Branch 0 taken 40 times.
✓ Branch 1 taken 26 times.
✓ Branch 2 taken 40 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 40 times.
 66 if (O.rows != L || O.cols != D || O.dtype != Dtype::FP32)
106 26 O.resize(L, D, Dtype::FP32);
107
2/4
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 66 times.
 66 if (L == 0 || D == 0) return;
108
109 66 const float* Xp = X.host_f32();
110 66 const float* Wqp = Wq.host_f32();
111 66 const float* Wkp = Wk.host_f32();
112 66 const float* Wvp = Wv.host_f32();
113 66 const float* Wop = Wo.host_f32();
114 66 const float* Rhp = rel_pos_h.host_f32();
115 66 const float* Rwp = rel_pos_w.host_f32();
116 66 float* Op = O.host_f32_mut();
117
118 // Per-head projections: Qh / Kh / Vh laid out (H*L, dh), each plus its bias.
119 66 std::vector<float> Qh(static_cast<size_t>(H) * L * dh);
120
1/2
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
 66 std::vector<float> Kh(static_cast<size_t>(H) * L * dh);
121
1/2
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
 66 std::vector<float> Vh(static_cast<size_t>(H) * L * dh);
122 264 auto project = [&](const float* W, const float* b, std::vector<float>& Out) {
123
2/2
✓ Branch 0 taken 759 times.
✓ Branch 1 taken 198 times.
 957 for (int hh = 0; hh < H; ++hh) {
124
2/2
✓ Branch 0 taken 16329 times.
✓ Branch 1 taken 759 times.
 17088 for (int i = 0; i < L; ++i) {
125 16329 const float* xr = Xp + static_cast<size_t>(i) * D;
126
2/2
✓ Branch 0 taken 130632 times.
✓ Branch 1 taken 16329 times.
 146961 for (int j = 0; j < dh; ++j) {
127 130632 const int o = hh * dh + j;
128 130632 const float* wr = W + static_cast<size_t>(o) * D;
129
2/2
✓ Branch 0 taken 110928 times.
✓ Branch 1 taken 19704 times.
 130632 float acc = b ? b[o] : 0.0f;
130
2/2
✓ Branch 0 taken 5579712 times.
✓ Branch 1 taken 130632 times.
 5710344 for (int k = 0; k < D; ++k) acc += xr[k] * wr[k];
131 130632 Out[(static_cast<size_t>(hh) * L + i) * dh + j] = acc;
132 130632 }
133 16329 }
134 759 }
135 198 };
136
1/2
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
 66 project(Wqp, bqp, Qh);
137
1/2
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
 66 project(Wkp, bkp, Kh);
138
1/2
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
 66 project(Wvp, bvp, Vh);
139
140
1/2
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
 66 std::vector<float> Yc(static_cast<size_t>(L) * D, 0.0f);
141
1/2
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
 66 std::vector<float> srow(L);
142
1/2
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
 66 std::vector<float> relh(grid_h); // per-key-row rel-h bias for current query
143
1/2
✓ Branch 0 taken 66 times.
✗ Branch 1 not taken.
 66 std::vector<float> relw(grid_w); // per-key-col rel-w bias for current query
144
145
2/2
✓ Branch 0 taken 253 times.
✓ Branch 1 taken 66 times.
 319 for (int hh = 0; hh < H; ++hh) {
146
2/2
✓ Branch 0 taken 5443 times.
✓ Branch 1 taken 253 times.
 5696 for (int i = 0; i < L; ++i) {
147 5443 const float* qr = &Qh[(static_cast<size_t>(hh) * L + i) * dh];
148 5443 const int qh = i / grid_w;
149 5443 const int qw = i % grid_w;
150
151 // Decompose the relative-position bias: relh[kh] = q . rel_pos_h[qh-kh+gh-1],
152 // relw[kw] = q . rel_pos_w[qw-kw+gw-1]. bias(kh,kw) = relh[kh] + relw[kw].
153
2/2
✓ Branch 0 taken 29883 times.
✓ Branch 1 taken 5443 times.
 35326 for (int kh = 0; kh < grid_h; ++kh) {
154 29883 const float* rr = Rhp + static_cast<size_t>(qh - kh + grid_h - 1) * dh;
155 29883 float acc = 0.0f;
156
2/2
✓ Branch 0 taken 239064 times.
✓ Branch 1 taken 29883 times.
 268947 for (int c = 0; c < dh; ++c) acc += qr[c] * rr[c];
157 29883 relh[kh] = acc;
158 29883 }
159
2/2
✓ Branch 0 taken 30549 times.
✓ Branch 1 taken 5443 times.
 35992 for (int kw = 0; kw < grid_w; ++kw) {
160 30549 const float* rr = Rwp + static_cast<size_t>(qw - kw + grid_w - 1) * dh;
161 30549 float acc = 0.0f;
162
2/2
✓ Branch 0 taken 244392 times.
✓ Branch 1 taken 30549 times.
 274941 for (int c = 0; c < dh; ++c) acc += qr[c] * rr[c];
163 30549 relw[kw] = acc;
164 30549 }
165
166 // scores = scale * Q.K + decomposed rel-pos bias
167 5443 float row_max = -1e30f;
168
2/2
✓ Branch 0 taken 188749 times.
✓ Branch 1 taken 5443 times.
 194192 for (int j = 0; j < L; ++j) {
169 188749 const float* kr = &Kh[(static_cast<size_t>(hh) * L + j) * dh];
170 188749 float s = 0.0f;
171
2/2
✓ Branch 0 taken 1509992 times.
✓ Branch 1 taken 188749 times.
 1698741 for (int c = 0; c < dh; ++c) s += qr[c] * kr[c];
172 188749 s = s * scale + relh[j / grid_w] + relw[j % grid_w];
173 188749 srow[j] = s;
174
2/2
✓ Branch 0 taken 165953 times.
✓ Branch 1 taken 22796 times.
 188749 if (s > row_max) row_max = s;
175 188749 }
176 // softmax
177 5443 float sum = 0.0f;
178
2/2
✓ Branch 0 taken 188749 times.
✓ Branch 1 taken 5443 times.
 194192 for (int j = 0; j < L; ++j) {
179 188749 const float e = std::exp(srow[j] - row_max);
180 188749 srow[j] = e;
181 188749 sum += e;
182 188749 }
183
1/2
✓ Branch 0 taken 5443 times.
✗ Branch 1 not taken.
 5443 const float inv = sum > 0.0f ? 1.0f / sum : 0.0f;
184 // weighted sum of V
185
2/2
✓ Branch 0 taken 43544 times.
✓ Branch 1 taken 5443 times.
 48987 for (int c = 0; c < dh; ++c) {
186 43544 float acc = 0.0f;
187
2/2
✓ Branch 0 taken 1509992 times.
✓ Branch 1 taken 43544 times.
 1553536 for (int j = 0; j < L; ++j)
188 3019984 acc += srow[j] * inv *
189 1509992 Vh[(static_cast<size_t>(hh) * L + j) * dh + c];
190 43544 Yc[static_cast<size_t>(i) * D + (hh * dh + c)] = acc;
191 43544 }
192 5443 }
193 253 }
194
195 // Output projection O = Yconcat @ Wo^T (+ bo).
196
2/2
✓ Branch 0 taken 1191 times.
✓ Branch 1 taken 66 times.
 1257 for (int i = 0; i < L; ++i) {
197 1191 const float* yr = &Yc[static_cast<size_t>(i) * D];
198
2/2
✓ Branch 0 taken 43544 times.
✓ Branch 1 taken 1191 times.
 44735 for (int c = 0; c < D; ++c) {
199 43544 const float* wr = Wop + static_cast<size_t>(c) * D;
200
2/2
✓ Branch 0 taken 36976 times.
✓ Branch 1 taken 6568 times.
 43544 float acc = bop ? bop[c] : 0.0f;
201
2/2
✓ Branch 0 taken 1859904 times.
✓ Branch 1 taken 43544 times.
 1903448 for (int k = 0; k < D; ++k) acc += yr[k] * wr[k];
202 43544 Op[static_cast<size_t>(i) * D + c] = acc;
203 43544 }
204 1191 }
205 66 }
206
207 // Windowed variant: split the (grid_h, grid_w) token grid into window x window
208 // tiles (zero-padding the bottom/right up to a multiple of window, SAM's
209 // window_partition), run the decomposed-rel-pos attention above independently
210 // per tile with the SAME weights / rel-pos, then crop the padding back off.
211 // Reuses the single-grid kernel above — the only new work is the gather/scatter
212 // across windows.
213 10 void self_attention_decomposed_rel_pos_windowed_forward(
214 const ::brotensor::Tensor& X,
215 const ::brotensor::Tensor& Wq, const ::brotensor::Tensor* bq,
216 const ::brotensor::Tensor& Wk, const ::brotensor::Tensor* bk,
217 const ::brotensor::Tensor& Wv, const ::brotensor::Tensor* bv,
218 const ::brotensor::Tensor& Wo, const ::brotensor::Tensor* bo,
219 const ::brotensor::Tensor& rel_pos_h,
220 const ::brotensor::Tensor& rel_pos_w,
221 int num_heads, int grid_h, int grid_w, int window, float scale,
222 ::brotensor::Tensor& O) {
223 using ::brotensor::Tensor;
224 10 const char* fn = "self_attention_decomposed_rel_pos_windowed_forward";
225 10 check_fp32(X, "X");
226 10 const int L = X.rows;
227 10 const int D = X.cols;
228
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (window <= 0)
229 ✗ throw std::runtime_error(std::string(fn) + ": window must be >= 1");
230
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (grid_h <= 0 || grid_w <= 0 || grid_h * grid_w != L)
231 ✗ throw std::runtime_error(std::string(fn) + ": grid_h*grid_w must equal X.rows");
232
233
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 10 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 10 if (O.rows != L || O.cols != D || O.dtype != Dtype::FP32)
234 10 O.resize(L, D, Dtype::FP32);
235
2/4
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 10 times.
 10 if (L == 0 || D == 0) return;
236
237 10 const int pad_h = (window - grid_h % window) % window;
238 10 const int pad_w = (window - grid_w % window) % window;
239 10 const int nw_h = (grid_h + pad_h) / window;
240 10 const int nw_w = (grid_w + pad_w) / window;
241 10 const int ww = window * window;
242
243 10 const float* Xp = X.host_f32();
244 10 float* Op = O.host_f32_mut();
245
246 10 Tensor win_in = Tensor::mat(ww, D); // host FP32, zeroed
247 10 Tensor win_out;
248
2/2
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 10 times.
 32 for (int nh = 0; nh < nw_h; ++nh) {
249
2/2
✓ Branch 0 taken 50 times.
✓ Branch 1 taken 22 times.
 72 for (int nw = 0; nw < nw_w; ++nw) {
250 // Gather this window's tokens (token-major), zeroing padded cells.
251
1/2
✓ Branch 0 taken 50 times.
✗ Branch 1 not taken.
 50 float* wi = win_in.host_f32_mut();
252
2/2
✓ Branch 0 taken 188 times.
✓ Branch 1 taken 50 times.
 238 for (int lh = 0; lh < window; ++lh) {
253 188 const int h = nh * window + lh;
254
2/2
✓ Branch 0 taken 716 times.
✓ Branch 1 taken 188 times.
 904 for (int lw = 0; lw < window; ++lw) {
255 716 const int w = nw * window + lw;
256 716 float* dst = wi + static_cast<size_t>(lh * window + lw) * D;
257
4/4
✓ Branch 0 taken 620 times.
✓ Branch 1 taken 96 times.
✓ Branch 2 taken 74 times.
✓ Branch 3 taken 546 times.
 716 if (h < grid_h && w < grid_w) {
258 546 const float* src =
259 546 Xp + static_cast<size_t>(h * grid_w + w) * D;
260
2/2
✓ Branch 0 taken 16400 times.
✓ Branch 1 taken 546 times.
 16946 for (int c = 0; c < D; ++c) dst[c] = src[c];
261 546 } else {
262
2/2
✓ Branch 0 taken 5136 times.
✓ Branch 1 taken 170 times.
 5306 for (int c = 0; c < D; ++c) dst[c] = 0.0f;
263 }
264 716 }
265 188 }
266
267
1/2
✓ Branch 0 taken 50 times.
✗ Branch 1 not taken.
 50 self_attention_decomposed_rel_pos_forward(
268 50 win_in, Wq, bq, Wk, bk, Wv, bv, Wo, bo, rel_pos_h, rel_pos_w,
269 50 num_heads, window, window, scale, win_out);
270
271 // Scatter back, dropping the padded tokens.
272
1/2
✓ Branch 0 taken 50 times.
✗ Branch 1 not taken.
 50 const float* wo = win_out.host_f32();
273
2/2
✓ Branch 0 taken 188 times.
✓ Branch 1 taken 50 times.
 238 for (int lh = 0; lh < window; ++lh) {
274 188 const int h = nh * window + lh;
275
2/2
✓ Branch 0 taken 26 times.
✓ Branch 1 taken 162 times.
 188 if (h >= grid_h) continue;
276
2/2
✓ Branch 0 taken 620 times.
✓ Branch 1 taken 162 times.
 782 for (int lw = 0; lw < window; ++lw) {
277 620 const int w = nw * window + lw;
278
2/2
✓ Branch 0 taken 74 times.
✓ Branch 1 taken 546 times.
 620 if (w >= grid_w) continue;
279 546 const float* src =
280 546 wo + static_cast<size_t>(lh * window + lw) * D;
281 546 float* dst = Op + static_cast<size_t>(h * grid_w + w) * D;
282
2/2
✓ Branch 0 taken 16400 times.
✓ Branch 1 taken 546 times.
 16946 for (int c = 0; c < D; ++c) dst[c] = src[c];
283 546 }
284 162 }
285 50 }
286 22 }
287 10 }
288
289 } // namespace brotensor::detail::cpu
290