GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 84.7% 111 / 0 / 131
Functions: 100.0% 4 / 0 / 4
Branches: 50.0% 92 / 0 / 184
src/cpu/resblock.cpp
Line Branch Exec Source
1 // CPU backend — diffusion ResBlock (CHUNK 6).
2 //
3 // Ground truth: src/cuda/resblock.cu — the FP32-weight resblock_forward and
4 // resblock_backward (NOT resblock_forward_int8w_fp16).
5 //
6 // DTYPE DECISIONS
7 // The CUDA resblock runs FP16 internally (GroupNorm/SiLU/conv2d all FP16,
8 // FP32 accumulators). The CPU backend is FP32-only (per CLAUDE.md). Both
9 // CPU impls run the FP32 scalar math; the parity test quantises inputs
10 // through FP16 so both backends start identical and compares with a loose
11 // FP16-scale tolerance (the test_resblock.cpp smoke-test envelope:
12 // atol=1e-2, rtol=1e-2). A diffusion ResBlock is a long composite chain
13 // (two GroupNorms + two 3x3 convs), so FP16 rounding compounds — a tight
14 // FP32 tolerance would not survive the FP16/FP32 cross-backend comparison.
15 //
16 // EXACT OP ORDER (verified against resblock.cu)
17 // forward:
18 // h1 = SiLU(GroupNorm(X, gamma1, beta1)) [C_in]
19 // h2 = conv2d(h1, W1, b1, 3x3 stride-1 pad-1) [C_out]
20 // if t_emb_shift: h2 += broadcast(t_emb_shift) — applied AFTER conv1,
21 // BEFORE GroupNorm2; (N, C_out) per-(n,c) or (C_out,) per-c shift.
22 // h3 = SiLU(GroupNorm(h2, gamma2, beta2)) [C_out]
23 // Y = conv2d(h3, W2, b2, 3x3 stride-1 pad-1) [C_out]
24 // skip = (Wskip ? conv2d(X, Wskip, bskip, 1x1) : X)
25 // Y += skip
26 // Wskip is required when C_in != C_out; null skip is the identity (C_in
27 // must equal C_out).
28 //
29 // ACCUMULATION (verified against resblock.cu's composition of the conv2d /
30 // group_norm backward kernels)
31 // * dX — OVERWRITTEN by the GN1 backward, then the skip
32 // path is added on top (dX += dY or dX += dX_skip).
33 // * dGamma1/dBeta1/2 — ACCUMULATE (+=) (group_norm_backward contract).
34 // * dW1/dW2/dWskip — ACCUMULATE (+=) (conv2d_backward_weight).
35 // * db1/db2/dbskip — ACCUMULATE (+=) (conv2d_backward_bias).
36 // * dt_emb_shift — ACCUMULATES (+=) (CUDA folds prev value when
37 // reducing dh2 over the spatial / batch axes).
38 // All grad buffers are caller-zeroed (or pre-seeded to test accumulation).
39
40 #include <brotensor/tensor.h>
41
42 #include <cmath>
43 #include <stdexcept>
44 #include <vector>
45
46 namespace brotensor::detail::cpu {
47
48 // CPU primitives from sibling chunk files.
49 void group_norm_forward(const ::brotensor::Tensor& X,
50 const ::brotensor::Tensor& gamma,
51 const ::brotensor::Tensor& beta,
52 int N, int C, int H, int W, int num_groups,
53 float eps, ::brotensor::Tensor& Y);
54 void group_norm_backward(const ::brotensor::Tensor& X,
55 const ::brotensor::Tensor& gamma,
56 const ::brotensor::Tensor& dY,
57 int N, int C, int H, int W, int num_groups, float eps,
58 ::brotensor::Tensor& dX, ::brotensor::Tensor& dGamma,
59 ::brotensor::Tensor& dBeta);
60 void silu_forward(const ::brotensor::Tensor& x, ::brotensor::Tensor& y);
61 void silu_backward(const ::brotensor::Tensor& x, const ::brotensor::Tensor& dY,
62 ::brotensor::Tensor& dX);
63 void conv2d_forward(const ::brotensor::Tensor& X, const ::brotensor::Tensor& Wt,
64 const ::brotensor::Tensor* bias,
65 int N, int C_in, int H, int W, int C_out, int kH, int kW,
66 int stride_h, int stride_w, int pad_h, int pad_w,
67 int dil_h, int dil_w, int groups, ::brotensor::Tensor& Y);
68 void conv2d_backward_input(const ::brotensor::Tensor& Wt,
69 const ::brotensor::Tensor& dY,
70 int N, int C_in, int H, int W,
71 int C_out, int kH, int kW,
72 int stride_h, int stride_w, int pad_h, int pad_w,
73 int dil_h, int dil_w, int groups,
74 ::brotensor::Tensor& dX);
75 void conv2d_backward_weight(const ::brotensor::Tensor& X,
76 const ::brotensor::Tensor& dY,
77 int N, int C_in, int H, int W,
78 int C_out, int kH, int kW,
79 int stride_h, int stride_w, int pad_h, int pad_w,
80 int dil_h, int dil_w, int groups,
81 ::brotensor::Tensor& dWt);
82 void conv2d_backward_bias(const ::brotensor::Tensor& dY,
83 int N, int C_out, int H_out, int W_out,
84 ::brotensor::Tensor& dB);
85 void add_inplace(::brotensor::Tensor& y, const ::brotensor::Tensor& x);
86
87 namespace {
88
89 using ::brotensor::Tensor;
90 using ::brotensor::Dtype;
91
92 // Resolve a t_emb_shift tensor's broadcast mode. Returns true and sets has_N
93 // if the shape is (N, C_out); false (per-channel) if (C_out,) / (C_out,1) /
94 // (1,C_out). Throws on any other shape — mirrors resblock.cu.
95 8 bool temb_has_N(const Tensor& t, int N, int C_out) {
96
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
 8 if (t.rows == N && t.cols == C_out) return true;
97 ✗ if ((t.rows == C_out && t.cols == 1) ||
98 ✗ (t.rows == 1 && t.cols == C_out) ||
99 ✗ t.size() == C_out)
100 ✗ return false;
101 ✗ throw std::runtime_error(
102 "resblock: t_emb_shift shape must be (N, C_out) or (C_out,)");
103 8 }
104
105 // h2[n, c, p] += t_emb_shift broadcast over the spatial axis.
106 8 void add_temb_shift(float* h2, const float* shift,
107 int N, int C_out, int spatial, bool has_N) {
108
2/2
✓ Branch 0 taken 10 times.
✓ Branch 1 taken 8 times.
 18 for (int n = 0; n < N; ++n) {
109
2/2
✓ Branch 0 taken 304 times.
✓ Branch 1 taken 10 times.
 314 for (int c = 0; c < C_out; ++c) {
110
1/2
✓ Branch 0 taken 304 times.
✗ Branch 1 not taken.
 304 const float sv = has_N ? shift[n * C_out + c] : shift[c];
111 304 float* row = h2 + (static_cast<std::size_t>(n) * C_out + c) * spatial;
112
2/2
✓ Branch 0 taken 17920 times.
✓ Branch 1 taken 304 times.
 18224 for (int p = 0; p < spatial; ++p) row[p] += sv;
113 304 }
114 10 }
115 8 }
116
117 } // namespace
118
119 // ─── resblock_forward ──────────────────────────────────────────────────────
120
121 8 void resblock_forward(const ::brotensor::Tensor& X,
122 const ::brotensor::Tensor& gamma1,
123 const ::brotensor::Tensor& beta1,
124 const ::brotensor::Tensor& W1,
125 const ::brotensor::Tensor* b1,
126 const ::brotensor::Tensor* t_emb_shift,
127 const ::brotensor::Tensor& gamma2,
128 const ::brotensor::Tensor& beta2,
129 const ::brotensor::Tensor& W2,
130 const ::brotensor::Tensor* b2,
131 const ::brotensor::Tensor* Wskip,
132 const ::brotensor::Tensor* bskip,
133 int N, int C_in, int C_out, int H, int W,
134 int num_groups, float eps,
135 ::brotensor::Tensor& Y) {
136
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (num_groups <= 0 || C_in % num_groups != 0 || C_out % num_groups != 0)
137 ✗ throw std::runtime_error("resblock_forward: num_groups must divide C_in and C_out");
138
3/4
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 8 if (Wskip == nullptr && C_in != C_out)
139 ✗ throw std::runtime_error("resblock_forward: Wskip required when C_in != C_out");
140 8 const int spatial = H * W;
141 8 const int out_cols = C_out * spatial;
142
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 != N || Y.cols != out_cols || Y.dtype != Dtype::FP32)
143 8 Y.resize(N, out_cols, Dtype::FP32);
144
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
 8 if (N == 0 || spatial == 0) return;
145
146 // Leg 1: GN1 → SiLU.
147 8 Tensor h1;
148
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 group_norm_forward(X, gamma1, beta1, N, C_in, H, W, num_groups, eps, h1);
149
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 silu_forward(h1, h1);
150
151 // Conv1: 3x3 same.
152 8 Tensor h2;
153
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 conv2d_forward(h1, W1, b1, N, C_in, H, W, C_out, 3, 3,
154 1, 1, 1, 1, 1, 1, 1, h2);
155
156 // Optional t_emb shift (after conv1, before GN2).
157
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 4 times.
 8 if (t_emb_shift) {
158
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 const bool has_N = temb_has_N(*t_emb_shift, N, C_out);
159
3/6
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 4 times.
✗ Branch 5 not taken.
 4 add_temb_shift(h2.host_f32_mut(), t_emb_shift->host_f32(),
160 4 N, C_out, spatial, has_N);
161 4 }
162
163 // Leg 2: GN2 → SiLU.
164 8 Tensor h3;
165
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 group_norm_forward(h2, gamma2, beta2, N, C_out, H, W, num_groups, eps, h3);
166
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 silu_forward(h3, h3);
167
168 // Conv2: 3x3 same → Y.
169
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 16 conv2d_forward(h3, W2, b2, N, C_out, H, W, C_out, 3, 3,
170 8 1, 1, 1, 1, 1, 1, 1, Y);
171
172 // Skip path.
173
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 4 times.
 8 if (Wskip == nullptr) {
174
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 add_inplace(Y, X);
175 4 } else {
176 4 Tensor skip;
177
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 conv2d_forward(X, *Wskip, bskip, N, C_in, H, W, C_out, 1, 1,
178 1, 1, 0, 0, 1, 1, 1, skip);
179
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 add_inplace(Y, skip);
180 4 }
181 8 }
182
183 // ─── resblock_backward ─────────────────────────────────────────────────────
184
185 8 void resblock_backward(const ::brotensor::Tensor& X,
186 const ::brotensor::Tensor& gamma1,
187 const ::brotensor::Tensor& beta1,
188 const ::brotensor::Tensor& W1,
189 const ::brotensor::Tensor* b1,
190 const ::brotensor::Tensor* t_emb_shift,
191 const ::brotensor::Tensor& gamma2,
192 const ::brotensor::Tensor& beta2,
193 const ::brotensor::Tensor& W2,
194 const ::brotensor::Tensor* b2,
195 const ::brotensor::Tensor* Wskip,
196 const ::brotensor::Tensor* bskip,
197 int N, int C_in, int C_out, int H, int W,
198 int num_groups, float eps,
199 const ::brotensor::Tensor& dY,
200 ::brotensor::Tensor& dX,
201 ::brotensor::Tensor& dGamma1,
202 ::brotensor::Tensor& dBeta1,
203 ::brotensor::Tensor& dW1,
204 ::brotensor::Tensor* db1,
205 ::brotensor::Tensor* dt_emb_shift,
206 ::brotensor::Tensor& dGamma2,
207 ::brotensor::Tensor& dBeta2,
208 ::brotensor::Tensor& dW2,
209 ::brotensor::Tensor* db2,
210 ::brotensor::Tensor* dWskip,
211 ::brotensor::Tensor* dbskip) {
212 // Forward biases b2/bskip do not affect any backward quantity (the conv2
213 // and skip-conv outputs are never recomputed in the backward, and a bias
214 // grad is a pure reduction of the downstream grad); accepted for API
215 // symmetry with the forward and resblock.cu. b1 IS used — conv1 is
216 // recomputed to obtain h2.
217 (void)b2; (void)bskip;
218
3/4
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 8 if (Wskip == nullptr && C_in != C_out)
219 ✗ throw std::runtime_error("resblock_backward: Wskip required when C_in != C_out");
220 8 const int spatial = H * W;
221
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (dY.rows != N || dY.cols != C_out * spatial)
222 ✗ throw std::runtime_error("resblock_backward: dY shape mismatch");
223
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.rows != N || dX.cols != C_in * spatial || dX.dtype != Dtype::FP32)
224 ✗ dX.resize(N, C_in * spatial, Dtype::FP32);
225
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
 8 if (N == 0 || spatial == 0) return;
