GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 94.8% 145 / 0 / 153
Functions: 100.0% 5 / 0 / 5
Branches: 50.6% 79 / 0 / 156
src/cpu/modulated_conv2d.cpp
Line Branch Exec Source
1 // ─── CPU modulated_conv2d (StyleGAN3) ───────────────────────────────────────
2 //
3 // The StyleGAN synthesis-layer core: per-sample style modulation of the conv
4 // weights, optional demodulation, then a standard stride-1 conv per sample.
5 // FP32 reference mirroring NVlabs `modulated_conv2d` (the pure-PyTorch path).
6 //
7 // Realized by looping the batch and reusing the validated CPU conv2d kernels
8 // (groups=1) on a per-sample weight — only the weight construction + demod is
9 // new. Layouts: X (N,C_in*H*W) NCHW; W (C_out, C_in*kH*kW) OIHW; s (N,C_in).
10 //
11 // w'[o,i,kh,kw] = W[o,i,kh,kw] * s[n,i]
12 // dcoef[n,o] = demodulate ? rsqrt(Σ_{i,kh,kw} w'^2 + eps) : 1
13 // w'' = w' * dcoef[n,o]
14 // Y[n] = conv2d(X[n], w'', pad, stride=1)
15 //
16 // Backward (per n; dw'' = conv2d_backward_weight(X[n],dY[n])):
17 // g[o] = Σ dw''[o,..] * w'[o,..]
18 // dw'[o] = demodulate ? dw''[o]*dcoef - g[o]*dcoef^3*w'[o] : dw''[o]
19 // dW[o] += Σ_n dw'[n,o] * s[n,i] (accumulate — caller zeros dW)
20 // ds[n,i] = Σ_{o,kh,kw} dw'[o,i,kh,kw] * W[o,i,kh,kw] (overwrite)
21 // dX[n] = conv2d_backward_input(w''[n], dY[n]) (overwrite)
22
23 #include <brotensor/tensor.h>
24
25 #include <cmath>
26 #include <stdexcept>
27 #include <string>
28 #include <vector>
29
30 namespace brotensor::detail::cpu {
31
32 // Reused CPU conv2d kernels (defined in conv2d.cpp, same namespace).
33 void conv2d_forward(const ::brotensor::Tensor& X, const ::brotensor::Tensor& Wt,
34 const ::brotensor::Tensor* bias,
35 int N, int C_in, int H, int W, int C_out, int kH, int kW,
36 int stride_h, int stride_w, int pad_h, int pad_w,
37 int dil_h, int dil_w, int groups, ::brotensor::Tensor& Y);
38 void conv2d_backward_input(const ::brotensor::Tensor& Wt,
39 const ::brotensor::Tensor& dY,
40 int N, int C_in, int H, int W,
41 int C_out, int kH, int kW,
42 int stride_h, int stride_w, int pad_h, int pad_w,
43 int dil_h, int dil_w, int groups,
44 ::brotensor::Tensor& dX);
45 void conv2d_backward_weight(const ::brotensor::Tensor& X,
46 const ::brotensor::Tensor& dY,
47 int N, int C_in, int H, int W,
48 int C_out, int kH, int kW,
49 int stride_h, int stride_w, int pad_h, int pad_w,
50 int dil_h, int dil_w, int groups,
51 ::brotensor::Tensor& dWt);
52
53 namespace {
54
55 92 inline void check_fp32(const ::brotensor::Tensor& t,
56 const char* op, const char* name) {
57
1/2
✓ Branch 0 taken 92 times.
✗ Branch 1 not taken.
 92 if (t.dtype != Dtype::FP32) {
58 ✗ throw std::runtime_error(std::string("brotensor: ") + op + ": " +
59 ✗ name + " must be FP32 (CPU backend is "
60 "FP32-only)");
61 }
62 92 }
63
64 // Non-owning CPU FP32 view over one row of a (N, cols) tensor.
65 64 inline ::brotensor::Tensor row_view(const ::brotensor::Tensor& T, int n, int cols) {
66 64 float* base = const_cast<float*>(T.host_f32()) + static_cast<size_t>(n) * cols;
67 64 return ::brotensor::Tensor::view(Device::CPU, base, 1, cols, Dtype::FP32);
68 }
69 40 inline ::brotensor::Tensor row_view_mut(::brotensor::Tensor& T, int n, int cols) {
70 40 float* base = T.host_f32_mut() + static_cast<size_t>(n) * cols;
71 40 return ::brotensor::Tensor::view(Device::CPU, base, 1, cols, Dtype::FP32);
72 }
73
74 } // namespace
75
76 8 void modulated_conv2d_forward(const ::brotensor::Tensor& X,
77 const ::brotensor::Tensor& W,
78 const ::brotensor::Tensor& s,
79 int N, int C_in, int H, int Wd,
80 int C_out, int kH, int kW,
81 int pad_h, int pad_w,
82 bool demodulate, float eps,
83 ::brotensor::Tensor& dcoef,
84 ::brotensor::Tensor& Y) {
85 8 check_fp32(X, "modulated_conv2d_forward", "X");
86 8 check_fp32(W, "modulated_conv2d_forward", "W");
87 8 check_fp32(s, "modulated_conv2d_forward", "s");
88 8 const int wk = C_in * kH * kW;
89
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (X.rows != N || X.cols != C_in * H * Wd)
90 ✗ throw std::runtime_error("modulated_conv2d_forward: X shape mismatch");
91
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (W.rows != C_out || W.cols != wk)
92 ✗ throw std::runtime_error("modulated_conv2d_forward: W shape mismatch");
93
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (s.rows != N || s.cols != C_in)
94 ✗ throw std::runtime_error("modulated_conv2d_forward: s shape mismatch");
95 8 const int H_out = H + 2 * pad_h - (kH - 1);
96 8 const int W_out = Wd + 2 * pad_w - (kW - 1);
97
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (H_out <= 0 || W_out <= 0)
98 ✗ throw std::runtime_error("modulated_conv2d_forward: non-positive output shape");
99
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 (dcoef.rows != N || dcoef.cols != C_out || dcoef.dtype != Dtype::FP32)
100 8 dcoef.resize(N, C_out, Dtype::FP32);
101 8 const int out_cols = C_out * H_out * W_out;
102
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)
103 8 Y.resize(N, out_cols, Dtype::FP32);
104
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
 8 if (N == 0 || out_cols == 0) return;
105
106 8 const float* Wp = W.host_f32();
107 8 const float* sp = s.host_f32();
108 8 float* dcp = dcoef.host_f32_mut();
109
110 8 ::brotensor::Tensor Wn = ::brotensor::Tensor::zeros_on(Device::CPU, C_out, wk);
111
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 float* Wnp = Wn.host_f32_mut();
112
113
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 8 times.
 24 for (int n = 0; n < N; ++n) {
114 16 const float* sn = sp + static_cast<size_t>(n) * C_in;
115
2/2
✓ Branch 0 taken 64 times.
✓ Branch 1 taken 16 times.
 80 for (int o = 0; o < C_out; ++o) {
116 64 const float* Wo = Wp + static_cast<size_t>(o) * wk;
117 64 float* Wno = Wnp + static_cast<size_t>(o) * wk;
118 // w'[o,..] = W[o,..]*s[n,i]; demod coefficient from its norm.
119 64 double ss = 0.0;
120
2/2
✓ Branch 0 taken 192 times.
✓ Branch 1 taken 64 times.
 256 for (int i = 0; i < C_in; ++i) {
121 192 const float si = sn[i];
122
2/2
✓ Branch 0 taken 1344 times.
✓ Branch 1 taken 192 times.
 1536 for (int t = 0; t < kH * kW; ++t) {
123 1344 const float wp = Wo[i * kH * kW + t] * si;
124 1344 Wno[i * kH * kW + t] = wp;
125 1344 ss += static_cast<double>(wp) * wp;
126 1344 }
127 192 }
128
2/2
✓ Branch 0 taken 40 times.
✓ Branch 1 taken 24 times.
 64 const float d = demodulate
129 40 ? 1.0f / std::sqrt(static_cast<float>(ss) + eps) : 1.0f;
130 64 dcp[static_cast<size_t>(n) * C_out + o] = d;
131
4/4
✓ Branch 0 taken 24 times.
✓ Branch 1 taken 40 times.
✓ Branch 2 taken 888 times.
✓ Branch 3 taken 40 times.
 952 if (demodulate) for (int t = 0; t < wk; ++t) Wno[t] *= d;
132 64 }
133
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 ::brotensor::Tensor Xn = row_view(X, n, C_in * H * Wd);
134
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 ::brotensor::Tensor Yn = row_view_mut(Y, n, out_cols);
135
2/4
✗ Branch 0 not taken.
✓ Branch 1 taken 16 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 16 times.
 32 conv2d_forward(Xn, Wn, nullptr, 1, C_in, H, Wd, C_out, kH, kW,
136 16 1, 1, pad_h, pad_w, 1, 1, 1, Yn);
137 16 }
138 8 }
139
140 12 void modulated_conv2d_backward(const ::brotensor::Tensor& X,
141 const ::brotensor::Tensor& W,
142 const ::brotensor::Tensor& s,
143 const ::brotensor::Tensor& dcoef,
144 const ::brotensor::Tensor& dY,
145 int N, int C_in, int H, int Wd,
146 int C_out, int kH, int kW,
147 int pad_h, int pad_w, bool demodulate, float eps,
148 ::brotensor::Tensor& dX,
149 ::brotensor::Tensor& dW,
150 ::brotensor::Tensor& ds) {
151 12 check_fp32(X, "modulated_conv2d_backward", "X");
152 12 check_fp32(W, "modulated_conv2d_backward", "W");
153 12 check_fp32(s, "modulated_conv2d_backward", "s");
154 12 check_fp32(dcoef, "modulated_conv2d_backward", "dcoef");
155 12 check_fp32(dY, "modulated_conv2d_backward", "dY");
156 (void)eps; // demod coefficient is precomputed (passed in as dcoef)
157 12 const int wk = C_in * kH * kW;
158 12 const int H_out = H + 2 * pad_h - (kH - 1);
159 12 const int W_out = Wd + 2 * pad_w - (kW - 1);
160
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (W.rows != C_out || W.cols != wk)
161 ✗ throw std::runtime_error("modulated_conv2d_backward: W shape mismatch");
162 // dW is an optional output: an uncommitted (data == nullptr) dW means
163 // "skip the weight gradient" (mirrors the CUDA backend — used by inversion,
164 // which freezes the weights). When committed it must be FP32 and shaped.
165 12 const bool want_dW = (dW.data != nullptr);
166
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 8 times.
 12 if (want_dW) {
167 8 check_fp32(dW, "modulated_conv2d_backward", "dW");
168
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (dW.rows != C_out || dW.cols != wk)
169 ✗ throw std::runtime_error("modulated_conv2d_backward: dW shape mismatch");
170 8 }
171
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 12 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 12 if (dX.rows != N || dX.cols != C_in * H * Wd || dX.dtype != Dtype::FP32)
172 12 dX.resize(N, C_in * H * Wd, Dtype::FP32);
173
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 12 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 12 if (ds.rows != N || ds.cols != C_in || ds.dtype != Dtype::FP32)
174 12 ds.resize(N, C_in, Dtype::FP32);
175
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (N == 0) return;
176
177 12 const float* Wp = W.host_f32();
178 12 const float* sp = s.host_f32();
179 12 const float* dcp = dcoef.host_f32();
180
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 4 times.
 12 float* dWp = want_dW ? dW.host_f32_mut() : nullptr; // accumulate (optional)
181 12 float* dsp = ds.host_f32_mut(); // overwrite
182
183 12 const int out_cols = C_out * H_out * W_out;
184 12 ::brotensor::Tensor Wpp = ::brotensor::Tensor::zeros_on(Device::CPU, C_out, wk); // w''
185
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 ::brotensor::Tensor Wpr = ::brotensor::Tensor::zeros_on(Device::CPU, C_out, wk); // w'
186
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 ::brotensor::Tensor dWpp = ::brotensor::Tensor::zeros_on(Device::CPU, C_out, wk);
187
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 std::vector<float> dWpr(static_cast<size_t>(C_out) * wk); // dw'
188
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 float* Wppp = Wpp.host_f32_mut();
189
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 float* Wprp = Wpr.host_f32_mut();
190
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 float* dWppp = dWpp.host_f32_mut();
191
192
2/2
✓ Branch 0 taken 24 times.
✓ Branch 1 taken 12 times.
 36 for (int n = 0; n < N; ++n) {
193 24 const float* sn = sp + static_cast<size_t>(n) * C_in;
194 24 const float* dcn = dcp + static_cast<size_t>(n) * C_out;
195 // Rebuild w' and w''.
196
2/2
✓ Branch 0 taken 96 times.
✓ Branch 1 taken 24 times.
 120 for (int o = 0; o < C_out; ++o) {
197 96 const float* Wo = Wp + static_cast<size_t>(o) * wk;
198 96 const float d = dcn[o];
199
2/2
✓ Branch 0 taken 288 times.
✓ Branch 1 taken 96 times.
 384 for (int i = 0; i < C_in; ++i) {
200 288 const float si = sn[i];
201
2/2
✓ Branch 0 taken 1824 times.
✓ Branch 1 taken 288 times.
 2112 for (int t = 0; t < kH * kW; ++t) {
202 1824 const int col = i * kH * kW + t;
203 1824 const float wp = Wo[col] * si;
204 1824 Wprp[static_cast<size_t>(o) * wk + col] = wp;
205 1824 Wppp[static_cast<size_t>(o) * wk + col] = wp * d;
206 1824 }
207 288 }
208 96 }
209
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 ::brotensor::Tensor Xn = row_view(X, n, C_in * H * Wd);
210
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 ::brotensor::Tensor dYn = row_view(dY, n, out_cols);
211 // dw'' (conv2d_backward_weight accumulates → zero first), and dX[n].
212
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 dWpp.zero();
213
2/4
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 24 times.
✗ Branch 3 not taken.
 48 conv2d_backward_weight(Xn, dYn, 1, C_in, H, Wd, C_out, kH, kW,
214 24 1, 1, pad_h, pad_w, 1, 1, 1, dWpp);
215
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 ::brotensor::Tensor dXn = row_view_mut(dX, n, C_in * H * Wd);
216
2/4
✗ Branch 0 not taken.
✓ Branch 1 taken 24 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 24 times.
 48 conv2d_backward_input(Wpp, dYn, 1, C_in, H, Wd, C_out, kH, kW,
217 24 1, 1, pad_h, pad_w, 1, 1, 1, dXn);
218 // Through demod → dw'.
219
2/2
✓ Branch 0 taken 96 times.
✓ Branch 1 taken 24 times.
 120 for (int o = 0; o < C_out; ++o) {
220 96 const size_t ob = static_cast<size_t>(o) * wk;
221
2/2
✓ Branch 0 taken 56 times.
✓ Branch 1 taken 40 times.
 96 if (demodulate) {
222 56 const float d = dcn[o];
223 56 double g = 0.0;
224
2/2
✓ Branch 0 taken 1128 times.
✓ Branch 1 taken 56 times.
 1184 for (int t = 0; t < wk; ++t)
225 1128 g += static_cast<double>(dWppp[ob + t]) * Wprp[ob + t];
226 56 const float gd3 = static_cast<float>(g) * d * d * d;
227
2/2
✓ Branch 0 taken 1128 times.
✓ Branch 1 taken 56 times.
 1184 for (int t = 0; t < wk; ++t)
228 1128 dWpr[ob + t] = dWppp[ob + t] * d - gd3 * Wprp[ob + t];
229 56 } else {
230
2/2
✓ Branch 0 taken 696 times.
✓ Branch 1 taken 40 times.
 736 for (int t = 0; t < wk; ++t) dWpr[ob + t] = dWppp[ob + t];
231 }
232 96 }
233 // Accumulate dW and write ds[n].
234 24 float* dsn = dsp + static_cast<size_t>(n) * C_in;
235
2/2
✓ Branch 0 taken 72 times.
✓ Branch 1 taken 24 times.
 96 for (int i = 0; i < C_in; ++i) dsn[i] = 0.0f;
236
2/2
✓ Branch 0 taken 96 times.
✓ Branch 1 taken 24 times.
 120 for (int o = 0; o < C_out; ++o) {
237 96 const float* Wo = Wp + static_cast<size_t>(o) * wk;
238 96 const size_t ob = static_cast<size_t>(o) * wk;
239
2/2
✓ Branch 0 taken 288 times.
✓ Branch 1 taken 96 times.
 384 for (int i = 0; i < C_in; ++i) {
240 288 const float si = sn[i];
241 288 double ds_acc = 0.0;
242
2/2
✓ Branch 0 taken 1824 times.
✓ Branch 1 taken 288 times.
 2112 for (int t = 0; t < kH * kW; ++t) {
243 1824 const int col = i * kH * kW + t;
244 1824 const float dwp = dWpr[ob + col];
245
2/2
✓ Branch 0 taken 480 times.
✓ Branch 1 taken 1344 times.
 1824 if (want_dW) dWp[ob + col] += dwp * si; // accumulate dW (optional)
246 1824 ds_acc += static_cast<double>(dwp) * Wo[col];
247 1824 }
248 288 dsn[i] += static_cast<float>(ds_acc); // overwrite ds[n,i]
249 288 }
250 96 }
251 24 }
252 12 }
253
254 } // namespace brotensor::detail::cpu
255