GCC Code Coverage Report

Directory: brotensor/
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 164 inline void check_fp32(const ::brotensor::Tensor& t,
56 const char* op, const char* name) {
57
1/2
✓ Branch 0 taken 164 times.
✗ Branch 1 not taken.
 164 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 164 }
63
64 // Non-owning CPU FP32 view over one row of a (N, cols) tensor.
65 100 inline ::brotensor::Tensor row_view(const ::brotensor::Tensor& T, int n, int cols) {
66 100 float* base = const_cast<float*>(T.host_f32()) + static_cast<size_t>(n) * cols;
67 100 return ::brotensor::Tensor::view(Device::CPU, base, 1, cols, Dtype::FP32);
68 }
69 64 inline ::brotensor::Tensor row_view_mut(::brotensor::Tensor& T, int n, int cols) {
70 64 float* base = T.host_f32_mut() + static_cast<size_t>(n) * cols;
71 64 return ::brotensor::Tensor::view(Device::CPU, base, 1, cols, Dtype::FP32);
72 }
73
74 } // namespace
75
76 16 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 16 check_fp32(X, "modulated_conv2d_forward", "X");
86 16 check_fp32(W, "modulated_conv2d_forward", "W");
87 16 check_fp32(s, "modulated_conv2d_forward", "s");
88 16 const int wk = C_in * kH * kW;
89
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 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 16 times.
✗ Branch 1 not taken.
 16 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 16 times.
✗ Branch 1 not taken.
 16 if (s.rows != N || s.cols != C_in)
94 ✗ throw std::runtime_error("modulated_conv2d_forward: s shape mismatch");
95 16 const int H_out = H + 2 * pad_h - (kH - 1);
96 16 const int W_out = Wd + 2 * pad_w - (kW - 1);
97
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 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 16 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 16 if (dcoef.rows != N || dcoef.cols != C_out || dcoef.dtype != Dtype::FP32)
100 16 dcoef.resize(N, C_out, Dtype::FP32);
101 16 const int out_cols = C_out * H_out * W_out;
102
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 16 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 16 if (Y.rows != N || Y.cols != out_cols || Y.dtype != Dtype::FP32)
103 16 Y.resize(N, out_cols, Dtype::FP32);
104
2/4
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 16 times.
 16 if (N == 0 || out_cols == 0) return;
105
106 16 const float* Wp = W.host_f32();
107 16 const float* sp = s.host_f32();
108 16 float* dcp = dcoef.host_f32_mut();
109
110 16 ::brotensor::Tensor Wn = ::brotensor::Tensor::zeros_on(Device::CPU, C_out, wk);
111
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 float* Wnp = Wn.host_f32_mut();
112
113
2/2
✓ Branch 0 taken 28 times.
✓ Branch 1 taken 16 times.
 44 for (int n = 0; n < N; ++n) {
114 28 const float* sn = sp + static_cast<size_t>(n) * C_in;
115
2/2
✓ Branch 0 taken 928 times.
✓ Branch 1 taken 28 times.
 956 for (int o = 0; o < C_out; ++o) {
116 928 const float* Wo = Wp + static_cast<size_t>(o) * wk;
117 928 float* Wno = Wnp + static_cast<size_t>(o) * wk;
118 // w'[o,..] = W[o,..]*s[n,i]; demod coefficient from its norm.
119 928 double ss = 0.0;
120
2/2
✓ Branch 0 taken 34752 times.
✓ Branch 1 taken 928 times.
 35680 for (int i = 0; i < C_in; ++i) {
121 34752 const float si = sn[i];
122
2/2
✓ Branch 0 taken 35904 times.
✓ Branch 1 taken 34752 times.
 70656 for (int t = 0; t < kH * kW; ++t) {
123 35904 const float wp = Wo[i * kH * kW + t] * si;
124 35904 Wno[i * kH * kW + t] = wp;
125 35904 ss += static_cast<double>(wp) * wp;
126 35904 }
127 34752 }
128
2/2
✓ Branch 0 taken 472 times.
✓ Branch 1 taken 456 times.
 928 const float d = demodulate
129 472 ? 1.0f / std::sqrt(static_cast<float>(ss) + eps) : 1.0f;
130 928 dcp[static_cast<size_t>(n) * C_out + o] = d;
131
4/4
✓ Branch 0 taken 456 times.
✓ Branch 1 taken 472 times.
✓ Branch 2 taken 18168 times.
✓ Branch 3 taken 472 times.
 19096 if (demodulate) for (int t = 0; t < wk; ++t) Wno[t] *= d;
132 928 }
133
1/2
✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
 28 ::brotensor::Tensor Xn = row_view(X, n, C_in * H * Wd);
134
1/2
✓ Branch 0 taken 28 times.
✗ Branch 1 not taken.
 28 ::brotensor::Tensor Yn = row_view_mut(Y, n, out_cols);
135
2/4
✗ Branch 0 not taken.
✓ Branch 1 taken 28 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 28 times.
 56 conv2d_forward(Xn, Wn, nullptr, 1, C_in, H, Wd, C_out, kH, kW,
136 28 1, 1, pad_h, pad_w, 1, 1, 1, Yn);
137 28 }
138 16 }
139
140 20 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 20 check_fp32(X, "modulated_conv2d_backward", "X");
152 20 check_fp32(W, "modulated_conv2d_backward", "W");
153 20 check_fp32(s, "modulated_conv2d_backward", "s");
154 20 check_fp32(dcoef, "modulated_conv2d_backward", "dcoef");
155 20 check_fp32(dY, "modulated_conv2d_backward", "dY");
156 (void)eps; // demod coefficient is precomputed (passed in as dcoef)
157 20 const int wk = C_in * kH * kW;
158 20 const int H_out = H + 2 * pad_h - (kH - 1);
159 20 const int W_out = Wd + 2 * pad_w - (kW - 1);
160
1/2
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
 20 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 20 const bool want_dW = (dW.data != nullptr);
