GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 72.0% 54 / 0 / 75
Functions: 100.0% 5 / 0 / 5
Branches: 38.9% 21 / 0 / 54
src/filtered_lrelu.cpp
Line Branch Exec Source
1 // ─── filtered_lrelu (StyleGAN3 alias-free nonlinearity) ─────────────────────
2 //
3 // The public entry points are thin dispatchers: if the resolved backend
4 // registered a fused `filtered_lrelu_forward/backward` vtable slot (CUDA does),
5 // they call it; otherwise they fall back to the device-agnostic COMPOSITE below
6 // — a sequence of the public bias_act + upfirdn2d ops, which run on whatever
7 // backend the operands live on (this is the only path on CPU/Metal). The
8 // composite mirrors NVlabs `_filtered_lrelu_ref` EXACTLY, including order:
9 //
10 // x = bias_act(x, b) # apply channel bias
11 // x = upfirdn2d(x, fu, up=up, pad=p, gain=up^2) # upsample
12 // x = bias_act(x, act=lrelu, gain, clamp) # bias-free lrelu+clamp
13 // x = upfirdn2d(x, fd, down=down) # downsample
14 //
15 // The bias is applied BEFORE the upsample (a linear bias_act), and the
16 // post-upsample bias_act carries only the lrelu — this is not interchangeable
17 // with biasing after the upsample. The backward reverses the chain.
18 //
19 // up_buf/act_buf are returned as caches: up_buf (the post-upsample tensor) is
20 // the input to the lrelu bias_act and is required by the backward; act_buf is
21 // kept for symmetry with the fused-kernel surface. A fused backend that does
22 // not reproduce these caches must keep whatever its own backward consumes —
23 // the cache contract is per-backend (see the CUDA kernel).
24
25 #include <brotensor/ops.h>
26 #include <brotensor/tensor.h>
27 #include <brotensor/detail/dispatch.h>
28
29 #include <stdexcept>
30
31 namespace brotensor {
32
33 // Composite fallback, also reused by the CUDA backend for configs its fused
34 // kernel does not cover. Declared here (not in the public header) and called
35 // from src/cuda/filtered_lrelu.cu via a matching forward declaration.
36 void filtered_lrelu_forward_composite(const Tensor& X, const Tensor& fu,
37 const Tensor& fd, const Tensor* b,
38 int N, int C, int H, int W, int up, int down,
39 int pad_x0, int pad_x1, int pad_y0, int pad_y1,
40 float gain, float slope, float clamp,
41 Tensor& up_buf, Tensor& act_buf, Tensor& Y);
42 void filtered_lrelu_backward_composite(const Tensor& dY, const Tensor& X,
43 const Tensor& fu, const Tensor& fd,
44 const Tensor* b, int N, int C, int H, int W,
45 int up, int down, int pad_x0, int pad_x1,
46 int pad_y0, int pad_y1, float gain, float slope,
47 float clamp, const Tensor& up_buf,
48 Tensor& dX, Tensor* dB);
49
50 namespace {
51
52 constexpr int ACT_LINEAR = 0;
53 constexpr int ACT_LRELU = 1;
54 constexpr float NO_CLAMP = -1.0f;
55
56 // Up-stage upfirdn2d output extent (down=1): (in*up + pad0 + pad1 - f) + 1.
57 28 inline int up_out(int in, int up, int pad0, int pad1, int fdim) {
58 28 return in * up + pad0 + pad1 - fdim + 1;
59 }
60
61 } // namespace
62
63 // ─── public dispatchers ─────────────────────────────────────────────────────
64
65 7 void filtered_lrelu_forward(const Tensor& X, const Tensor& fu, const Tensor& fd,
66 const Tensor* b, int N, int C, int H, int W,
67 int up, int down, int pad_x0, int pad_x1,
68 int pad_y0, int pad_y1, float gain, float slope,
69 float clamp, Tensor& up_buf, Tensor& act_buf,
70 Tensor& Y) {
71 7 const auto& v = detail::dispatch(X, fu, fd, up_buf, act_buf, Y);
72
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 7 times.
 7 if (v.filtered_lrelu_forward) {
73 ✗ detail::adopt_output(up_buf, X.device);
74 ✗ detail::adopt_output(act_buf, X.device);
75 ✗ detail::adopt_output(Y, X.device);
76 ✗ v.filtered_lrelu_forward(X, fu, fd, b, N, C, H, W, up, down,
77 ✗ pad_x0, pad_x1, pad_y0, pad_y1, gain, slope,
78 ✗ clamp, up_buf, act_buf, Y);
79 ✗ return;
80 }
81 14 filtered_lrelu_forward_composite(X, fu, fd, b, N, C, H, W, up, down,
82 7 pad_x0, pad_x1, pad_y0, pad_y1, gain, slope,
83 7 clamp, up_buf, act_buf, Y);
84 7 }
85
86 7 void filtered_lrelu_backward(const Tensor& dY, const Tensor& X,
87 const Tensor& fu, const Tensor& fd,
88 const Tensor* b, int N, int C, int H, int W,
89 int up, int down, int pad_x0, int pad_x1,
90 int pad_y0, int pad_y1, float gain, float slope,
91 float clamp, const Tensor& up_buf,
92 Tensor& dX, Tensor* dB) {
93 7 const auto& v = detail::dispatch(dY, X, fu, fd, up_buf, dX);
94
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 7 times.
 7 if (v.filtered_lrelu_backward) {
95 ✗ detail::adopt_output(dX, X.device);
96 ✗ if (dB) detail::adopt_output(*dB, X.device);
97 ✗ v.filtered_lrelu_backward(dY, X, fu, fd, b, N, C, H, W, up, down,
98 ✗ pad_x0, pad_x1, pad_y0, pad_y1, gain, slope,
99 ✗ clamp, up_buf, dX, dB);
100 ✗ return;
101 }
102 14 filtered_lrelu_backward_composite(dY, X, fu, fd, b, N, C, H, W, up, down,
103 7 pad_x0, pad_x1, pad_y0, pad_y1, gain, slope,
104 7 clamp, up_buf, dX, dB);
105 7 }
106
107 // ─── composite implementation (fallback / CPU / Metal) ──────────────────────
108
109 7 void filtered_lrelu_forward_composite(const Tensor& X, const Tensor& fu, const Tensor& fd,
110 const Tensor* b, int N, int C, int H, int W,
111 int up, int down, int pad_x0, int pad_x1,
112 int pad_y0, int pad_y1, float gain, float slope,
113 float clamp, Tensor& up_buf, Tensor& act_buf,
114 Tensor& Y) {
115 7 const int fuH = fu.rows, fuW = fu.cols;
116 7 const int fdH = fd.rows, fdW = fd.cols;
117
118 // 1. Channel bias (linear bias_act), at the input rate.
119 7 Tensor pre;
120
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 bias_act_forward(X, b, N, C, H * W, ACT_LINEAR, 0.0f, 1.0f, NO_CLAMP, pre);
121
122 // 2. Upsample with gain = up^2.
123
2/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
 14 upfirdn2d_forward(pre, fu, N, C, H, W, fuH, fuW,
124 7 up, up, 1, 1, pad_x0, pad_x1, pad_y0, pad_y1,
125 7 /*flip=*/false, static_cast<float>(up * up), up_buf);
126
127
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 const int Huo = up_out(H, up, pad_y0, pad_y1, fuH);
128
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 const int Wuo = up_out(W, up, pad_x0, pad_x1, fuW);
129
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (up_buf.cols != C * Huo * Wuo) {
130 ✗ throw std::runtime_error("filtered_lrelu_forward: upsample dim mismatch");
131 }
132
133 // 3. Bias-free leaky ReLU + clamp at the 2x rate.
134
2/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
 14 bias_act_forward(up_buf, nullptr, N, C, Huo * Wuo, ACT_LRELU, slope,
135 7 gain, clamp, act_buf);
136
137 // 4. Downsample (gain 1, no padding).
138
2/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
 14 upfirdn2d_forward(act_buf, fd, N, C, Huo, Wuo, fdH, fdW,
139 7 1, 1, down, down, 0, 0, 0, 0,
140 7 /*flip=*/false, 1.0f, Y);
141 7 }
142
143 7 void filtered_lrelu_backward_composite(const Tensor& dY, const Tensor& X,
144 const Tensor& fu, const Tensor& fd,
145 const Tensor* b, int N, int C, int H, int W,
146 int up, int down, int pad_x0, int pad_x1,
147 int pad_y0, int pad_y1, float gain, float slope,
148 float clamp, const Tensor& up_buf,
149 Tensor& dX, Tensor* dB) {
150 7 const int fuH = fu.rows, fuW = fu.cols;
151 7 const int fdH = fd.rows, fdW = fd.cols;
152 7 const int Huo = up_out(H, up, pad_y0, pad_y1, fuH);
153 7 const int Wuo = up_out(W, up, pad_x0, pad_x1, fuW);
154
155 // up_buf is the post-upsample (pre-lrelu) tensor — the lrelu backward needs
156 // it. The fused forward skips producing it (it's the buffer fusion avoids),
157 // so when the caller hands us an uncommitted up_buf we recompute it here
158 // from X exactly as the forward did (bias → up-FIR). The cache is thus
159 // optional: a populated up_buf is used directly, an empty one is rebuilt.
160 7 Tensor up_buf_local;
161 7 const Tensor* up_ptr = &up_buf;
162
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (up_buf.data == nullptr) {
163 ✗ Tensor pre;
164 ✗ bias_act_forward(X, b, N, C, H * W, ACT_LINEAR, 0.0f, 1.0f, NO_CLAMP, pre);
165 ✗ upfirdn2d_forward(pre, fu, N, C, H, W, fuH, fuW,
166 ✗ up, up, 1, 1, pad_x0, pad_x1, pad_y0, pad_y1,
167 ✗ /*flip=*/false, static_cast<float>(up * up), up_buf_local);
168 ✗ up_ptr = &up_buf_local;
169 ✗ }
170
171 // 4'. Through the downsample.
172 7 Tensor d_act;
173
2/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
 14 upfirdn2d_backward(dY, fd, N, C, Huo, Wuo, fdH, fdW,
174 7 1, 1, down, down, 0, 0, 0, 0,
175 /*flip=*/false, 1.0f, d_act);
176
177 // 3'. Through the leaky ReLU (no bias gradient here).
178 7 Tensor d_up;
179
2/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
 14 bias_act_backward(d_act, *up_ptr, nullptr, N, C, Huo * Wuo, ACT_LRELU, slope,
180 7 gain, clamp, d_up, nullptr);
181
182 // 2'. Through the upsample.
183 7 Tensor d_pre;
184
2/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
 14 upfirdn2d_backward(d_up, fu, N, C, H, W, fuH, fuW,
185 7 up, up, 1, 1, pad_x0, pad_x1, pad_y0, pad_y1,
186 7 /*flip=*/false, static_cast<float>(up * up), d_pre);
187
188 // 1'. Through the channel bias — dX (overwrite) and dB (accumulate).
189
2/4
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 7 times.
✗ Branch 3 not taken.
 14 bias_act_backward(d_pre, X, b, N, C, H * W, ACT_LINEAR, 0.0f, 1.0f,
190 7 NO_CLAMP, dX, dB);
191 7 }
192
193 } // namespace brotensor
194