226
227 // ── Recompute forward intermediates. ──
228 8 Tensor h1_pre_silu, h1;
229
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 group_norm_forward(X, gamma1, beta1, N, C_in, H, W, num_groups, eps,
230 h1_pre_silu);
231
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 silu_forward(h1_pre_silu, h1);
232
233 8 Tensor h2;
234
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 conv2d_forward(h1, W1, b1, N, C_in, H, W, C_out, 3, 3,
235 1, 1, 1, 1, 1, 1, 1, h2);
236 8 bool temb_N = false;
237
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 4 times.
 8 if (t_emb_shift) {
238
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 temb_N = temb_has_N(*t_emb_shift, N, C_out);
239
3/6
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 4 times.
✗ Branch 5 not taken.
 4 add_temb_shift(h2.host_f32_mut(), t_emb_shift->host_f32(),
240 4 N, C_out, spatial, temb_N);
241 4 }
242
243 8 Tensor h3_pre_silu, h3;
244
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 group_norm_forward(h2, gamma2, beta2, N, C_out, H, W, num_groups, eps,
245 h3_pre_silu);
246
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 silu_forward(h3_pre_silu, h3);
247
248 // ── Conv2 backward: dh3 (input grad), dW2 +=, db2 +=. ──
249 8 Tensor dh3;
250
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 conv2d_backward_input(W2, dY, N, C_out, H, W, C_out, 3, 3,
251 1, 1, 1, 1, 1, 1, 1, dh3);
252
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 16 conv2d_backward_weight(h3, dY, N, C_out, H, W, C_out, 3, 3,
253 8 1, 1, 1, 1, 1, 1, 1, dW2);
254
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 if (db2) conv2d_backward_bias(dY, N, C_out, H, W, *db2);
255
256 // ── SiLU2 backward over h3_pre_silu. ──
257 8 Tensor dh3_pre_silu;
258
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 silu_backward(h3_pre_silu, dh3, dh3_pre_silu);
259
260 // ── GN2 backward: dh2 (overwritten), dGamma2/dBeta2 accumulate. ──
261 8 Tensor dh2;
262
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 16 group_norm_backward(h2, gamma2, dh3_pre_silu, N, C_out, H, W,
263 8 num_groups, eps, dh2, dGamma2, dBeta2);
264
265 // ── t_emb_shift backward: reduce dh2 over the spatial (and, for the
266 // per-channel case, batch) axes; accumulate into dt_emb_shift. ──
267
3/4
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 8 if (t_emb_shift && dt_emb_shift) {
268
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 const float* dh2p = dh2.host_f32();
269
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 float* dtp = dt_emb_shift->host_f32_mut();
270
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (temb_N) {
271
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 4 times.
 8 for (int n = 0; n < N; ++n) {
272
2/2
✓ Branch 0 taken 48 times.
✓ Branch 1 taken 4 times.
 52 for (int c = 0; c < C_out; ++c) {
273 48 float acc = 0.0f;
274 48 const float* row =
275 96 dh2p + (static_cast<std::size_t>(n) * C_out + c) *
276 48 spatial;
277
2/2
✓ Branch 0 taken 1536 times.
✓ Branch 1 taken 48 times.
 1584 for (int p = 0; p < spatial; ++p) acc += row[p];
278 48 dtp[n * C_out + c] += acc; // accumulate
279 48 }
280 4 }
281 4 } else {
282 ✗ for (int c = 0; c < C_out; ++c) {
283 ✗ float acc = 0.0f;
284 ✗ for (int n = 0; n < N; ++n) {
285 ✗ const float* row =
286 ✗ dh2p + (static_cast<std::size_t>(n) * C_out + c) *
287 ✗ spatial;
288 ✗ for (int p = 0; p < spatial; ++p) acc += row[p];
289 ✗ }
290 ✗ dtp[c] += acc; // accumulate
291 ✗ }
292 }
293 4 }
294
295 // ── Conv1 backward: dh1, dW1 +=, db1 +=. ──
296 8 Tensor dh1;
297
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 conv2d_backward_input(W1, dh2, N, C_in, H, W, C_out, 3, 3,
298 1, 1, 1, 1, 1, 1, 1, dh1);
299
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 16 conv2d_backward_weight(h1, dh2, N, C_in, H, W, C_out, 3, 3,
300 8 1, 1, 1, 1, 1, 1, 1, dW1);
301
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 if (db1) conv2d_backward_bias(dh2, N, C_out, H, W, *db1);
302
303 // ── SiLU1 backward over h1_pre_silu. ──
304 8 Tensor dh1_pre_silu;
305
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 silu_backward(h1_pre_silu, dh1, dh1_pre_silu);
306
307 // ── GN1 backward: dX (overwritten), dGamma1/dBeta1 accumulate. ──
308
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 16 group_norm_backward(X, gamma1, dh1_pre_silu, N, C_in, H, W,
309 8 num_groups, eps, dX, dGamma1, dBeta1);
310
311 // ── Skip path backward, then sum into dX. ──
312
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 4 times.
 8 if (Wskip == nullptr) {
313
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 add_inplace(dX, dY); // identity skip: dX += dY.
314 4 } else {
315 4 Tensor dX_skip;
316
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 conv2d_backward_input(*Wskip, dY, N, C_in, H, W, C_out, 1, 1,
317 1, 1, 0, 0, 1, 1, 1, dX_skip);
318
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (dWskip)
319
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 8 conv2d_backward_weight(X, dY, N, C_in, H, W, C_out, 1, 1,
320 4 1, 1, 0, 0, 1, 1, 1, *dWskip);
321
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 if (dbskip)
322
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 conv2d_backward_bias(dY, N, C_out, H, W, *dbskip);
323
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 add_inplace(dX, dX_skip);
324 4 }
325 8 }
326
327 } // namespace brotensor::detail::cpu
328