166
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 16 times.
 20 if (want_dW) {
167 16 check_fp32(dW, "modulated_conv2d_backward", "dW");
168
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (dW.rows != C_out || dW.cols != wk)
169 ✗ throw std::runtime_error("modulated_conv2d_backward: dW shape mismatch");
170 16 }
171
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 20 if (dX.rows != N || dX.cols != C_in * H * Wd || dX.dtype != Dtype::FP32)
172 20 dX.resize(N, C_in * H * Wd, Dtype::FP32);
173
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 20 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 20 if (ds.rows != N || ds.cols != C_in || ds.dtype != Dtype::FP32)
174 20 ds.resize(N, C_in, Dtype::FP32);
175
1/2
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
 20 if (N == 0) return;
176
177 20 const float* Wp = W.host_f32();
178 20 const float* sp = s.host_f32();
179 20 const float* dcp = dcoef.host_f32();
180
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 4 times.
 20 float* dWp = want_dW ? dW.host_f32_mut() : nullptr; // accumulate (optional)
181 20 float* dsp = ds.host_f32_mut(); // overwrite
182
183 20 const int out_cols = C_out * H_out * W_out;
184 20 ::brotensor::Tensor Wpp = ::brotensor::Tensor::zeros_on(Device::CPU, C_out, wk); // w''
185
1/2
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
 20 ::brotensor::Tensor Wpr = ::brotensor::Tensor::zeros_on(Device::CPU, C_out, wk); // w'
186
1/2
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
 20 ::brotensor::Tensor dWpp = ::brotensor::Tensor::zeros_on(Device::CPU, C_out, wk);
187
1/2
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
 20 std::vector<float> dWpr(static_cast<size_t>(C_out) * wk); // dw'
188
1/2
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
 20 float* Wppp = Wpp.host_f32_mut();
189
1/2
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
 20 float* Wprp = Wpr.host_f32_mut();
190
1/2
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
 20 float* dWppp = dWpp.host_f32_mut();
