GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 94.2% 131 / 0 / 139
Functions: 100.0% 7 / 0 / 7
Branches: 48.8% 41 / 0 / 84
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
18 #include <cmath>
19 #include <cstddef>
20
21 namespace brotensor::detail::cpu {
22
23 22 void layernorm_forward_inference_batched(const ::brotensor::Tensor& X_RD,
24 const ::brotensor::Tensor& gamma,
25 const ::brotensor::Tensor& beta,
26 ::brotensor::Tensor& Y_RD,
27 float eps) {
28 22 const int R = X_RD.rows;
29 22 const int D = X_RD.cols;
30
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 22 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 22 if (Y_RD.rows != R || Y_RD.cols != D ||
31 ✗ Y_RD.dtype != ::brotensor::Dtype::FP32) {
32 22 Y_RD.resize(R, D, ::brotensor::Dtype::FP32);
33 22 }
34
2/4
✓ Branch 0 taken 22 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 22 times.
 22 if (R == 0 || D == 0) return;
35
36 22 const float* xp = X_RD.host_f32();
37 22 const float* gp = gamma.host_f32();
38 22 const float* bp = beta.host_f32();
39 22 float* yp = Y_RD.host_f32_mut();
40
41 22 const float invD = 1.0f / static_cast<float>(D);
42 // Each row owns Y's row exclusively (gamma/beta are read-only shared),
43 // so this parallelizes across rows with no cross-thread writes.
44
1/2
✓ Branch 0 taken 22 times.
✗ Branch 1 not taken.
 56 parallel_for(static_cast<std::size_t>(R), [&](std::size_t rowi) {
45 34 const int row = static_cast<int>(rowi);
46 34 const float* xr = xp + static_cast<std::size_t>(row) * D;
47 34 float* yr = yp + static_cast<std::size_t>(row) * D;
48
49 // Two passes: mean, then the sum of squared deviations from it.
50 // E[x^2] - E[x]^2 is one pass but cancels catastrophically once a row's
51 // mean dominates its spread — a CLIP text row whose embedding sits near
52 // 395 has both terms at ~1.6e5 in FP32, so their difference is noise and
53 // lands negative about as often as not, making rstd NaN. Summing
54 // deviations keeps every term non-negative, so var cannot go negative.
55 34 float sum = 0.0f;
56
2/2
✓ Branch 0 taken 3028 times.
✓ Branch 1 taken 34 times.
 3062 for (int i = 0; i < D; ++i) sum += xr[i];
57 34 const float mean = sum * invD;
58
59 34 float sumsq = 0.0f;
60
2/2
✓ Branch 0 taken 3028 times.
✓ Branch 1 taken 34 times.
 3062 for (int i = 0; i < D; ++i) {
61 3028 const float d = xr[i] - mean;
62 3028 sumsq += d * d;
63 3028 }
64 34 const float var = sumsq * invD;
65 34 const float rstd = 1.0f / std::sqrt(var + eps);
66
67
2/2
✓ Branch 0 taken 3028 times.
✓ Branch 1 taken 34 times.
 3062 for (int i = 0; i < D; ++i) {
68 3028 const float xh = (xr[i] - mean) * rstd;
69 3028 yr[i] = gp[i] * xh + bp[i];
70 3028 }
71 34 });
72 22 }
73
74 8 void layernorm_forward_batched_with_caches(const ::brotensor::Tensor& X_RD,
75 const ::brotensor::Tensor& gamma,
76 const ::brotensor::Tensor& beta,
77 ::brotensor::Tensor& Y_RD,
78 ::brotensor::Tensor& Xhat_RD,
79 ::brotensor::Tensor& Mean_R,
80 ::brotensor::Tensor& Rstd_R,
81 float eps) {
82 8 const int R = X_RD.rows;
83 8 const int D = X_RD.cols;
84
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 ||
85 ✗ Y_RD.dtype != ::brotensor::Dtype::FP32) {
86 8 Y_RD.resize(R, D, ::brotensor::Dtype::FP32);
87 8 }
88
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 ||
89 ✗ Xhat_RD.dtype != ::brotensor::Dtype::FP32) {
90 8 Xhat_RD.resize(R, D, ::brotensor::Dtype::FP32);
91 8 }
92
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 ||
93 ✗ Mean_R.dtype != ::brotensor::Dtype::FP32) {
94 8 Mean_R.resize(R, 1, ::brotensor::Dtype::FP32);
95 8 }
96
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 ||
97 ✗ Rstd_R.dtype != ::brotensor::Dtype::FP32) {
98 8 Rstd_R.resize(R, 1, ::brotensor::Dtype::FP32);
99 8 }
100
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;
101
102 8 const float* xp = X_RD.host_f32();
103 8 const float* gp = gamma.host_f32();
104 8 const float* bp = beta.host_f32();
105 8 float* yp = Y_RD.host_f32_mut();
106 8 float* hp = Xhat_RD.host_f32_mut();
107 8 float* mp = Mean_R.host_f32_mut();
108 8 float* sp = Rstd_R.host_f32_mut();
109
110 8 const float invD = 1.0f / static_cast<float>(D);
111 // Each row owns its own slice of Y/Xhat/Mean/Rstd exclusively, so this
112 // parallelizes across rows with no cross-thread writes.
113
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 133 parallel_for(static_cast<std::size_t>(R), [&](std::size_t rowi) {
114 125 const int row = static_cast<int>(rowi);
115 125 const float* xr = xp + static_cast<std::size_t>(row) * D;
116 125 float* yr = yp + static_cast<std::size_t>(row) * D;
117 125 float* hr = hp + static_cast<std::size_t>(row) * D;
118
119 // Two-pass variance, for the cancellation reason spelled out in
120 // layernorm_forward_inference_batched above.
121 125 float sum = 0.0f;
122
2/2
✓ Branch 0 taken 23139 times.
✓ Branch 1 taken 125 times.
 23264 for (int i = 0; i < D; ++i) sum += xr[i];
123 125 const float mean = sum * invD;
124
125 125 float sumsq = 0.0f;
126
2/2
✓ Branch 0 taken 23527 times.
✓ Branch 1 taken 125 times.
 23652 for (int i = 0; i < D; ++i) {
127 23527 const float d = xr[i] - mean;
128 23527 sumsq += d * d;
129 23527 }
130 125 const float var = sumsq * invD;
131 125 const float rstd = 1.0f / std::sqrt(var + eps);
132
133 125 mp[row] = mean;
134 125 sp[row] = rstd;
135
136
2/2
✓ Branch 0 taken 20000 times.
✓ Branch 1 taken 125 times.
 20125 for (int i = 0; i < D; ++i) {
137 20000 const float xh = (xr[i] - mean) * rstd;
138 20000 hr[i] = xh;
139 20000 yr[i] = gp[i] * xh + bp[i];
140 20000 }
141 125 });
142 8 }
143
144 8 void layernorm_backward_batched_with_caches(const ::brotensor::Tensor& dY_RD,
145 const ::brotensor::Tensor& Xhat_RD,
146 const ::brotensor::Tensor& gamma,
147 const ::brotensor::Tensor& Rstd_R,
148 ::brotensor::Tensor& dX_RD,
149 ::brotensor::Tensor& dGamma,
150 ::brotensor::Tensor& dBeta) {
151 8 const int R = dY_RD.rows;
152 8 const int D = dY_RD.cols;
153
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 ||
154 8 dX_RD.dtype != ::brotensor::Dtype::FP32) {
155 ✗ dX_RD.resize(R, D, ::brotensor::Dtype::FP32);
156 ✗ }
157
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;
158
159 8 const float* dyp = dY_RD.host_f32();
160 8 const float* hp = Xhat_RD.host_f32();
161 8 const float* gp = gamma.host_f32();
162 8 const float* sp = Rstd_R.host_f32();
163 8 float* dxp = dX_RD.host_f32_mut();
164 8 float* dgp = dGamma.host_f32_mut();
165 8 float* dbp = dBeta.host_f32_mut();
166
167 8 const float nf = static_cast<float>(D);
168
169 // dGamma/dBeta accumulate dgp[i]/dbp[i] summed over every row — a shared
170 // reduction across the row axis, not disjoint per row. So only the dX
171 // pass (fully owned by its own row) is parallelized here; dGamma/dBeta
172 // are accumulated afterwards in a separate, single-threaded pass (cheap:
173 // dyr/hr are already materialized, no recomputation needed).
174
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 132 parallel_for(static_cast<std::size_t>(R), [&](std::size_t rowi) {
175 124 const int row = static_cast<int>(rowi);
176 124 const float* dyr = dyp + static_cast<std::size_t>(row) * D;
177 124 const float* hr = hp + static_cast<std::size_t>(row) * D;
178 124 float* dxr = dxp + static_cast<std::size_t>(row) * D;
179 124 const float rstd = sp[row];
180
181 124 float sum_dxh = 0.0f;
182 124 float sum_dxh_xhat = 0.0f;
183
2/2
✓ Branch 0 taken 19182 times.
✓ Branch 1 taken 124 times.
 19306 for (int i = 0; i < D; ++i) {
184 19182 const float dxh = dyr[i] * gp[i];
185 19182 sum_dxh += dxh;
186 19182 sum_dxh_xhat += dxh * hr[i];
187 19182 }
188 124 const float scale = rstd / nf;
189
2/2
✓ Branch 0 taken 19409 times.
✓ Branch 1 taken 124 times.
 19533 for (int i = 0; i < D; ++i) {
190 19409 const float dxh = dyr[i] * gp[i];
191 19409 dxr[i] = scale * (nf * dxh - sum_dxh - hr[i] * sum_dxh_xhat);
192 19409 }
193 124 });
194
195 // Single-threaded: dGamma/dBeta reduce across rows, so they cannot be
196 // split across threads without a race on dgp[i]/dbp[i].
197
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 125 times.
 133 for (int row = 0; row < R; ++row) {
198 125 const float* dyr = dyp + static_cast<std::size_t>(row) * D;
199 125 const float* hr = hp + static_cast<std::size_t>(row) * D;
200
2/2
✓ Branch 0 taken 26193 times.
✓ Branch 1 taken 125 times.
 26318 for (int i = 0; i < D; ++i) {
201 26193 const float g = dyr[i];
202 26193 dgp[i] += g * hr[i];
203 26193 dbp[i] += g;
204 26193 }
205 125 }
206 8 }
207
208 5 void build_causal_mask_row(int L, int q, ::brotensor::Tensor& mask) {
209
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 ||
210 ✗ mask.dtype != ::brotensor::Dtype::FP32) {
211 5 mask.resize(L, 1, ::brotensor::Dtype::FP32);
212 5 }
213
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
 5 if (L <= 0) return;
214 5 float* mp = mask.host_f32_mut();
215
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;
216 5 }
217
218 } // namespace brotensor::detail::cpu
219