GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 95.5% 169 / 0 / 177
Functions: 100.0% 9 / 0 / 9
Branches: 57.8% 67 / 0 / 116
src/cpu/layernorm_inference.cpp
Line Branch Exec Source
1 // ─── CPU inference-batched layernorm + causal-mask helper (CHUNK 1) ────────
2 //
3 // FP32 scalar host implementations.
4 //
5 // layernorm_forward_inference_batched — per-row LayerNorm over an (R, D)
6 // tensor. Ports src/cuda/layernorm.cu's
7 // layernorm_forward_inference_batched_kernel: mean/variance per row,
8 // rstd = 1/sqrt(var + eps), y = gamma * xhat + beta. FP32 only.
9 //
10 // build_causal_mask_row — fills an (L, 1) mask: mask[k] = (k <= q) ? 1 : 0.
11 // Ports src/cuda/elementwise.cu's causal_mask_row_kernel. Its only
12 // Tensor operand is the output; the dispatcher resolves the device
13 // from `mask` itself (see src/ops.cpp::build_causal_mask_row).
14
15 #include <brotensor/tensor.h>
16 #include <brotensor/detail/cpu/thread_pool.h>
17 #include "cpu_jit.h"
18
19 #include <cmath>
20 #include <cstddef>
21
22 namespace brotensor::detail::cpu {
23
24 271 void layernorm_forward_inference_batched(const ::brotensor::Tensor& X_RD,
25 const ::brotensor::Tensor& gamma,
26 const ::brotensor::Tensor* beta,
27 ::brotensor::Tensor& Y_RD,
28 float eps) {
29 271 const int R = X_RD.rows;
30 271 const int D = X_RD.cols;
31
4/6
✓ Branch 0 taken 240 times.
✓ Branch 1 taken 31 times.
✓ Branch 2 taken 240 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 240 times.
 271 if (Y_RD.rows != R || Y_RD.cols != D ||
32 240 Y_RD.dtype != ::brotensor::Dtype::FP32) {
33 31 Y_RD.resize(R, D, ::brotensor::Dtype::FP32);
34 31 }
35
2/4
✓ Branch 0 taken 271 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 271 times.
 271 if (R == 0 || D == 0) return;
36
37 271 const float* xp = X_RD.host_f32();
38 271 const float* gp = gamma.host_f32();
39
2/2
✓ Branch 0 taken 267 times.
✓ Branch 1 taken 4 times.
 271 const float* bp = beta ? beta->host_f32() : nullptr;
40 271 float* yp = Y_RD.host_f32_mut();
41
42
5/6
✓ Branch 0 taken 267 times.
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 267 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 1 time.
✓ Branch 5 taken 266 times.
 271 if (bp && jit::is_jit_available() && D >= 8) {
43 266 jit::layernorm_forward_inference_batched(xp, gp, bp, eps, yp, R, D);
44 266 return;
45 }
46
47 5 const float invD = 1.0f / static_cast<float>(D);
48 // Each row owns Y's row exclusively (gamma/beta are read-only shared),
49 // so this parallelizes across rows with no cross-thread writes.
50
1/2
✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
 417 parallel_for(static_cast<std::size_t>(R), [&](std::size_t rowi) {
51 412 const int row = static_cast<int>(rowi);
52 412 const float* xr = xp + static_cast<std::size_t>(row) * D;
53 412 float* yr = yp + static_cast<std::size_t>(row) * D;
54
55 // Two passes: mean, then the sum of squared deviations from it.
56 // E[x^2] - E[x]^2 is one pass but cancels catastrophically once a row's
57 // mean dominates its spread — a CLIP text row whose embedding sits near
58 // 395 has both terms at ~1.6e5 in FP32, so their difference is noise and
59 // lands negative about as often as not, making rstd NaN. Summing
60 // deviations keeps every term non-negative, so var cannot go negative.
61 412 float sum = 0.0f;
62
2/2
✓ Branch 0 taken 294545 times.
✓ Branch 1 taken 412 times.
 294957 for (int i = 0; i < D; ++i) sum += xr[i];
63 412 const float mean = sum * invD;
64
65 412 float sumsq = 0.0f;
66
2/2
✓ Branch 0 taken 286453 times.
✓ Branch 1 taken 412 times.
 286865 for (int i = 0; i < D; ++i) {
67 286453 const float d = xr[i] - mean;
68 286453 sumsq += d * d;
69 286453 }
70 412 const float var = sumsq * invD;
71 412 const float rstd = 1.0f / std::sqrt(var + eps);
72
73
2/2
✓ Branch 0 taken 286447 times.
✓ Branch 1 taken 412 times.
 286859 for (int i = 0; i < D; ++i) {
74 286447 const float xh = (xr[i] - mean) * rstd;
75
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 286446 times.
 286447 yr[i] = bp ? (gp[i] * xh + bp[i]) : (gp[i] * xh);
76 286447 }
77 412 });
78 271 }
79
80 2 void layernorm_forward_inference_batched_fp16(const ::brotensor::Tensor& X_RD,
81 const ::brotensor::Tensor& gamma,
82 const ::brotensor::Tensor* beta,
83 ::brotensor::Tensor& Y_RD,
84 float eps) {
85 2 const int R = X_RD.rows;
86 2 const int D = X_RD.cols;
87
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_RD.rows != R || Y_RD.cols != D ||
88 ✗ Y_RD.dtype != ::brotensor::Dtype::FP16) {
89 2 Y_RD.resize(R, D, ::brotensor::Dtype::FP16);
90 2 }
91
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 if (R == 0 || D == 0) return;
92
93 2 const uint16_t* xp = X_RD.host_fp16();
94 2 const uint16_t* gp = gamma.host_fp16();
95
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 1 time.
 2 const uint16_t* bp = beta ? beta->host_fp16() : nullptr;
96 2 uint16_t* yp = Y_RD.host_fp16_mut();
97
98 2 const float invD = 1.0f / static_cast<float>(D);
99
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 18 parallel_for(static_cast<std::size_t>(R), [&](std::size_t rowi) {
100 16 const int row = static_cast<int>(rowi);
101 16 const uint16_t* xr = xp + static_cast<std::size_t>(row) * D;
102 16 uint16_t* yr = yp + static_cast<std::size_t>(row) * D;
103
104 16 float sum = 0.0f;
105
2/2
✓ Branch 0 taken 781 times.
✓ Branch 1 taken 16 times.
 797 for (int i = 0; i < D; ++i) sum += ::brotensor::fp16_bits_to_fp32(xr[i]);
106 16 const float mean = sum * invD;
107
108 16 float sumsq = 0.0f;
109
2/2
✓ Branch 0 taken 769 times.
✓ Branch 1 taken 16 times.
 785 for (int i = 0; i < D; ++i) {
110 769 const float d = ::brotensor::fp16_bits_to_fp32(xr[i]) - mean;
111 769 sumsq += d * d;
112 769 }
113 16 const float var = sumsq * invD;
114 16 const float rstd = 1.0f / std::sqrt(var + eps);
115
116
2/2
✓ Branch 0 taken 795 times.
✓ Branch 1 taken 16 times.
 811 for (int i = 0; i < D; ++i) {
117 795 const float xh = (::brotensor::fp16_bits_to_fp32(xr[i]) - mean) * rstd;
118 795 const float g = ::brotensor::fp16_bits_to_fp32(gp[i]);
119
2/2
✓ Branch 0 taken 283 times.
✓ Branch 1 taken 512 times.
 795 const float b = bp ? ::brotensor::fp16_bits_to_fp32(bp[i]) : 0.0f;
120 795 yr[i] = ::brotensor::fp32_to_fp16_bits(g * xh + b);
121 795 }
122 16 });
123 2 }
124
125 8 void layernorm_forward_batched_with_caches(const ::brotensor::Tensor& X_RD,
126 const ::brotensor::Tensor& gamma,
127 const ::brotensor::Tensor& beta,
128 ::brotensor::Tensor& Y_RD,
129 ::brotensor::Tensor& Xhat_RD,
130 ::brotensor::Tensor& Mean_R,
131 ::brotensor::Tensor& Rstd_R,
132 float eps) {
133 8 const int R = X_RD.rows;
134 8 const int D = X_RD.cols;
135
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_RD.rows != R || Y_RD.cols != D ||
136 ✗ Y_RD.dtype != ::brotensor::Dtype::FP32) {
137 8 Y_RD.resize(R, D, ::brotensor::Dtype::FP32);
138 8 }
139
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 (Xhat_RD.rows != R || Xhat_RD.cols != D ||
140 ✗ Xhat_RD.dtype != ::brotensor::Dtype::FP32) {
141 8 Xhat_RD.resize(R, D, ::brotensor::Dtype::FP32);
142 8 }
143
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 (Mean_R.rows != R || Mean_R.cols != 1 ||
144 ✗ Mean_R.dtype != ::brotensor::Dtype::FP32) {
145 8 Mean_R.resize(R, 1, ::brotensor::Dtype::FP32);
146 8 }
147
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 (Rstd_R.rows != R || Rstd_R.cols != 1 ||
148 ✗ Rstd_R.dtype != ::brotensor::Dtype::FP32) {
149 8 Rstd_R.resize(R, 1, ::brotensor::Dtype::FP32);
150 8 }
151
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
 8 if (R == 0 || D == 0) return;
