GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 96.5% 249 / 0 / 258
Functions: 87.5% 7 / 0 / 8
Branches: 57.4% 163 / 0 / 284
src/cpu/packed_encoder.cpp
Line Branch Exec Source
1 // ─── CPU packed-encoder ops ─────────────────────────────────────────────────
2 //
3 // FP32 reference implementations of the ops a packed (variable-length,
4 // unpadded) encoder batch needs — the parity oracle for the CUDA kernels in
5 // src/cuda/flash_attention_packed.cu and src/cuda/packed_encoder.cu:
6 //
7 // flash_attention_packed_qkv_forward — bidirectional (optionally windowed)
8 // self-attention off a fused (L, 3*H*hd) QKV, per-row sequence bounds.
9 // rope_qkv_packed_inplace — RoPE on the Q/K sections, per-row pos.
10 // segment_softmax_stats — [top1, top1-top2, norm. entropy, k/255]
11 // per variable-length logit segment.
12
13 #include <brotensor/tensor.h>
14
15 #include <algorithm>
16 #include <cmath>
17 #include <cstdint>
18 #include <stdexcept>
19 #include <string>
20 #include <vector>
21
22 namespace brotensor::detail::cpu {
23
24 namespace {
25
26 ✗ [[noreturn]] void fail(const char* op, const std::string& reason) {
27 ✗ throw std::runtime_error(std::string("brotensor: ") + op + ": " + reason);
28 ✗ }
29
30 1806 void need_fp32(const ::brotensor::Tensor& t, const char* op, const char* name) {
31
1/6
✓ Branch 0 taken 1806 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 1806 if (t.dtype != Dtype::FP32) fail(op, std::string(name) + " must be FP32 (CPU backend is FP32-only)");
32 1806 }
33
34 } // namespace
35
36 1192 void flash_attention_packed_qkv_forward(const ::brotensor::Tensor& QKV,
37 const ::brotensor::Tensor& seq_bounds,
38 int num_heads, int window,
39 ::brotensor::Tensor& O) {
40 1192 constexpr const char* op = "flash_attention_packed_qkv_forward";
41 1192 need_fp32(QKV, op, "QKV");
42
1/4
✓ Branch 0 taken 1192 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 1192 if (seq_bounds.dtype != Dtype::INT32) fail(op, "seq_bounds must be INT32");
43
1/2
✓ Branch 0 taken 1192 times.
✗ Branch 1 not taken.
 1192 if (num_heads <= 0 || QKV.cols % (3 * num_heads) != 0) {
44 ✗ fail(op, "QKV.cols must be 3 * num_heads * head_dim");
45 }
46 1192 const int L = QKV.rows;
47
1/4
✓ Branch 0 taken 1192 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 1192 if (seq_bounds.rows != L || seq_bounds.cols != 2) fail(op, "seq_bounds must be (L, 2)");
48 1192 const int D = QKV.cols / 3;
49 1192 const int hd = D / num_heads;
50
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 1192 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 1192 if (O.rows != L || O.cols != D || O.dtype != Dtype::FP32) O.resize(L, D, Dtype::FP32);
51
1/2
✓ Branch 0 taken 1192 times.
✗ Branch 1 not taken.
 1192 if (L == 0) return;
52
53 1192 const float* qkv = QKV.host_f32();
54 1192 const int32_t* b = static_cast<const int32_t*>(seq_bounds.host_raw());
55 1192 float* out = O.host_f32_mut();
56 1192 const float scale = 1.0f / std::sqrt(static_cast<float>(hd));
57
2/2
✓ Branch 0 taken 679 times.
✓ Branch 1 taken 513 times.
 1192 const int hw = window > 0 ? window / 2 : -1;
58 1192 std::vector<float> s;
59
2/2
✓ Branch 0 taken 57242 times.
✓ Branch 1 taken 1192 times.
 58434 for (int r = 0; r < L; ++r) {
60 57242 const int start = b[2 * r], end = b[2 * r + 1];
61
0/4
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 57242 if (start < 0 || end > L || r < start || r >= end) fail(op, "row outside its seq_bounds");
62
3/4
✓ Branch 0 taken 36045 times.
✓ Branch 1 taken 21197 times.
✓ Branch 2 taken 36045 times.
✗ Branch 3 not taken.
 57242 const int lo = hw >= 0 ? std::max(start, r - hw) : start;
63
3/4
✓ Branch 0 taken 36045 times.
✓ Branch 1 taken 21197 times.
✓ Branch 2 taken 36045 times.
✗ Branch 3 not taken.
 57242 const int hi = hw >= 0 ? std::min(end, r + hw + 1) : end;
64
1/2
✓ Branch 0 taken 57242 times.
✗ Branch 1 not taken.
 57242 s.assign(static_cast<std::size_t>(hi - lo), 0.0f);
65
2/2
✓ Branch 0 taken 118342 times.
✓ Branch 1 taken 57242 times.
 175584 for (int h = 0; h < num_heads; ++h) {
66 118342 const float* q = qkv + static_cast<std::size_t>(r) * 3 * D + h * hd;
67 118342 float mx = -INFINITY;
68
2/2
✓ Branch 0 taken 118342 times.
✓ Branch 1 taken 3517272 times.
 3635614 for (int j = lo; j < hi; ++j) {
69 3517272 const float* k = qkv + static_cast<std::size_t>(j) * 3 * D + D + h * hd;
70 3517272 float dot = 0.0f;
71
2/2
✓ Branch 0 taken 161536824 times.
✓ Branch 1 taken 3517272 times.
 165054096 for (int d = 0; d < hd; ++d) dot += q[d] * k[d];
72 3517272 s[static_cast<std::size_t>(j - lo)] = dot * scale;
73
1/2
✓ Branch 0 taken 3517272 times.
✗ Branch 1 not taken.
 3517272 mx = std::max(mx, dot * scale);
74 3517272 }
75 118342 float sum = 0.0f;
76
2/2
✓ Branch 0 taken 3517272 times.
✓ Branch 1 taken 118342 times.
 3635614 for (float& v : s) {
77 3517272 v = std::exp(v - mx);
78 3517272 sum += v;
79 }
80 118342 float* o = out + static_cast<std::size_t>(r) * D + h * hd;
81
2/2
✓ Branch 0 taken 1631032 times.
✓ Branch 1 taken 118342 times.
 1749374 for (int d = 0; d < hd; ++d) o[d] = 0.0f;
82
2/2
✓ Branch 0 taken 3517272 times.
✓ Branch 1 taken 118342 times.
 3635614 for (int j = lo; j < hi; ++j) {
83 3517272 const float* v = qkv + static_cast<std::size_t>(j) * 3 * D + 2 * D + h * hd;
84 3517272 const float p = s[static_cast<std::size_t>(j - lo)] / sum;
85
2/2
✓ Branch 0 taken 161536824 times.
✓ Branch 1 taken 3517272 times.
 165054096 for (int d = 0; d < hd; ++d) o[d] += p * v[d];
86 3517272 }
87 118342 }
88 57242 }
89 1192 }
90
91 8 void flash_attention_packed_qkv_backward(const ::brotensor::Tensor& QKV, const ::brotensor::Tensor& dO,
92 const ::brotensor::Tensor& seq_bounds, int num_heads, int window,
93 ::brotensor::Tensor& dQKV) {
94 8 constexpr const char* op = "flash_attention_packed_qkv_backward";
95 8 need_fp32(QKV, op, "QKV");
96 8 need_fp32(dO, op, "dO");
97
1/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 8 if (seq_bounds.dtype != Dtype::INT32) fail(op, "seq_bounds must be INT32");
98
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (num_heads <= 0 || QKV.cols % (3 * num_heads) != 0) {
99 ✗ fail(op, "QKV.cols must be 3 * num_heads * head_dim");
100 }
101 8 const int L = QKV.rows;
102 8 const int D = QKV.cols / 3;
103 8 const int hd = D / num_heads;
104
1/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 8 if (seq_bounds.rows != L || seq_bounds.cols != 2) fail(op, "seq_bounds must be (L, 2)");
105
1/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 8 if (dO.rows != L || dO.cols != D) fail(op, "dO must be (L, num_heads*head_dim)");
106
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 8 if (dQKV.rows != L || dQKV.cols != 3 * D || dQKV.dtype != Dtype::FP32) dQKV.resize(L, 3 * D, Dtype::FP32);
107 8 dQKV.zero();
108
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (L == 0) return;
109
110 8 const float* qkv = QKV.host_f32();
111 8 const float* go = dO.host_f32();
112 8 const int32_t* b = static_cast<const int32_t*>(seq_bounds.host_raw());
113 8 float* g = dQKV.host_f32_mut();
114 8 const float scale = 1.0f / std::sqrt(static_cast<float>(hd));
115
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 4 times.
 8 const int hw = window > 0 ? window / 2 : -1;
116 8 std::vector<float> p;
117
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 std::vector<float> o(static_cast<std::size_t>(hd));
118
2/2
✓ Branch 0 taken 2272 times.
✓ Branch 1 taken 8 times.
 2280 for (int r = 0; r < L; ++r) {
119 2272 const int start = b[2 * r], end = b[2 * r + 1];
120
0/4
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 2272 if (start < 0 || end > L || r < start || r >= end) fail(op, "row outside its seq_bounds");
121
3/4
✓ Branch 0 taken 1435 times.
✓ Branch 1 taken 837 times.
✓ Branch 2 taken 1435 times.
✗ Branch 3 not taken.
 2272 const int lo = hw >= 0 ? std::max(start, r - hw) : start;
122
3/4
✓ Branch 0 taken 1435 times.
✓ Branch 1 taken 837 times.
✓ Branch 2 taken 1435 times.
✗ Branch 3 not taken.
 2272 const int hi = hw >= 0 ? std::min(end, r + hw + 1) : end;
123
1/2
✓ Branch 0 taken 2272 times.
✗ Branch 1 not taken.
 2272 p.assign(static_cast<std::size_t>(hi - lo), 0.0f);
124
2/2
✓ Branch 0 taken 7208 times.
✓ Branch 1 taken 2272 times.
 9480 for (int h = 0; h < num_heads; ++h) {
125 7208 const float* q = qkv + static_cast<std::size_t>(r) * 3 * D + h * hd;
126 7208 const float* dout = go + static_cast<std::size_t>(r) * D + h * hd;
127 7208 float mx = -INFINITY;
128
2/2
✓ Branch 0 taken 7208 times.
✓ Branch 1 taken 579592 times.
 586800 for (int j = lo; j < hi; ++j) {
129 579592 const float* k = qkv + static_cast<std::size_t>(j) * 3 * D + D + h * hd;
130 579592 float dot = 0.0f;
131
2/2
✓ Branch 0 taken 36700256 times.
✓ Branch 1 taken 579592 times.
 37279848 for (int d = 0; d < hd; ++d) dot += q[d] * k[d];
132 579592 p[static_cast<std::size_t>(j - lo)] = dot * scale;
133
1/2
✓ Branch 0 taken 579592 times.
✗ Branch 1 not taken.
 579592 mx = std::max(mx, dot * scale);
134 579592 }
135 7208 float sum = 0.0f;
136
2/2
✓ Branch 0 taken 579592 times.
✓ Branch 1 taken 7208 times.
 586800 for (float& v : p) {
137 579592 v = std::exp(v - mx);
138 579592 sum += v;
139 }
140
1/2
✓ Branch 0 taken 7208 times.
✗ Branch 1 not taken.
 7208 std::fill(o.begin(), o.end(), 0.0f);
141
2/2
✓ Branch 0 taken 579592 times.
✓ Branch 1 taken 7208 times.
 586800 for (int j = lo; j < hi; ++j) {
142 579592 float& pj = p[static_cast<std::size_t>(j - lo)];
143 579592 pj /= sum;
144 579592 const float* v = qkv + static_cast<std::size_t>(j) * 3 * D + 2 * D + h * hd;
145
2/2
✓ Branch 0 taken 36700256 times.
✓ Branch 1 taken 579592 times.
 37279848 for (int d = 0; d < hd; ++d) o[static_cast<std::size_t>(d)] += pj * v[d];
146 579592 }
147 7208 float drow = 0.0f; // D_r = dO_r . O_r = sum_j P_rj dP_rj
148
2/2
✓ Branch 0 taken 396384 times.
✓ Branch 1 taken 7208 times.
 403592 for (int d = 0; d < hd; ++d) drow += dout[d] * o[static_cast<std::size_t>(d)];
149 7208 float* dq = g + static_cast<std::size_t>(r) * 3 * D + h * hd;
150
2/2
✓ Branch 0 taken 579592 times.
✓ Branch 1 taken 7208 times.
 586800 for (int j = lo; j < hi; ++j) {
151 579592 const float pj = p[static_cast<std::size_t>(j - lo)];
152 579592 const float* k = qkv + static_cast<std::size_t>(j) * 3 * D + D + h * hd;
153 579592 const float* v = qkv + static_cast<std::size_t>(j) * 3 * D + 2 * D + h * hd;
154 579592 float* dk = g + static_cast<std::size_t>(j) * 3 * D + D + h * hd;
155 579592 float* dv = g + static_cast<std::size_t>(j) * 3 * D + 2 * D + h * hd;
156 579592 float dp = 0.0f;
157
2/2
✓ Branch 0 taken 36700256 times.
✓ Branch 1 taken 579592 times.
 37279848 for (int d = 0; d < hd; ++d) dp += dout[d] * v[d];
158 579592 const float ds = pj * (dp - drow) * scale;
159
2/2
✓ Branch 0 taken 36700256 times.
✓ Branch 1 taken 579592 times.
 37279848 for (int d = 0; d < hd; ++d) {
160 36700256 dv[d] += pj * dout[d];
161 36700256 dq[d] += ds * k[d];
162 36700256 dk[d] += ds * q[d];
163 36700256 }
164 579592 }
165 7208 }
166 2272 }
167 8 }
168
169 2 void rope_qkv_packed_inplace(::brotensor::Tensor& QKV, const ::brotensor::Tensor& cos_tbl,
170 const ::brotensor::Tensor& sin_tbl, const ::brotensor::Tensor& pos,
171 int num_heads, int head_dim) {
172 2 constexpr const char* op = "rope_qkv_packed_inplace";
173 2 need_fp32(QKV, op, "QKV");
174 2 need_fp32(cos_tbl, op, "cos_tbl");
175 2 need_fp32(sin_tbl, op, "sin_tbl");
176
1/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 2 if (pos.dtype != Dtype::INT32 || pos.rows != QKV.rows) fail(op, "pos must be (L, 1) INT32");
177
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (head_dim <= 0 || (head_dim & 1) || num_heads <= 0 || QKV.cols != 3 * num_heads * head_dim) {
178 ✗ fail(op, "QKV.cols must be 3 * num_heads * head_dim with head_dim even");
179 }
180 2 const int half = head_dim / 2;
181
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (cos_tbl.cols != half || sin_tbl.cols != half || sin_tbl.rows != cos_tbl.rows) {
182 ✗ fail(op, "cos_tbl / sin_tbl must be (P, head_dim/2)");
183 }
184 2 const int L = QKV.rows;
185 2 const int D = num_heads * head_dim;
186 2 float* x = QKV.host_f32_mut();
187 2 const float* ct = cos_tbl.host_f32();
188 2 const float* st = sin_tbl.host_f32();
189 2 const int32_t* p = static_cast<const int32_t*>(pos.host_raw());
190
2/2
✓ Branch 0 taken 81 times.
✓ Branch 1 taken 2 times.
 83 for (int r = 0; r < L; ++r) {
191 81 const int pr = p[r];
192
1/4
✓ Branch 0 taken 81 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 81 if (pr < 0 || pr >= cos_tbl.rows) fail(op, "pos outside the cos/sin table");
193 81 float* row = x + static_cast<std::size_t>(r) * 3 * D;
194
2/2
✓ Branch 0 taken 162 times.
✓ Branch 1 taken 81 times.
 243 for (int sec = 0; sec < 2; ++sec) {
195
2/2
✓ Branch 0 taken 1638 times.
✓ Branch 1 taken 162 times.
 1800 for (int h = 0; h < num_heads; ++h) {
196 1638 float* v = row + sec * D + h * head_dim;
197
2/2
✓ Branch 0 taken 49560 times.
✓ Branch 1 taken 1638 times.
 51198 for (int i = 0; i < half; ++i) {
198 49560 const float c = ct[static_cast<std::size_t>(pr) * half + i];
199 49560 const float s = st[static_cast<std::size_t>(pr) * half + i];
200 49560 const float x0 = v[2 * i], x1 = v[2 * i + 1];
201 49560 v[2 * i] = x0 * c - x1 * s;
202 49560 v[2 * i + 1] = x0 * s + x1 * c;
203 49560 }
204 1638 }
205 162 }
206 81 }
207 2 }
208
209 1 void segment_softmax_stats(const ::brotensor::Tensor& logits,
210 const ::brotensor::Tensor& seg_offsets,
211 ::brotensor::Tensor& out) {
212 1 constexpr const char* op = "segment_softmax_stats";
213 1 need_fp32(logits, op, "logits");
214
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 if (seg_offsets.dtype != Dtype::INT32 || seg_offsets.rows < 1) {
215 ✗ fail(op, "seg_offsets must be (S+1, 1) INT32");
216 }
217 1 const int S = seg_offsets.rows - 1;
218
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 1 time.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 1 if (out.rows != S || out.cols != 4 || out.dtype != Dtype::FP32) out.resize(S, 4, Dtype::FP32);
219 1 const float* l = logits.host_f32();
220 1 const int32_t* off = static_cast<const int32_t*>(seg_offsets.host_raw());
221 1 float* o = out.host_f32_mut();
222 1 std::vector<float> p;
223
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 1 time.
 5 for (int s = 0; s < S; ++s) {
224 4 const int a = off[s], b = off[s + 1];
225 4 float* row = o + static_cast<std::size_t>(s) * 4;
226
1/6
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 4 if (a < 0 || b < a || b > logits.size()) fail(op, "seg_offsets out of range");
227
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 3 times.
 4 if (b == a) {
228 1 row[0] = row[1] = row[2] = row[3] = 0.0f;
229 1 continue;
230 }
231 3 const int n = b - a;
232 3 float mx = -INFINITY;
233
3/4
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 14 times.
✓ Branch 2 taken 14 times.
✗ Branch 3 not taken.
 17 for (int i = a; i < b; ++i) mx = std::max(mx, l[i]);
234
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 p.assign(static_cast<std::size_t>(n), 0.0f);
235 3 float sum = 0.0f;
236
2/2
✓ Branch 0 taken 14 times.
✓ Branch 1 taken 3 times.
 17 for (int i = 0; i < n; ++i) {
237 14 p[static_cast<std::size_t>(i)] = std::exp(l[a + i] - mx);
238 14 sum += p[static_cast<std::size_t>(i)];
239 14 }
240 3 float top1 = 0.0f, top2 = 0.0f, ent = 0.0f;
241
2/2
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 14 times.
 17 for (float& v : p) {
242 14 v /= sum;
243
1/2
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
 14 ent -= v * std::log(std::max(v, 1e-9f));
244
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 10 times.
 14 if (v > top1) {
245 4 top2 = top1;
246 4 top1 = v;
247
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 2 times.
 14 } else if (v > top2) {
248 2 top2 = v;
249 2 }
250 }
251
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 const float k = std::max(2.0f, static_cast<float>(n));
252 3 row[0] = top1;
253 3 row[1] = top1 - top2;
254 3 row[2] = ent / std::log(k);
255 3 row[3] = k / 255.0f;
256 3 }
257 1 }
258
259 namespace {
260
261 22054806 float apply_act(float v, int act) {
262
5/6
✓ Branch 0 taken 293103 times.
✓ Branch 1 taken 6893772 times.
✓ Branch 2 taken 1165087 times.
✓ Branch 3 taken 13686460 times.
✓ Branch 4 taken 16384 times.
✗ Branch 5 not taken.
 22054806 switch (act) {
263
1/2
✓ Branch 0 taken 16384 times.
✗ Branch 1 not taken.
 16384 case 1: return v > 0.0f ? v : 0.0f;
264 293103 case 2: return 0.5f * v * (1.0f + std::tanh(0.7978845608f * (v + 0.044715f * v * v * v)));
265 6893772 case 3: return 0.5f * v * (1.0f + std::erf(v * 0.70710678118f));
266 1165087 case 4: return v / (1.0f + std::exp(-v));
267 ✗ case 5: return v / (1.0f + std::exp(-1.702f * v));
268 13686460 default: return v;
269 }
270 22054806 }
271
272 } // namespace
273
274 221 void linear_forward_batched_ex(const ::brotensor::Tensor& W, const ::brotensor::Tensor* bias,
275 const ::brotensor::Tensor& X, int act, int epilogue_flags,
276 ::brotensor::Tensor* /*workspace: CPU never splits*/,
277 ::brotensor::Tensor& Y) {
278 221 constexpr const char* op = "linear_forward_batched_ex";
279 221 const int epilogue = epilogue_flags & ~16; // kLinearEpiFastAccum: FP32 here regardless
280 221 need_fp32(W, op, "W");
281 221 need_fp32(X, op, "X");
282
2/2
✓ Branch 0 taken 149 times.
✓ Branch 1 taken 72 times.
 221 if (bias) need_fp32(*bias, op, "bias");
283 221 const int N = W.rows, K = W.cols, M = X.rows;
284
1/4
✓ Branch 0 taken 221 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 221 if (X.cols != K) fail(op, "X.cols must equal W.cols");
285
3/6
✓ Branch 0 taken 149 times.
✓ Branch 1 taken 72 times.
✓ Branch 2 taken 149 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 221 if (bias && static_cast<long long>(bias->rows) * bias->cols != N) fail(op, "bias size must equal W.rows");
286
1/4
✓ Branch 0 taken 221 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 221 if (epilogue < 0 || epilogue > 3) fail(op, "unknown epilogue");
287
5/8
✓ Branch 0 taken 186 times.
✓ Branch 1 taken 35 times.
✓ Branch 2 taken 33 times.
✓ Branch 3 taken 153 times.
✓ Branch 4 taken 68 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
 221 if ((epilogue == 2 || epilogue == 3) && (act != 0 || N % 2 != 0)) fail(op, "geglu/swiglu needs act 0 and even W.rows");
288
4/4
✓ Branch 0 taken 186 times.
✓ Branch 1 taken 35 times.
✓ Branch 2 taken 33 times.
✓ Branch 3 taken 153 times.
 221 const int out_cols = (epilogue == 2 || epilogue == 3) ? N / 2 : N;
289
2/2
✓ Branch 0 taken 59 times.
✓ Branch 1 taken 162 times.
 221 if (epilogue == 1) {
290
1/4
✓ Branch 0 taken 59 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 59 if (Y.rows != M || Y.cols != N || Y.dtype != Dtype::FP32) fail(op, "accumulate needs Y (B, out) FP32");
291
5/6
✓ Branch 0 taken 160 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 147 times.
✓ Branch 3 taken 13 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 147 times.
 221 } else if (Y.rows != M || Y.cols != out_cols || Y.dtype != Dtype::FP32) {
292 15 Y.resize(M, out_cols, Dtype::FP32);
293 15 }
294
1/2
✓ Branch 0 taken 221 times.
✗ Branch 1 not taken.
 221 if (M == 0) return;
295 221 const float* w = W.host_f32();
296 221 const float* x = X.host_f32();
297
2/2
✓ Branch 0 taken 149 times.
✓ Branch 1 taken 72 times.
 221 const float* bb = bias ? bias->host_f32() : nullptr;
298 221 float* y = Y.host_f32_mut();
299 221 std::vector<float> r(static_cast<std::size_t>(N));
300
2/2
✓ Branch 0 taken 20322 times.
✓ Branch 1 taken 221 times.
 20543 for (int m = 0; m < M; ++m) {
301 20322 const float* xr = x + static_cast<std::size_t>(m) * K;
302
2/2
✓ Branch 0 taken 19200857 times.
✓ Branch 1 taken 20322 times.
 19221179 for (int n = 0; n < N; ++n) {
303 19200857 const float* wr = w + static_cast<std::size_t>(n) * K;
304 19200857 float acc = 0.0f;
305
2/2
✓ Branch 0 taken 11832275629 times.
✓ Branch 1 taken 19200857 times.
 11851476486 for (int k = 0; k < K; ++k) acc += xr[k] * wr[k];
306
3/4
✓ Branch 0 taken 9854819 times.
✓ Branch 1 taken 9346038 times.
✓ Branch 2 taken 19200857 times.
✗ Branch 3 not taken.
 19200857 r[static_cast<std::size_t>(n)] = apply_act(acc + (bb ? bb[n] : 0.0f), act);
307 19200857 }
308 20322 float* yr = y + static_cast<std::size_t>(m) * out_cols;
309
2/2
✓ Branch 0 taken 11350 times.
✓ Branch 1 taken 8972 times.
 20322 if (epilogue == 0) {
310
2/2
✓ Branch 0 taken 8124783 times.
✓ Branch 1 taken 8972 times.
 8133755 for (int n = 0; n < N; ++n) yr[n] = r[static_cast<std::size_t>(n)];
311
2/2
✓ Branch 0 taken 5913 times.
✓ Branch 1 taken 5437 times.
 20322 } else if (epilogue == 1) {
312
2/2
✓ Branch 0 taken 5368176 times.
✓ Branch 1 taken 5913 times.
 5374089 for (int n = 0; n < N; ++n) yr[n] += r[static_cast<std::size_t>(n)];
313
2/2
✓ Branch 0 taken 4724 times.
✓ Branch 1 taken 713 times.
 11350 } else if (epilogue == 2) {
314
2/2
✓ Branch 0 taken 4724 times.
✓ Branch 1 taken 2661749 times.
 2666473 for (int j = 0; j < out_cols; ++j)
315
1/2
✓ Branch 0 taken 2661749 times.
✗ Branch 1 not taken.
 2661749 yr[j] = r[static_cast<std::size_t>(2 * j)] * apply_act(r[static_cast<std::size_t>(2 * j + 1)], 3);
316 4724 } else {
317
2/2
✓ Branch 0 taken 713 times.
✓ Branch 1 taken 192200 times.
 192913 for (int j = 0; j < out_cols; ++j)
318
1/2
✓ Branch 0 taken 192200 times.
✗ Branch 1 not taken.
 192200 yr[j] = apply_act(r[static_cast<std::size_t>(j)], 4) * r[static_cast<std::size_t>(out_cols + j)];
319 }
320 20322 }
321 221 }
322
323 } // namespace brotensor::detail::cpu
324