191
192
2/2
✓ Branch 0 taken 36 times.
✓ Branch 1 taken 20 times.
 56 for (int n = 0; n < N; ++n) {
193 36 const float* sn = sp + static_cast<size_t>(n) * C_in;
194 36 const float* dcn = dcp + static_cast<size_t>(n) * C_out;
195 // Rebuild w' and w''.
196
2/2
✓ Branch 0 taken 960 times.
✓ Branch 1 taken 36 times.
 996 for (int o = 0; o < C_out; ++o) {
197 960 const float* Wo = Wp + static_cast<size_t>(o) * wk;
198 960 const float d = dcn[o];
199
2/2
✓ Branch 0 taken 34848 times.
✓ Branch 1 taken 960 times.
 35808 for (int i = 0; i < C_in; ++i) {
200 34848 const float si = sn[i];
201
2/2
✓ Branch 0 taken 36384 times.
✓ Branch 1 taken 34848 times.
 71232 for (int t = 0; t < kH * kW; ++t) {
202 36384 const int col = i * kH * kW + t;
203 36384 const float wp = Wo[col] * si;
204 36384 Wprp[static_cast<size_t>(o) * wk + col] = wp;
205 36384 Wppp[static_cast<size_t>(o) * wk + col] = wp * d;
206 36384 }
207 34848 }
208 960 }
209
1/2
✓ Branch 0 taken 36 times.
✗ Branch 1 not taken.
 36 ::brotensor::Tensor Xn = row_view(X, n, C_in * H * Wd);
210
1/2
✓ Branch 0 taken 36 times.
✗ Branch 1 not taken.
 36 ::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 36 times.
✗ Branch 1 not taken.
 36 dWpp.zero();
213
2/4
✓ Branch 0 taken 36 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 36 times.
✗ Branch 3 not taken.
 72 conv2d_backward_weight(Xn, dYn, 1, C_in, H, Wd, C_out, kH, kW,
214 36 1, 1, pad_h, pad_w, 1, 1, 1, dWpp);
215
1/2
✓ Branch 0 taken 36 times.
✗ Branch 1 not taken.
 36 ::brotensor::Tensor dXn = row_view_mut(dX, n, C_in * H * Wd);
216
2/4
✗ Branch 0 not taken.
✓ Branch 1 taken 36 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 36 times.
 72 conv2d_backward_input(Wpp, dYn, 1, C_in, H, Wd, C_out, kH, kW,
217 36 1, 1, pad_h, pad_w, 1, 1, 1, dXn);
218 // Through demod → dw'.
219
2/2
✓ Branch 0 taken 960 times.
✓ Branch 1 taken 36 times.
 996 for (int o = 0; o < C_out; ++o) {
220 960 const size_t ob = static_cast<size_t>(o) * wk;
221
2/2
✓ Branch 0 taken 488 times.
✓ Branch 1 taken 472 times.
 960 if (demodulate) {
222 488 const float d = dcn[o];
223 488 double g = 0.0;
224
2/2
✓ Branch 0 taken 18408 times.
✓ Branch 1 taken 488 times.
 18896 for (int t = 0; t < wk; ++t)
225 18408 g += static_cast<double>(dWppp[ob + t]) * Wprp[ob + t];
226 488 const float gd3 = static_cast<float>(g) * d * d * d;
227
2/2
✓ Branch 0 taken 18408 times.
✓ Branch 1 taken 488 times.
 18896 for (int t = 0; t < wk; ++t)
228 18408 dWpr[ob + t] = dWppp[ob + t] * d - gd3 * Wprp[ob + t];
229 488 } else {
230
2/2
✓ Branch 0 taken 17976 times.
✓ Branch 1 taken 472 times.
 18448 for (int t = 0; t < wk; ++t) dWpr[ob + t] = dWppp[ob + t];
231 }
232 960 }
233 // Accumulate dW and write ds[n].
234 36 float* dsn = dsp + static_cast<size_t>(n) * C_in;
235
2/2
✓ Branch 0 taken 552 times.
✓ Branch 1 taken 36 times.
 588 for (int i = 0; i < C_in; ++i) dsn[i] = 0.0f;
236
2/2
✓ Branch 0 taken 960 times.
✓ Branch 1 taken 36 times.
 996 for (int o = 0; o < C_out; ++o) {
237 960 const float* Wo = Wp + static_cast<size_t>(o) * wk;
238 960 const size_t ob = static_cast<size_t>(o) * wk;
239
2/2
✓ Branch 0 taken 34848 times.
✓ Branch 1 taken 960 times.
 35808 for (int i = 0; i < C_in; ++i) {
240 34848 const float si = sn[i];
241 34848 double ds_acc = 0.0;
242
2/2
✓ Branch 0 taken 36384 times.
✓ Branch 1 taken 34848 times.
 71232 for (int t = 0; t < kH * kW; ++t) {
243 36384 const int col = i * kH * kW + t;
244 36384 const float dwp = dWpr[ob + col];
245
2/2
✓ Branch 0 taken 480 times.
✓ Branch 1 taken 35904 times.
 36384 if (want_dW) dWp[ob + col] += dwp * si; // accumulate dW (optional)
246 36384 ds_acc += static_cast<double>(dwp) * Wo[col];
247 36384 }
248 34848 dsn[i] += static_cast<float>(ds_acc); // overwrite ds[n,i]
249 34848 }
250 960 }
251 36 }
252 20 }
253
254 } // namespace brotensor::detail::cpu
255