GCC Code Coverage Report

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