GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 88.5% 170 / 0 / 192
Functions: 100.0% 7 / 0 / 7
Branches: 39.3% 70 / 0 / 178
src/cpu/rope.cpp
Line Branch Exec Source
1 // ─── CPU RoPE ops (CHUNK 2) ────────────────────────────────────────────────
2 //
3 // FP32 scalar host implementations. Ports src/cuda/rope.cu — FP32 path only.
4 //
5 // Rotary position embedding: per head, rotate consecutive pairs
6 // (x_{2i}, x_{2i+1}) by angle theta = pos * theta_base^{-2i/head_dim}, where
7 // pos = row + seq_offset.
8 //
9 // X / Y layout: (L, num_heads * head_dim), row-major. Within a row, heads
10 // are contiguous head_dim-sized blocks; within a head the pairs are
11 // (2i, 2i+1).
12 // forward: Y[2i] = x0*c - x1*s ; Y[2i+1] = x0*s + x1*c
13 // backward: dX[2i] = dy0*c + dy1*s ; dX[2i+1] = -dy0*s + dy1*c (R(θ)^T)
14 //
15 // Both directions overwrite their output (the GPU kernels write directly).
16
17 #include <brotensor/tensor.h>
18
19 #include <cmath>
20 #include <stdexcept>
21 #include <string>
22 #include <vector>
23
24 namespace brotensor::detail::cpu {
25
26 namespace {
27
28 569 inline float rope_theta(int pair_i, int head_dim, float base) {
29 // theta_i = base^{-2i/head_dim} = exp(-2i/hd * log(base)).
30 1707 return std::exp(-static_cast<float>(2 * pair_i) /
31 1138 static_cast<float>(head_dim) * std::log(base));
32 }
33
34 } // namespace
35
36 9 void rope_forward(const ::brotensor::Tensor& X, int head_dim, int num_heads,
37 int seq_offset, float theta_base, ::brotensor::Tensor& Y) {
38
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (head_dim <= 0 || (head_dim & 1) != 0) {
39 ✗ throw std::runtime_error("rope_forward: head_dim must be a positive even integer");
40 }
41
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (num_heads <= 0) {
42 ✗ throw std::runtime_error("rope_forward: num_heads must be positive");
43 }
44
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (X.cols != num_heads * head_dim) {
45 ✗ throw std::runtime_error("rope_forward: X.cols != num_heads * head_dim");
46 }
47 9 const int L = X.rows;
48
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 9 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 9 if (Y.rows != L || Y.cols != X.cols) Y.resize(L, X.cols);
49 9 const int half = head_dim / 2;
50
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (L * num_heads * half == 0) return;
51 9 const int D = num_heads * head_dim;
52 9 const float* Xp = X.host_f32();
53 9 float* Yp = Y.host_f32_mut();
54 9 std::vector<float> cos_row(half);
55
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 std::vector<float> sin_row(half);
56
2/2
✓ Branch 0 taken 9 times.
✓ Branch 1 taken 71 times.
 80 for (int row = 0; row < L; ++row) {
57 71 const int pos = row + seq_offset;
58
2/2
✓ Branch 0 taken 71 times.
✓ Branch 1 taken 381 times.
 452 for (int i = 0; i < half; ++i) {
59 381 const float theta =
60
1/2
✓ Branch 0 taken 381 times.
✗ Branch 1 not taken.
 381 static_cast<float>(pos) * rope_theta(i, head_dim, theta_base);
61 381 cos_row[i] = std::cos(theta);
62 381 sin_row[i] = std::sin(theta);
63 381 }
64
2/2
✓ Branch 0 taken 186 times.
✓ Branch 1 taken 71 times.
 257 for (int h = 0; h < num_heads; ++h) {
65 186 const int base_off = row * D + h * head_dim;
66
2/2
✓ Branch 0 taken 859 times.
✓ Branch 1 taken 186 times.
 1045 for (int i = 0; i < half; ++i) {
67 859 const float c = cos_row[i];
68 859 const float s = sin_row[i];
69 859 const float x0 = Xp[base_off + 2 * i];
70 859 const float x1 = Xp[base_off + 2 * i + 1];
71 859 Yp[base_off + 2 * i] = x0 * c - x1 * s;
72 859 Yp[base_off + 2 * i + 1] = x0 * s + x1 * c;
73 859 }
74 186 }
75 71 }
76 9 }
77
78 4 void rope_backward(const ::brotensor::Tensor& dY, int head_dim, int num_heads,
79 int seq_offset, float theta_base, ::brotensor::Tensor& dX) {
80
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (head_dim <= 0 || (head_dim & 1) != 0) {
81 ✗ throw std::runtime_error("rope_backward: head_dim must be a positive even integer");
82 }
83
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (num_heads <= 0) {
84 ✗ throw std::runtime_error("rope_backward: num_heads must be positive");
85 }
86
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (dY.cols != num_heads * head_dim) {
87 ✗ throw std::runtime_error("rope_backward: dY.cols != num_heads * head_dim");
88 }
89 4 const int L = dY.rows;
90
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 4 if (dX.rows != L || dX.cols != dY.cols) dX.resize(L, dY.cols);
91 4 const int half = head_dim / 2;
92
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (L * num_heads * half == 0) return;
93 4 const int D = num_heads * head_dim;
94 4 const float* dYp = dY.host_f32();
95 4 float* dXp = dX.host_f32_mut();
96 4 std::vector<float> cos_row(half);
97
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 std::vector<float> sin_row(half);
98
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 33 times.
 37 for (int row = 0; row < L; ++row) {
99 33 const int pos = row + seq_offset;
100
2/2
✓ Branch 0 taken 33 times.
✓ Branch 1 taken 188 times.
 221 for (int i = 0; i < half; ++i) {
101 188 const float theta =
102
1/2
✓ Branch 0 taken 188 times.
✗ Branch 1 not taken.
 188 static_cast<float>(pos) * rope_theta(i, head_dim, theta_base);
103 188 cos_row[i] = std::cos(theta);
104 188 sin_row[i] = std::sin(theta);
105 188 }
106
2/2
✓ Branch 0 taken 82 times.
✓ Branch 1 taken 33 times.
 115 for (int h = 0; h < num_heads; ++h) {
107 82 const int base_off = row * D + h * head_dim;
108
2/2
✓ Branch 0 taken 408 times.
✓ Branch 1 taken 82 times.
 490 for (int i = 0; i < half; ++i) {
109 408 const float c = cos_row[i];
110 408 const float s = sin_row[i];
111 408 const float dy0 = dYp[base_off + 2 * i];
112 408 const float dy1 = dYp[base_off + 2 * i + 1];
113 // Inverse rotation (transpose of R(θ)).
114 408 dXp[base_off + 2 * i] = dy0 * c + dy1 * s;
115 408 dXp[base_off + 2 * i + 1] = -dy0 * s + dy1 * c;
116 408 }
117 82 }
118 33 }
119 4 }
120
121 // ─── RoPE with explicit cos/sin tables ─────────────────────────────────────
122 //
123 // Ports src/cuda/rope.cu — rope_apply / rope_apply_backward, FP32 path.
124 // cos_tbl / sin_tbl are (L, head_dim/2): one rotation angle per (row, pair),
125 // shared across all heads. The caller owns position semantics (this is what
126 // supports 2D axial RoPE). Both directions overwrite their output.
127
128 namespace {
129
130 10 inline void check_tables(const ::brotensor::Tensor& cos_tbl,
131 const ::brotensor::Tensor& sin_tbl,
132 const char* op, int L, int half) {
133
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (cos_tbl.dtype != Dtype::FP32 || sin_tbl.dtype != Dtype::FP32) {
134 ✗ throw std::runtime_error(std::string(op) + ": cos_tbl / sin_tbl must be FP32");
135 }
136
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (cos_tbl.size() != L * half || sin_tbl.size() != L * half) {
137 ✗ throw std::runtime_error(std::string(op) +
138 ": cos_tbl / sin_tbl must each be (L, head_dim/2)");
139 }
140 10 }
141
142 } // namespace
143
144 8 void rope_apply(const ::brotensor::Tensor& X, const ::brotensor::Tensor& cos_tbl,
145 const ::brotensor::Tensor& sin_tbl, int head_dim, int num_heads,
146 ::brotensor::Tensor& Y) {
147
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (head_dim <= 0 || (head_dim & 1) != 0) {
148 ✗ throw std::runtime_error("rope_apply: head_dim must be a positive even integer");
149 }
150
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (num_heads <= 0) {
151 ✗ throw std::runtime_error("rope_apply: num_heads must be positive");
152 }
153
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (X.dtype != Dtype::FP32) {
154 ✗ throw std::runtime_error("rope_apply: X must be FP32 (CPU backend is FP32-only)");
155 }
156
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (X.cols != num_heads * head_dim) {
157 ✗ throw std::runtime_error("rope_apply: X.cols != num_heads * head_dim");
158 }
159 8 const int L = X.rows;
160 8 const int half = head_dim / 2;
161 8 check_tables(cos_tbl, sin_tbl, "rope_apply", L, half);
162
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 (Y.rows != L || Y.cols != X.cols || Y.dtype != Dtype::FP32) {
163 8 Y.resize(L, X.cols, Dtype::FP32);
164 8 }
165
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (L * num_heads * half == 0) return;
166 8 const int D = num_heads * head_dim;
167 8 const float* Xp = X.host_f32();
168 8 const float* Cp = cos_tbl.host_f32();
169 8 const float* Sp = sin_tbl.host_f32();
170 8 float* Yp = Y.host_f32_mut();
171
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 66 times.
 74 for (int row = 0; row < L; ++row) {
172
2/2
✓ Branch 0 taken 176 times.
✓ Branch 1 taken 66 times.
 242 for (int h = 0; h < num_heads; ++h) {
173 176 const int base_off = row * D + h * head_dim;
174
2/2
✓ Branch 0 taken 816 times.
✓ Branch 1 taken 176 times.
 992 for (int i = 0; i < half; ++i) {
175 816 const float c = Cp[row * half + i];
176 816 const float s = Sp[row * half + i];
177 816 const float x0 = Xp[base_off + 2 * i];
178 816 const float x1 = Xp[base_off + 2 * i + 1];
179 816 Yp[base_off + 2 * i] = x0 * c - x1 * s;
180 816 Yp[base_off + 2 * i + 1] = x0 * s + x1 * c;
181 816 }
182 176 }
183 66 }
184 8 }
185
186 // ─── RoPE with explicit PER-HEAD cos/sin tables ───────────────────────────
187 //
188 // Like rope_apply, but cos_tbl / sin_tbl are (L*num_heads, head_dim/2): one
189 // angle per (row, head, pair). Rows are head-minor within each token, i.e. the
190 // table row for (row, h) is (row*num_heads + h). Inference-only.
191
192 2 void rope_apply_perhead(const ::brotensor::Tensor& X,
193 const ::brotensor::Tensor& cos_tbl,
194 const ::brotensor::Tensor& sin_tbl,
195 int head_dim, int num_heads, ::brotensor::Tensor& Y) {
196
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (head_dim <= 0 || (head_dim & 1) != 0) {
197 ✗ throw std::runtime_error("rope_apply_perhead: head_dim must be a positive even integer");
198 }
199
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (num_heads <= 0) {
200 ✗ throw std::runtime_error("rope_apply_perhead: num_heads must be positive");
201 }
202
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (X.dtype != Dtype::FP32) {
203 ✗ throw std::runtime_error("rope_apply_perhead: X must be FP32 (CPU backend is FP32-only)");
204 }
205
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (X.cols != num_heads * head_dim) {
206 ✗ throw std::runtime_error("rope_apply_perhead: X.cols != num_heads * head_dim");
207 }
208 2 const int L = X.rows;
209 2 const int half = head_dim / 2;
210
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (cos_tbl.dtype != Dtype::FP32 || sin_tbl.dtype != Dtype::FP32) {
211 ✗ throw std::runtime_error("rope_apply_perhead: cos_tbl / sin_tbl must be FP32");
212 }
213
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (cos_tbl.size() != L * num_heads * half || sin_tbl.size() != L * num_heads * half) {
214 ✗ throw std::runtime_error(
215 "rope_apply_perhead: cos_tbl / sin_tbl must each be (L*num_heads, head_dim/2)");
216 }
217
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 (Y.rows != L || Y.cols != X.cols || Y.dtype != Dtype::FP32) {
218 2 Y.resize(L, X.cols, Dtype::FP32);
219 2 }
220
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (L * num_heads * half == 0) return;
221 2 const int D = num_heads * head_dim;
222 2 const float* Xp = X.host_f32();
223 2 const float* Cp = cos_tbl.host_f32();
224 2 const float* Sp = sin_tbl.host_f32();
225 2 float* Yp = Y.host_f32_mut();
226
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 10 times.
 12 for (int row = 0; row < L; ++row) {
227
2/2
✓ Branch 0 taken 30 times.
✓ Branch 1 taken 10 times.
 40 for (int h = 0; h < num_heads; ++h) {
228 30 const int base_off = row * D + h * head_dim;
229 30 const int tbl_off = (row * num_heads + h) * half;
230
2/2
✓ Branch 0 taken 120 times.
✓ Branch 1 taken 30 times.
 150 for (int i = 0; i < half; ++i) {
231 120 const float c = Cp[tbl_off + i];
232 120 const float s = Sp[tbl_off + i];
233 120 const float x0 = Xp[base_off + 2 * i];
234 120 const float x1 = Xp[base_off + 2 * i + 1];
235 120 Yp[base_off + 2 * i] = x0 * c - x1 * s;
236 120 Yp[base_off + 2 * i + 1] = x0 * s + x1 * c;
237 120 }
238 30 }
239 10 }
240 2 }
241
242 2 void rope_apply_backward(const ::brotensor::Tensor& dY,
243 const ::brotensor::Tensor& cos_tbl,
244 const ::brotensor::Tensor& sin_tbl,
245 int head_dim, int num_heads, ::brotensor::Tensor& dX) {
246
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (head_dim <= 0 || (head_dim & 1) != 0) {
247 ✗ throw std::runtime_error("rope_apply_backward: head_dim must be a positive even integer");
248 }
249
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (num_heads <= 0) {
250 ✗ throw std::runtime_error("rope_apply_backward: num_heads must be positive");
251 }
252
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (dY.dtype != Dtype::FP32) {
253 ✗ throw std::runtime_error("rope_apply_backward: dY must be FP32 (CPU backend is FP32-only)");
254 }
255
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (dY.cols != num_heads * head_dim) {
256 ✗ throw std::runtime_error("rope_apply_backward: dY.cols != num_heads * head_dim");
257 }
258 2 const int L = dY.rows;
259 2 const int half = head_dim / 2;
260 2 check_tables(cos_tbl, sin_tbl, "rope_apply_backward", L, half);
261
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 (dX.rows != L || dX.cols != dY.cols || dX.dtype != Dtype::FP32) {
262 2 dX.resize(L, dY.cols, Dtype::FP32);
263 2 }
264
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (L * num_heads * half == 0) return;
265 2 const int D = num_heads * head_dim;
266 2 const float* dYp = dY.host_f32();
267 2 const float* Cp = cos_tbl.host_f32();
268 2 const float* Sp = sin_tbl.host_f32();
269 2 float* dXp = dX.host_f32_mut();
270
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 18 times.
 20 for (int row = 0; row < L; ++row) {
271
2/2
✓ Branch 0 taken 60 times.
✓ Branch 1 taken 18 times.
 78 for (int h = 0; h < num_heads; ++h) {
272 60 const int base_off = row * D + h * head_dim;
273
2/2
✓ Branch 0 taken 288 times.
✓ Branch 1 taken 60 times.
 348 for (int i = 0; i < half; ++i) {
274 288 const float c = Cp[row * half + i];
275 288 const float s = Sp[row * half + i];
276 288 const float dy0 = dYp[base_off + 2 * i];
277 288 const float dy1 = dYp[base_off + 2 * i + 1];
278 288 dXp[base_off + 2 * i] = dy0 * c + dy1 * s;
279 288 dXp[base_off + 2 * i + 1] = -dy0 * s + dy1 * c;
280 288 }
281 60 }
282 18 }
283 2 }
284
285 } // namespace brotensor::detail::cpu
286