GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 96.2% 176 / 0 / 183
Functions: 100.0% 8 / 0 / 8
Branches: 38.1% 45 / 0 / 118
src/cpu/batch_norm.cpp
Line Branch Exec Source
1 // ─── CPU BatchNorm ops ─────────────────────────────────────────────────────
2 //
3 // Standard NCHW BatchNorm, FP32-only (CPU backend is FP32-only). Statistics
4 // are reduced over (N, H, W) for each channel — i.e. one mean / var per C.
5 // This is the variant pretrained ResNet / DETR-ResNet50 / classic
6 // Mask2Former backbones use; differs from GroupNorm (which reduces over
7 // a (channels_per_group, H, W) tile within a single sample).
8 //
9 // Three slots:
10 //
11 // batch_norm_forward — training. Computes per-channel batch mean/var
12 // over (N, H, W); writes Y; updates running_mean /
13 // running_var in place via momentum; saves
14 // batch mean and rstd for the backward pass.
15 // Convention:
16 // running = (1 - momentum) * running
17 // + momentum * batch_stat
18 // (PyTorch's nn.BatchNorm2d convention; the "batch"
19 // variance fed into running_var is the *unbiased*
20 // estimator — the sum of squared deviations from the
21 // channel mean over M/(M-1).)
22 //
23 // batch_norm_inference — uses running_mean / running_var; pure forward;
24 // no state mutation. This is what loaded pretrained
25 // checkpoints want during inference.
26 //
27 // batch_norm_backward — given X + saved batch mean/rstd from forward,
28 // computes dX (overwritten) plus dGamma / dBeta
29 // (accumulated; caller zeros).
30 //
31 // Reduction width M = N * H * W. Running-var uses unbiased denom (M-1) when
32 // M > 1; for M==1 it stays equal to the biased value (matches PyTorch which
33 // just uses biased when bessel correction is undefined). Forward Y and
34 // backward dX use the *biased* var (= 1/M sum (x-mean)^2) — same as PyTorch.
35
36 #include <brotensor/tensor.h>
37 #include <brotensor/detail/cpu/thread_pool.h>
38
39 #include <cmath>
40 #include <cstddef>
41 #include <stdexcept>
42 #include <string>
43 #include <vector>
44
45 namespace brotensor::detail::cpu {
46
47 namespace {
48
49 86 inline void check_fp32(const ::brotensor::Tensor& t,
50 const char* op, const char* name) {
51
1/2
✓ Branch 0 taken 86 times.
✗ Branch 1 not taken.
 86 if (t.dtype != Dtype::FP32) {
52 ✗ throw std::runtime_error(std::string(op) + ": " + name +
53 " must be FP32 (CPU backend is FP32-only)");
54 }
55 86 }
56
57 61 inline void check_per_channel(const ::brotensor::Tensor& t,
58 int C, const char* op, const char* name) {
59
1/2
✓ Branch 0 taken 61 times.
✗ Branch 1 not taken.
 61 if (t.size() != C) {
60 ✗ throw std::runtime_error(std::string(op) + ": " + name +
61 " must have C elements");
62 }
63 61 }
64
65 } // namespace
66
67 7 void batch_norm_forward(const ::brotensor::Tensor& X,
68 const ::brotensor::Tensor& gamma,
69 const ::brotensor::Tensor& beta,
70 ::brotensor::Tensor& running_mean,
71 ::brotensor::Tensor& running_var,
72 int N, int C, int H, int W,
73 float eps, float momentum,
74 ::brotensor::Tensor& Y,
75 ::brotensor::Tensor& saved_mean,
76 ::brotensor::Tensor& saved_rstd) {
77 7 check_fp32(X, "batch_norm_forward", "X");
78 7 check_fp32(gamma, "batch_norm_forward", "gamma");
79 7 check_fp32(beta, "batch_norm_forward", "beta");
80 7 check_fp32(running_mean, "batch_norm_forward", "running_mean");
81 7 check_fp32(running_var, "batch_norm_forward", "running_var");
82 7 check_per_channel(gamma, C, "batch_norm_forward", "gamma");
83 7 check_per_channel(beta, C, "batch_norm_forward", "beta");
84 7 check_per_channel(running_mean, C, "batch_norm_forward", "running_mean");
85 7 check_per_channel(running_var, C, "batch_norm_forward", "running_var");
86
87 7 const int spatial = H * W;
88 7 const int cols = C * spatial;
89
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 7 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 7 if (Y.rows != N || Y.cols != cols || Y.dtype != Dtype::FP32) {
90 7 Y.resize(N, cols, Dtype::FP32);
91 7 }
92
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 7 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 7 if (saved_mean.rows != C || saved_mean.cols != 1 ||
93 ✗ saved_mean.dtype != Dtype::FP32) {
94 7 saved_mean.resize(C, 1, Dtype::FP32);
95 7 }
96
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 7 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 7 if (saved_rstd.rows != C || saved_rstd.cols != 1 ||
97 ✗ saved_rstd.dtype != Dtype::FP32) {
98 7 saved_rstd.resize(C, 1, Dtype::FP32);
99 7 }
100
2/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 7 times.
 7 if (N == 0 || cols == 0) return;
101
102 7 const float* Xp = X.host_f32();
103 7 const float* gp = gamma.host_f32();
104 7 const float* bp = beta.host_f32();
105 7 float* rmp = running_mean.host_f32_mut();
106 7 float* rvp = running_var.host_f32_mut();
107 7 float* Yp = Y.host_f32_mut();
108 7 float* smp = saved_mean.host_f32_mut();
109 7 float* srp = saved_rstd.host_f32_mut();
110
111 7 const int M = N * spatial;
112 7 const float inv_M = 1.0f / static_cast<float>(M);
113 // Bessel correction factor for the running-var update. For M==1 we leave
114 // the biased estimate (matches PyTorch's behaviour).
115
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 1 time.
 7 const float bessel = (M > 1) ? static_cast<float>(M) /
116 6 static_cast<float>(M - 1)
117 : 1.0f;
118
119 // Each channel c owns rmp[c]/rvp[c]/smp[c]/srp[c] and Y's per-channel
120 // slice exclusively (X/gamma/beta are read-only shared, and no other c
121 // touches these locations), so this parallelizes across c with no
122 // cross-thread writes.
123
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 36 parallel_for(static_cast<std::size_t>(C), [&](std::size_t ci) {
124 29 const int c = static_cast<int>(ci);
125 // Pass 1: the channel's mean across (N, H, W), then the sum of squared
126 // deviations from it. E[x^2] - E[x]^2 is one pass but cancels
127 // catastrophically once a channel's mean dwarfs its spread — a
128 // near-constant channel puts both terms on the same large value, so
129 // their FP32 difference is noise and can come out negative, making rstd
130 // NaN. Deviations are non-negative by construction, so var cannot go
131 // negative.
132 29 float sum = 0.0f;
133
2/2
✓ Branch 0 taken 70 times.
✓ Branch 1 taken 29 times.
 99 for (int n = 0; n < N; ++n) {
134 70 const float* x_chan = Xp + (n * C + c) * spatial;
135
2/2
✓ Branch 0 taken 2545 times.
✓ Branch 1 taken 70 times.
 2615 for (int s = 0; s < spatial; ++s) sum += x_chan[s];
136 70 }
137 29 const float mean = sum * inv_M;
138
139 29 float sumsq = 0.0f;
140
2/2
✓ Branch 0 taken 67 times.
✓ Branch 1 taken 29 times.
 96 for (int n = 0; n < N; ++n) {
141 67 const float* x_chan = Xp + (n * C + c) * spatial;
142
2/2
✓ Branch 0 taken 2534 times.
✓ Branch 1 taken 67 times.
 2601 for (int s = 0; s < spatial; ++s) {
143 2534 const float d = x_chan[s] - mean;
144 2534 sumsq += d * d;
145 2534 }
146 67 }
147 29 const float var_b = sumsq * inv_M; // biased
148 29 const float rstd = 1.0f / std::sqrt(var_b + eps);
149 29 const float var_unb = var_b * bessel; // unbiased
150
151 // Save for backward.
152 29 smp[c] = mean;
153 29 srp[c] = rstd;
154
155 // Update running stats (PyTorch convention).
156 29 rmp[c] = (1.0f - momentum) * rmp[c] + momentum * mean;
157 29 rvp[c] = (1.0f - momentum) * rvp[c] + momentum * var_unb;
158
159 // Pass 2: write Y = (x - mean) * rstd * gamma + beta.
160 29 const float gv = gp[c];
161 29 const float bv = bp[c];
162
2/2
✓ Branch 0 taken 69 times.
✓ Branch 1 taken 29 times.
 98 for (int n = 0; n < N; ++n) {
163 69 const float* x_chan = Xp + (n * C + c) * spatial;
164 69 float* y_chan = Yp + (n * C + c) * spatial;
165
2/2
✓ Branch 0 taken 2567 times.
✓ Branch 1 taken 69 times.
 2636 for (int s = 0; s < spatial; ++s) {
166 2567 y_chan[s] = (x_chan[s] - mean) * rstd * gv + bv;
167 2567 }
168 69 }
169 29 });
170 7 }
171
172 6 void batch_norm_inference(const ::brotensor::Tensor& X,
173 const ::brotensor::Tensor& gamma,
174 const ::brotensor::Tensor& beta,
175 const ::brotensor::Tensor& running_mean,
176 const ::brotensor::Tensor& running_var,
177 int N, int C, int H, int W,
178 float eps,
179 ::brotensor::Tensor& Y) {
180 6 check_fp32(X, "batch_norm_inference", "X");
181 6 check_fp32(gamma, "batch_norm_inference", "gamma");
182 6 check_fp32(beta, "batch_norm_inference", "beta");
183 6 check_fp32(running_mean, "batch_norm_inference", "running_mean");
184 6 check_fp32(running_var, "batch_norm_inference", "running_var");
185 6 check_per_channel(gamma, C, "batch_norm_inference", "gamma");
186 6 check_per_channel(beta, C, "batch_norm_inference", "beta");
187 6 check_per_channel(running_mean, C, "batch_norm_inference", "running_mean");
188 6 check_per_channel(running_var, C, "batch_norm_inference", "running_var");
189
190 6 const int spatial = H * W;
191 6 const int cols = C * spatial;
192
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 6 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 6 if (Y.rows != N || Y.cols != cols || Y.dtype != Dtype::FP32) {
193 6 Y.resize(N, cols, Dtype::FP32);
194 6 }
195
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 6 times.
 6 if (N == 0 || cols == 0) return;
196
197 6 const float* Xp = X.host_f32();
198 6 const float* gp = gamma.host_f32();
199 6 const float* bp = beta.host_f32();
200 6 const float* rmp = running_mean.host_f32();
201 6 const float* rvp = running_var.host_f32();
202 6 float* Yp = Y.host_f32_mut();
203
204 // Precompute per-channel affine y = x * a_c + b_c.
205 // (x - mu) / sqrt(var + eps) * gamma + beta
206 // = x * (gamma / sqrt(var+eps)) + (beta - mu * gamma / sqrt(var+eps))
207 // Each channel c owns Y's per-channel slice exclusively, so this
208 // parallelizes across c with no cross-thread writes.
209
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 41 parallel_for(static_cast<std::size_t>(C), [&](std::size_t ci) {
210 35 const int c = static_cast<int>(ci);
211 35 const float inv = 1.0f / std::sqrt(rvp[c] + eps);
212 35 const float a = gp[c] * inv;
213 35 const float b = bp[c] - rmp[c] * a;
214
2/2
✓ Branch 0 taken 51 times.
✓ Branch 1 taken 35 times.
 86 for (int n = 0; n < N; ++n) {
215 51 const float* x_chan = Xp + (n * C + c) * spatial;
216 51 float* y_chan = Yp + (n * C + c) * spatial;
217
2/2
✓ Branch 0 taken 1944 times.
✓ Branch 1 taken 51 times.
 1995 for (int s = 0; s < spatial; ++s) {
218 1944 y_chan[s] = x_chan[s] * a + b;
219 1944 }
220 51 }
221 35 });
222 6 }
223
224 3 void batch_norm_backward(const ::brotensor::Tensor& X,
225 const ::brotensor::Tensor& gamma,
226 const ::brotensor::Tensor& saved_mean,
227 const ::brotensor::Tensor& saved_rstd,
228 const ::brotensor::Tensor& dY,
229 int N, int C, int H, int W,
230 ::brotensor::Tensor& dX,
231 ::brotensor::Tensor& dGamma,
232 ::brotensor::Tensor& dBeta) {
233 3 check_fp32(X, "batch_norm_backward", "X");
234 3 check_fp32(gamma, "batch_norm_backward", "gamma");
235 3 check_fp32(saved_mean, "batch_norm_backward", "saved_mean");
236 3 check_fp32(saved_rstd, "batch_norm_backward", "saved_rstd");
237 3 check_fp32(dY, "batch_norm_backward", "dY");
238 3 check_fp32(dGamma, "batch_norm_backward", "dGamma");
239 3 check_fp32(dBeta, "batch_norm_backward", "dBeta");
240 3 check_per_channel(gamma, C, "batch_norm_backward", "gamma");
241 3 check_per_channel(saved_mean, C, "batch_norm_backward", "saved_mean");
242 3 check_per_channel(saved_rstd, C, "batch_norm_backward", "saved_rstd");
243
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 6 if (dGamma.rows != C || dGamma.cols != 1 ||
244 3 dBeta.rows != C || dBeta.cols != 1) {
245 ✗ throw std::runtime_error("batch_norm_backward: dGamma/dBeta must be (C,1)");
246 }
247
248 3 const int spatial = H * W;
249 3 const int cols = C * spatial;
250
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (dY.rows != N || dY.cols != cols) {
251 ✗ throw std::runtime_error("batch_norm_backward: dY shape mismatch");
252 }
253
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (X.rows != N || X.cols != cols) {
254 ✗ throw std::runtime_error("batch_norm_backward: X shape mismatch");
255 }
256
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 3 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 3 if (dX.rows != N || dX.cols != cols || dX.dtype != Dtype::FP32) {
257 3 dX.resize(N, cols, Dtype::FP32);
258 3 }
259
2/4
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 3 times.
 3 if (N == 0 || cols == 0) return;
260
261 3 const float* Xp = X.host_f32();
262 3 const float* gp = gamma.host_f32();
263 3 const float* mp = saved_mean.host_f32();
264 3 const float* rp = saved_rstd.host_f32();
265 3 const float* dYp = dY.host_f32();
266 3 float* dXp = dX.host_f32_mut();
267 3 float* dGp = dGamma.host_f32_mut();
268 3 float* dBp = dBeta.host_f32_mut();
269
270 3 const int M = N * spatial;
271 3 const float inv_M = 1.0f / static_cast<float>(M);
272
273 // Per-channel: derive dGamma, dBeta, and the two reduction sums used
274 // by the dX formula. Then a second pass over (N, H, W) writes dX.
275 //
276 // xhat = (x - mean) * rstd
277 // dxhat = dY * gamma
278 // dGamma_c += sum (dY * xhat)
279 // dBeta_c += sum dY
280 // dX = rstd * (dxhat - (sum dxhat + xhat * sum(dxhat*xhat)) / M)
281 //
282 // The outer axis here is already C (channel), and — unlike GroupNorm's
283 // batch axis — each channel c owns dGp[c]/dBp[c] and dX's per-channel
284 // slice exclusively: no other c ever touches these locations, so the
285 // per-channel reduction over (N, H, W) is entirely local to this c's
286 // work and safe to run concurrently with other c's. The one thing that
287 // must change to parallelize: xhat_buf was a single buffer shared and
288 // reused across every c in the original serial loop, so it's now
289 // declared LOCAL to the lambda (one per c) instead — safe even across
290 // threads since each invocation gets its own buffer.
291
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 19 parallel_for(static_cast<std::size_t>(C), [&](std::size_t ci) {
292 16 const int c = static_cast<int>(ci);
293 16 std::vector<float> xhat_buf(static_cast<std::size_t>(M));
294
295 16 const float mean = mp[c];
296 16 const float rstd = rp[c];
297 16 const float gv = gp[c];
298
299 16 float sum_dY = 0.0f;
300 16 float sum_dY_xh = 0.0f;
301
2/2
✓ Branch 0 taken 43 times.
✓ Branch 1 taken 16 times.
 59 for (int n = 0; n < N; ++n) {
302 43 const float* x_chan = Xp + (n * C + c) * spatial;
303 43 const float* dy_chan = dYp + (n * C + c) * spatial;
304 43 float* xh_chan = xhat_buf.data() + n * spatial;
305
2/2
✓ Branch 0 taken 1853 times.
✓ Branch 1 taken 43 times.
 1896 for (int s = 0; s < spatial; ++s) {
306 1853 const float xh = (x_chan[s] - mean) * rstd;
307 1853 xh_chan[s] = xh;
308 1853 sum_dY += dy_chan[s];
309 1853 sum_dY_xh += dy_chan[s] * xh;
310 1853 }
311 43 }
312
313 16 dGp[c] += sum_dY_xh; // accumulate
314 16 dBp[c] += sum_dY; // accumulate
315
316 // For dX we need sum1 = sum dxhat = gv * sum_dY,
317 // and sum2 = sum (dxhat * xhat) = gv * sum_dY_xh.
318 16 const float sum1 = gv * sum_dY;
319 16 const float sum2 = gv * sum_dY_xh;
320
321
2/2
✓ Branch 0 taken 43 times.
✓ Branch 1 taken 16 times.
 59 for (int n = 0; n < N; ++n) {
322 43 const float* dy_chan = dYp + (n * C + c) * spatial;
323 43 float* dx_chan = dXp + (n * C + c) * spatial;
324 43 const float* xh_chan = xhat_buf.data() + n * spatial;
325
2/2
✓ Branch 0 taken 1841 times.
✓ Branch 1 taken 43 times.
 1884 for (int s = 0; s < spatial; ++s) {
326 1841 const float xh = xh_chan[s];
327 1841 const float dxh = dy_chan[s] * gv;
328 1841 dx_chan[s] = rstd * (dxh - (sum1 + xh * sum2) * inv_M);
329 1841 }
330 43 }
331 16 });
332 3 }
333
334 } // namespace brotensor::detail::cpu
335