152
153 8 const float* xp = X_RD.host_f32();
154 8 const float* gp = gamma.host_f32();
155 8 const float* bp = beta.host_f32();
156 8 float* yp = Y_RD.host_f32_mut();
157 8 float* hp = Xhat_RD.host_f32_mut();
158 8 float* mp = Mean_R.host_f32_mut();
159 8 float* sp = Rstd_R.host_f32_mut();
160
161 8 const float invD = 1.0f / static_cast<float>(D);
162 // Each row owns its own slice of Y/Xhat/Mean/Rstd exclusively, so this
163 // parallelizes across rows with no cross-thread writes.
164
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 132 parallel_for(static_cast<std::size_t>(R), [&](std::size_t rowi) {
165 124 const int row = static_cast<int>(rowi);
166 124 const float* xr = xp + static_cast<std::size_t>(row) * D;
167 124 float* yr = yp + static_cast<std::size_t>(row) * D;
168 124 float* hr = hp + static_cast<std::size_t>(row) * D;
169
170 // Two-pass variance, for the cancellation reason spelled out in
171 // layernorm_forward_inference_batched above.
172 124 float sum = 0.0f;
173
2/2
✓ Branch 0 taken 20409 times.
✓ Branch 1 taken 124 times.
 20533 for (int i = 0; i < D; ++i) sum += xr[i];
174 124 const float mean = sum * invD;
175
176 124 float sumsq = 0.0f;
177
2/2
✓ Branch 0 taken 21109 times.
✓ Branch 1 taken 124 times.
 21233 for (int i = 0; i < D; ++i) {
178 21109 const float d = xr[i] - mean;
179 21109 sumsq += d * d;
180 21109 }
181 124 const float var = sumsq * invD;
182 124 const float rstd = 1.0f / std::sqrt(var + eps);
183
184 124 mp[row] = mean;
185 124 sp[row] = rstd;
186
187
2/2
✓ Branch 0 taken 17491 times.
✓ Branch 1 taken 124 times.
 17615 for (int i = 0; i < D; ++i) {
188 17491 const float xh = (xr[i] - mean) * rstd;
189 17491 hr[i] = xh;
190 17491 yr[i] = gp[i] * xh + bp[i];
191 17491 }
192 124 });
193 8 }
194
195 8 void layernorm_backward_batched_with_caches(const ::brotensor::Tensor& dY_RD,
196 const ::brotensor::Tensor& Xhat_RD,
197 const ::brotensor::Tensor& gamma,
198 const ::brotensor::Tensor& Rstd_R,
199 ::brotensor::Tensor& dX_RD,
200 ::brotensor::Tensor& dGamma,
201 ::brotensor::Tensor& dBeta) {
202 8 const int R = dY_RD.rows;
203 8 const int D = dY_RD.cols;
204
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 (dX_RD.rows != R || dX_RD.cols != D ||
205 8 dX_RD.dtype != ::brotensor::Dtype::FP32) {
206 ✗ dX_RD.resize(R, D, ::brotensor::Dtype::FP32);
207 ✗ }
208
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
 8 if (R == 0 || D == 0) return;
209
210 8 const float* dyp = dY_RD.host_f32();
211 8 const float* hp = Xhat_RD.host_f32();
212 8 const float* gp = gamma.host_f32();
213 8 const float* sp = Rstd_R.host_f32();
214 8 float* dxp = dX_RD.host_f32_mut();
215 8 float* dgp = dGamma.host_f32_mut();
216 8 float* dbp = dBeta.host_f32_mut();
217
218 8 const float nf = static_cast<float>(D);
219
220 // dGamma/dBeta accumulate dgp[i]/dbp[i] summed over every row — a shared
221 // reduction across the row axis, not disjoint per row. So only the dX
222 // pass (fully owned by its own row) is parallelized here; dGamma/dBeta
223 // are accumulated afterwards in a separate, single-threaded pass (cheap:
224 // dyr/hr are already materialized, no recomputation needed).
225
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 133 parallel_for(static_cast<std::size_t>(R), [&](std::size_t rowi) {
226 125 const int row = static_cast<int>(rowi);
227 125 const float* dyr = dyp + static_cast<std::size_t>(row) * D;
228 125 const float* hr = hp + static_cast<std::size_t>(row) * D;
229 125 float* dxr = dxp + static_cast<std::size_t>(row) * D;
230 125 const float rstd = sp[row];
231
232 125 float sum_dxh = 0.0f;
233 125 float sum_dxh_xhat = 0.0f;
234
2/2
✓ Branch 0 taken 15064 times.
✓ Branch 1 taken 125 times.
 15189 for (int i = 0; i < D; ++i) {
235 15064 const float dxh = dyr[i] * gp[i];
236 15064 sum_dxh += dxh;
237 15064 sum_dxh_xhat += dxh * hr[i];
238 15064 }
239 125 const float scale = rstd / nf;
240
2/2
✓ Branch 0 taken 13579 times.
✓ Branch 1 taken 125 times.
 13704 for (int i = 0; i < D; ++i) {
241 13579 const float dxh = dyr[i] * gp[i];
242 13579 dxr[i] = scale * (nf * dxh - sum_dxh - hr[i] * sum_dxh_xhat);
243 13579 }
244 125 });
245
246 // Single-threaded: dGamma/dBeta reduce across rows, so they cannot be
247 // split across threads without a race on dgp[i]/dbp[i].
248
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 125 times.
 133 for (int row = 0; row < R; ++row) {
249 125 const float* dyr = dyp + static_cast<std::size_t>(row) * D;
250 125 const float* hr = hp + static_cast<std::size_t>(row) * D;
251
2/2
✓ Branch 0 taken 26193 times.
✓ Branch 1 taken 125 times.
 26318 for (int i = 0; i < D; ++i) {
252 26193 const float g = dyr[i];
253 26193 dgp[i] += g * hr[i];
254 26193 dbp[i] += g;
255 26193 }
256 125 }
257 8 }
258
259 5 void build_causal_mask_row(int L, int q, ::brotensor::Tensor& mask) {
260
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 5 if (mask.rows != L || mask.cols != 1 ||
261 ✗ mask.dtype != ::brotensor::Dtype::FP32) {
262 5 mask.resize(L, 1, ::brotensor::Dtype::FP32);
263 5 }
264
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
 5 if (L <= 0) return;
265 5 float* mp = mask.host_f32_mut();
266
2/2
✓ Branch 0 taken 5 times.
✓ Branch 1 taken 325 times.
 330 for (int k = 0; k < L; ++k) mp[k] = (k <= q) ? 1.0f : 0.0f;
267 5 }
268
269 } // namespace brotensor::detail::cpu
270