GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 91.6% 568 / 0 / 620
Functions: 100.0% 67 / 0 / 67
Branches: 54.7% 117 / 0 / 214
src/ops_spatial.cpp
Line Branch Exec Source
1 // brotensor public op wrappers — 2D/3D spatial, conv, norm, pool, StyleGAN.
2
3 #include <brotensor/ops.h>
4 #include <brotensor/tensor.h>
5 #include <brotensor/detail/dispatch.h>
6
7 #include <stdexcept>
8
9 namespace brotensor {
10
11 // ─── Conv2d ────────────────────────────────────────────────────────────────
12
13 104 void conv2d_forward(const Tensor& X, const Tensor& Wt, const Tensor* bias,
14 int N, int C_in, int H, int W,
15 int C_out, int kH, int kW,
16 int stride_h, int stride_w,
17 int pad_h, int pad_w,
18 int dil_h, int dil_w,
19 int groups, Tensor& Y) {
20 104 BROTENSOR_CHECK_SPATIAL_2D("conv2d_forward", X, N, C_in, H, W);
21 104 const auto& v = detail::dispatch_with_opts(X, Wt, {bias, &Y});
22
1/2
✓ Branch 0 taken 104 times.
✗ Branch 1 not taken.
 104 if (!v.conv2d_forward) detail::throw_not_implemented("conv2d_forward", X.device);
23 104 detail::adopt_output(Y, X.device);
24 208 v.conv2d_forward(X, Wt, bias, N, C_in, H, W, C_out, kH, kW,
25 104 stride_h, stride_w, pad_h, pad_w, dil_h, dil_w, groups, Y);
26 104 }
27
28 54 void conv2d_backward_input(const Tensor& Wt, const Tensor& dY,
29 int N, int C_in, int H, int W,
30 int C_out, int kH, int kW,
31 int stride_h, int stride_w,
32 int pad_h, int pad_w,
33 int dil_h, int dil_w,
34 int groups, Tensor& dX) {
35
1/2
✓ Branch 0 taken 54 times.
✗ Branch 1 not taken.
 54 if (N <= 0 || C_in <= 0 || H <= 0 || W <= 0 || C_out <= 0) {
36 ✗ detail::throw_spatial_2d_error("conv2d_backward_input", "dY",
37 ✗ N, C_in, H, W, dY.rows, dY.cols);
38 }
39
6/12
✓ Branch 0 taken 54 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 54 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 54 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 54 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 54 times.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 54 times.
 54 if (kH > 0 && kW > 0 && stride_h > 0 && stride_w > 0 && dil_h > 0 && dil_w > 0) {
40 54 const int H_out = (H + 2 * pad_h - dil_h * (kH - 1) - 1) / stride_h + 1;
41 54 const int W_out = (W + 2 * pad_w - dil_w * (kW - 1) - 1) / stride_w + 1;
42
2/4
✓ Branch 0 taken 54 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 54 times.
 54 if (H_out > 0 && W_out > 0) {
43 54 BROTENSOR_CHECK_SPATIAL_2D("conv2d_backward_input", dY, N, C_out, H_out, W_out);
44 54 }
45 54 }
46 54 const auto& v = detail::dispatch(Wt, dY, dX);
47
1/2
✓ Branch 0 taken 54 times.
✗ Branch 1 not taken.
 54 if (!v.conv2d_backward_input) detail::throw_not_implemented("conv2d_backward_input", Wt.device);
48 54 detail::adopt_output(dX, Wt.device);
49 108 v.conv2d_backward_input(Wt, dY, N, C_in, H, W, C_out, kH, kW,
50 54 stride_h, stride_w, pad_h, pad_w, dil_h, dil_w, groups, dX);
51 54 }
52
53 44 void conv2d_backward_weight(const Tensor& X, const Tensor& dY,
54 int N, int C_in, int H, int W,
55 int C_out, int kH, int kW,
56 int stride_h, int stride_w,
57 int pad_h, int pad_w,
58 int dil_h, int dil_w,
59 int groups, Tensor& dWt) {
60 44 BROTENSOR_CHECK_SPATIAL_2D("conv2d_backward_weight", X, N, C_in, H, W);
61 44 const auto& v = detail::dispatch(X, dY, dWt);
62
1/2
✓ Branch 0 taken 44 times.
✗ Branch 1 not taken.
 44 if (!v.conv2d_backward_weight) detail::throw_not_implemented("conv2d_backward_weight", X.device);
63 44 detail::adopt_output(dWt, X.device);
64 88 v.conv2d_backward_weight(X, dY, N, C_in, H, W, C_out, kH, kW,
65 44 stride_h, stride_w, pad_h, pad_w, dil_h, dil_w, groups, dWt);
66 44 }
67
68 33 void conv2d_backward_bias(const Tensor& dY,
69 int N, int C_out, int H_out, int W_out,
70 Tensor& dB) {
71 33 BROTENSOR_CHECK_SPATIAL_2D("conv2d_backward_bias", dY, N, C_out, H_out, W_out);
72 33 const auto& v = detail::dispatch(dY, dB);
73
1/2
✓ Branch 0 taken 33 times.
✗ Branch 1 not taken.
 33 if (!v.conv2d_backward_bias) detail::throw_not_implemented("conv2d_backward_bias", dY.device);
74 33 detail::adopt_output(dB, dY.device);
75 33 v.conv2d_backward_bias(dY, N, C_out, H_out, W_out, dB);
76 33 }
77
78 10 void conv2d_int8w_fp16_forward(const Tensor& X,
79 const Tensor& W_int8, const Tensor& scales,
80 const Tensor* bias,
81 int N, int C_in, int H, int W,
82 int C_out, int kH, int kW,
83 int stride_h, int stride_w,
84 int pad_h, int pad_w,
85 int dil_h, int dil_w, int groups,
86 Tensor& Y) {
87 10 BROTENSOR_CHECK_SPATIAL_2D("conv2d_int8w_fp16_forward", X, N, C_in, H, W);
88 10 const auto& v = detail::dispatch_with_opts(X, W_int8, {&scales, bias, &Y});
89
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.conv2d_int8w_fp16_forward)
90 ✗ detail::throw_not_implemented("conv2d_int8w_fp16_forward", X.device);
91 10 detail::adopt_output(Y, X.device);
92 20 v.conv2d_int8w_fp16_forward(X, W_int8, scales, bias,
93 10 N, C_in, H, W, C_out, kH, kW,
94 10 stride_h, stride_w, pad_h, pad_w,
95 10 dil_h, dil_w, groups, Y);
96 10 }
97
98 18 void deform_conv2d_forward(const Tensor& X, const Tensor& offset,
99 const Tensor* mask, const Tensor& Wt,
100 const Tensor* bias,
101 int N, int C_in, int H, int W,
102 int C_out, int kH, int kW,
103 int stride_h, int stride_w,
104 int pad_h, int pad_w,
105 int dil_h, int dil_w,
106 int groups, int deform_groups, Tensor& Y) {
107 18 BROTENSOR_CHECK_SPATIAL_2D("deform_conv2d_forward", X, N, C_in, H, W);
108 18 const auto& v = detail::dispatch_with_opts(X, Wt, {&offset, mask, bias, &Y});
109
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 if (!v.deform_conv2d_forward) detail::throw_not_implemented("deform_conv2d_forward", X.device);
110 18 detail::adopt_output(Y, X.device);
111 36 v.deform_conv2d_forward(X, offset, mask, Wt, bias, N, C_in, H, W,
112 18 C_out, kH, kW, stride_h, stride_w, pad_h, pad_w,
113 18 dil_h, dil_w, groups, deform_groups, Y);
114 18 }
115
116 // ─── Conv3d ────────────────────────────────────────────────────────────────
117
118 45 void conv3d_forward(const Tensor& X, const Tensor& Wt, const Tensor* bias,
119 int N, int C_in, int T, int H, int W,
120 int C_out, int kT, int kH, int kW,
121 int stride_t, int stride_h, int stride_w,
122 int pad_t, int pad_h, int pad_w,
123 int dil_t, int dil_h, int dil_w,
124 int groups, Tensor& Y) {
125 45 BROTENSOR_CHECK_SPATIAL_3D("conv3d_forward", X, N, C_in, T, H, W);
126 45 const auto& v = detail::dispatch_with_opts(X, Wt, {bias, &Y});
127
1/2
✓ Branch 0 taken 45 times.
✗ Branch 1 not taken.
 45 if (!v.conv3d_forward) detail::throw_not_implemented("conv3d_forward", X.device);
128 45 detail::adopt_output(Y, X.device);
129 90 v.conv3d_forward(X, Wt, bias, N, C_in, T, H, W, C_out, kT, kH, kW,
130 45 stride_t, stride_h, stride_w,
131 45 pad_t, pad_h, pad_w,
132 45 dil_t, dil_h, dil_w, groups, Y);
133 45 }
134
135 9 void conv3d_int8w_fp16_forward(const Tensor& X, const Tensor& W_int8,
136 const Tensor& scales, const Tensor* bias,
137 int N, int C_in, int T, int H, int W,
138 int C_out, int kT, int kH, int kW,
139 int stride_t, int stride_h, int stride_w,
140 int pad_t, int pad_h, int pad_w,
141 int dil_t, int dil_h, int dil_w,
142 int groups, Tensor& Y) {
143 9 BROTENSOR_CHECK_SPATIAL_3D("conv3d_int8w_fp16_forward", X, N, C_in, T, H, W);
144 9 const auto& v = detail::dispatch_with_opts(X, W_int8, {&scales, bias, &Y});
145
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 1 time.
 9 if (!v.conv3d_int8w_fp16_forward)
146 1 detail::throw_not_implemented("conv3d_int8w_fp16_forward", X.device);
147 8 detail::adopt_output(Y, X.device);
148 16 v.conv3d_int8w_fp16_forward(X, W_int8, scales, bias,
149 8 N, C_in, T, H, W, C_out, kT, kH, kW,
150 8 stride_t, stride_h, stride_w,
151 8 pad_t, pad_h, pad_w,
152 8 dil_t, dil_h, dil_w, groups, Y);
153 8 }
154
155 // ─── GroupNorm ─────────────────────────────────────────────────────────────
156
157 41 void group_norm_forward(const Tensor& X, const Tensor& gamma, const Tensor& beta,
158 int N, int C, int H, int W, int num_groups,
159 float eps, Tensor& Y) {
160 41 BROTENSOR_CHECK_SPATIAL_2D("group_norm_forward", X, N, C, H, W);
161 41 const auto& v = detail::dispatch(X, gamma, beta, Y);
162
1/2
✓ Branch 0 taken 41 times.
✗ Branch 1 not taken.
 41 if (!v.group_norm_forward) detail::throw_not_implemented("group_norm_forward", X.device);
163 41 detail::adopt_output(Y, X.device);
164 41 v.group_norm_forward(X, gamma, beta, N, C, H, W, num_groups, eps, Y);
165 41 }
166
167 32 void group_norm_backward(const Tensor& X, const Tensor& gamma, const Tensor& dY,
168 int N, int C, int H, int W, int num_groups, float eps,
169 Tensor& dX, Tensor& dGamma, Tensor& dBeta) {
170 32 BROTENSOR_CHECK_SPATIAL_2D("group_norm_backward", X, N, C, H, W);
171 32 BROTENSOR_CHECK_SPATIAL_2D("group_norm_backward", dY, N, C, H, W);
172 32 const auto& v = detail::dispatch(X, gamma, dY, dX, dGamma, dBeta);
173
1/2
✓ Branch 0 taken 32 times.
✗ Branch 1 not taken.
 32 if (!v.group_norm_backward) detail::throw_not_implemented("group_norm_backward", X.device);
174 32 detail::adopt_output(dX, X.device);
175 32 detail::adopt_output(dGamma, X.device);
176 32 detail::adopt_output(dBeta, X.device);
177 64 v.group_norm_backward(X, gamma, dY, N, C, H, W, num_groups, eps,
178 32 dX, dGamma, dBeta);
179 32 }
180
181 // ─── Activations ───────────────────────────────────────────────────────────
182
183 575 void silu_forward(const Tensor& x, Tensor& y) {
184 575 const auto& v = detail::dispatch(x, y);
185
1/2
✓ Branch 0 taken 575 times.
✗ Branch 1 not taken.
 575 if (!v.silu_forward) detail::throw_not_implemented("silu_forward", x.device);
186 575 detail::adopt_output(y, x.device);
187 575 v.silu_forward(x, y);
188 575 }
189 12 void silu_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
190 12 const auto& v = detail::dispatch(x, dY, dX);
191
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.silu_backward) detail::throw_not_implemented("silu_backward", x.device);
192 12 detail::adopt_output(dX, x.device);
193 12 v.silu_backward(x, dY, dX);
194 12 }
195 12 void gelu_forward(const Tensor& x, Tensor& y) {
196 12 const auto& v = detail::dispatch(x, y);
197
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.gelu_forward) detail::throw_not_implemented("gelu_forward", x.device);
198 12 detail::adopt_output(y, x.device);
199 12 v.gelu_forward(x, y);
200 12 }
201 8 void gelu_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
202 8 const auto& v = detail::dispatch(x, dY, dX);
203
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.gelu_backward) detail::throw_not_implemented("gelu_backward", x.device);
204 8 detail::adopt_output(dX, x.device);
205 8 v.gelu_backward(x, dY, dX);
206 8 }
207 10 void gelu_exact_forward(const Tensor& x, Tensor& y) {
208 10 const auto& v = detail::dispatch(x, y);
209
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.gelu_exact_forward) detail::throw_not_implemented("gelu_exact_forward", x.device);
210 10 detail::adopt_output(y, x.device);
211 10 v.gelu_exact_forward(x, y);
212 10 }
213 8 void gelu_exact_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
214 8 const auto& v = detail::dispatch(x, dY, dX);
215
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.gelu_exact_backward) detail::throw_not_implemented("gelu_exact_backward", x.device);
216 8 detail::adopt_output(dX, x.device);
217 8 v.gelu_exact_backward(x, dY, dX);
218 8 }
219 10 void quick_gelu_forward(const Tensor& x, Tensor& y) {
220 10 const auto& v = detail::dispatch(x, y);
221
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.quick_gelu_forward) detail::throw_not_implemented("quick_gelu_forward", x.device);
222 10 detail::adopt_output(y, x.device);
223 10 v.quick_gelu_forward(x, y);
224 10 }
225 8 void quick_gelu_backward(const Tensor& x, const Tensor& dY, Tensor& dX) {
226 8 const auto& v = detail::dispatch(x, dY, dX);
227
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.quick_gelu_backward) detail::throw_not_implemented("quick_gelu_backward", x.device);
228 8 detail::adopt_output(dX, x.device);
229 8 v.quick_gelu_backward(x, dY, dX);
230 8 }
231
232 // ─── Resample ──────────────────────────────────────────────────────────────
233
234 22 void upsample_nearest_2x(const Tensor& X, int N, int C, int H, int W, Tensor& Y) {
235 22 BROTENSOR_CHECK_SPATIAL_2D("upsample_nearest_2x", X, N, C, H, W);
236 22 const auto& v = detail::dispatch(X, Y);
237
1/2
✓ Branch 0 taken 22 times.
✗ Branch 1 not taken.
 22 if (!v.upsample_nearest_2x) detail::throw_not_implemented("upsample_nearest_2x", X.device);
238 22 detail::adopt_output(Y, X.device);
239 22 v.upsample_nearest_2x(X, N, C, H, W, Y);
240 22 }
241 19 void upsample_bilinear_2x(const Tensor& X, int N, int C, int H, int W, Tensor& Y) {
242 19 BROTENSOR_CHECK_SPATIAL_2D("upsample_bilinear_2x", X, N, C, H, W);
243 19 const auto& v = detail::dispatch(X, Y);
244
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 if (!v.upsample_bilinear_2x) detail::throw_not_implemented("upsample_bilinear_2x", X.device);
245 19 detail::adopt_output(Y, X.device);
246 19 v.upsample_bilinear_2x(X, N, C, H, W, Y);
247 19 }
248 19 void downsample_avg_2x(const Tensor& X, int N, int C, int H, int W, Tensor& Y) {
249 19 BROTENSOR_CHECK_SPATIAL_2D("downsample_avg_2x", X, N, C, H, W);
250 19 const auto& v = detail::dispatch(X, Y);
251
1/2
✓ Branch 0 taken 19 times.
✗ Branch 1 not taken.
 19 if (!v.downsample_avg_2x) detail::throw_not_implemented("downsample_avg_2x", X.device);
252 19 detail::adopt_output(Y, X.device);
253 19 v.downsample_avg_2x(X, N, C, H, W, Y);
254 19 }
255 16 void upsample_nearest_2x_backward(const Tensor& dY, int N, int C, int H, int W, Tensor& dX) {
256
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (N <= 0 || C <= 0 || H <= 0 || W <= 0) {
257 ✗ detail::throw_spatial_2d_error("upsample_nearest_2x_backward", "dY",
258 ✗ N, C, H, W, dY.rows, dY.cols);
259 }
260 16 BROTENSOR_CHECK_SPATIAL_2D("upsample_nearest_2x_backward", dY, N, C, 2 * H, 2 * W);
261 16 const auto& v = detail::dispatch(dY, dX);
262
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.upsample_nearest_2x_backward) detail::throw_not_implemented("upsample_nearest_2x_backward", dY.device);
263 16 detail::adopt_output(dX, dY.device);
264 16 v.upsample_nearest_2x_backward(dY, N, C, H, W, dX);
265 16 }
266 16 void upsample_bilinear_2x_backward(const Tensor& dY, int N, int C, int H, int W, Tensor& dX) {
267
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (N <= 0 || C <= 0 || H <= 0 || W <= 0) {
268 ✗ detail::throw_spatial_2d_error("upsample_bilinear_2x_backward", "dY",
269 ✗ N, C, H, W, dY.rows, dY.cols);
270 }
271 16 BROTENSOR_CHECK_SPATIAL_2D("upsample_bilinear_2x_backward", dY, N, C, 2 * H, 2 * W);
272 16 const auto& v = detail::dispatch(dY, dX);
273
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.upsample_bilinear_2x_backward) detail::throw_not_implemented("upsample_bilinear_2x_backward", dY.device);
274 16 detail::adopt_output(dX, dY.device);
275 16 v.upsample_bilinear_2x_backward(dY, N, C, H, W, dX);
276 16 }
277 16 void downsample_avg_2x_backward(const Tensor& dY, int N, int C, int H, int W, Tensor& dX) {
278
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (N <= 0 || C <= 0 || H <= 0 || W <= 0) {
279 ✗ detail::throw_spatial_2d_error("downsample_avg_2x_backward", "dY",
280 ✗ N, C, H, W, dY.rows, dY.cols);
281 }
282
2/4
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 16 times.
 16 if (H % 2 == 0 && W % 2 == 0) {
283 16 BROTENSOR_CHECK_SPATIAL_2D("downsample_avg_2x_backward", dY, N, C, H / 2, W / 2);
284 16 }
285 16 const auto& v = detail::dispatch(dY, dX);
286
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.downsample_avg_2x_backward) detail::throw_not_implemented("downsample_avg_2x_backward", dY.device);
287 16 detail::adopt_output(dX, dY.device);
288 16 v.downsample_avg_2x_backward(dY, N, C, H, W, dX);
289 16 }
290
291 86 void interp2d_forward(const Tensor& X,
292 int N, int C, int H_in, int W_in, int H_out, int W_out,
293 int mode, Tensor& Y) {
294 86 BROTENSOR_CHECK_SPATIAL_2D("interp2d_forward", X, N, C, H_in, W_in);
295 86 const auto& v = detail::dispatch(X, Y);
296
1/2
✓ Branch 0 taken 86 times.
✗ Branch 1 not taken.
 86 if (!v.interp2d_forward) detail::throw_not_implemented("interp2d_forward", X.device);
297 86 detail::adopt_output(Y, X.device);
298 86 v.interp2d_forward(X, N, C, H_in, W_in, H_out, W_out, mode, Y);
299 86 }
300 17 void interp2d_backward(const Tensor& dY,
301 int N, int C, int H_in, int W_in, int H_out, int W_out,
302 int mode, Tensor& dX) {
303
1/2
✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
 17 if (N <= 0 || C <= 0 || H_in <= 0 || W_in <= 0) {
304 ✗ detail::throw_spatial_2d_error("interp2d_backward", "dY",
305 ✗ N, C, H_in, W_in, dY.rows, dY.cols);
306 }
307 17 BROTENSOR_CHECK_SPATIAL_2D("interp2d_backward", dY, N, C, H_out, W_out);
308 17 const auto& v = detail::dispatch(dY, dX);
309
1/2
✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
 17 if (!v.interp2d_backward) detail::throw_not_implemented("interp2d_backward", dY.device);
310 17 detail::adopt_output(dX, dY.device);
311 17 v.interp2d_backward(dY, N, C, H_in, W_in, H_out, W_out, mode, dX);
312 17 }
313 22 void interp2d_align_corners_forward(const Tensor& X,
314 int N, int C, int H_in, int W_in,
315 int H_out, int W_out, int mode, Tensor& Y) {
316 22 BROTENSOR_CHECK_SPATIAL_2D("interp2d_align_corners_forward", X, N, C, H_in, W_in);
317 22 const auto& v = detail::dispatch(X, Y);
318
1/2
✓ Branch 0 taken 22 times.
✗ Branch 1 not taken.
 22 if (!v.interp2d_align_corners_forward)
319 ✗ detail::throw_not_implemented("interp2d_align_corners_forward", X.device);
320 22 detail::adopt_output(Y, X.device);
321 22 v.interp2d_align_corners_forward(X, N, C, H_in, W_in, H_out, W_out, mode, Y);
322 22 }
323
324 65 void pad2d_forward(const Tensor& X, int N, int C, int H, int W,
325 int pad_top, int pad_bottom, int pad_left, int pad_right,
326 int mode, Tensor& Y) {
327 65 BROTENSOR_CHECK_SPATIAL_2D("pad2d_forward", X, N, C, H, W);
328 65 const auto& v = detail::dispatch(X, Y);
329
1/2
✓ Branch 0 taken 65 times.
✗ Branch 1 not taken.
 65 if (!v.pad2d_forward) detail::throw_not_implemented("pad2d_forward", X.device);
330 65 detail::adopt_output(Y, X.device);
331 65 v.pad2d_forward(X, N, C, H, W, pad_top, pad_bottom, pad_left, pad_right, mode, Y);
332 65 }
333 13 void pad2d_backward(const Tensor& dY, int N, int C, int H, int W,
334 int pad_top, int pad_bottom, int pad_left, int pad_right,
335 int mode, Tensor& dX) {
336
1/2
✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
 13 if (N <= 0 || C <= 0 || H <= 0 || W <= 0) {
337 ✗ detail::throw_spatial_2d_error("pad2d_backward", "dY",
338 ✗ N, C, H, W, dY.rows, dY.cols);
339 }
340 13 const int H_out = H + pad_top + pad_bottom;
341 13 const int W_out = W + pad_left + pad_right;
342
2/4
✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 13 times.
 13 if (H_out > 0 && W_out > 0) {
343 13 BROTENSOR_CHECK_SPATIAL_2D("pad2d_backward", dY, N, C, H_out, W_out);
344 13 }
345 13 const auto& v = detail::dispatch(dY, dX);
346
1/2
✓ Branch 0 taken 13 times.
✗ Branch 1 not taken.
 13 if (!v.pad2d_backward) detail::throw_not_implemented("pad2d_backward", dY.device);
347 13 detail::adopt_output(dX, dY.device);
348 13 v.pad2d_backward(dY, N, C, H, W, pad_top, pad_bottom, pad_left, pad_right, mode, dX);
349 13 }
350
351 16 void slice2d_forward(const Tensor& X, int N, int C, int H, int W,
352 int h0, int w0, int H_out, int W_out, Tensor& Y) {
353 16 BROTENSOR_CHECK_SPATIAL_2D("slice2d_forward", X, N, C, H, W);
354 16 const auto& v = detail::dispatch(X, Y);
355
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.slice2d_forward) detail::throw_not_implemented("slice2d_forward", X.device);
356 16 detail::adopt_output(Y, X.device);
357 16 v.slice2d_forward(X, N, C, H, W, h0, w0, H_out, W_out, Y);
358 16 }
359 8 void slice2d_backward(const Tensor& dY, int N, int C, int H, int W,
360 int h0, int w0, int H_out, int W_out, Tensor& dX) {
361
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (N <= 0 || C <= 0 || H <= 0 || W <= 0) {
362 ✗ detail::throw_spatial_2d_error("slice2d_backward", "dY",
363 ✗ N, C, H, W, dY.rows, dY.cols);
364 }
365
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 8 times.
 8 if (H_out > 0 && W_out > 0) {
366 8 BROTENSOR_CHECK_SPATIAL_2D("slice2d_backward", dY, N, C, H_out, W_out);
367 8 }
368 8 const auto& v = detail::dispatch(dY, dX);
369
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.slice2d_backward) detail::throw_not_implemented("slice2d_backward", dY.device);
370 8 detail::adopt_output(dX, dY.device);
371 8 v.slice2d_backward(dY, N, C, H, W, h0, w0, H_out, W_out, dX);
372 8 }
373
374 10 void unfold2d_forward(const Tensor& X, int N, int C, int H, int W,
375 int kH, int kW, int stride_h, int stride_w,
376 int pad_top, int pad_bottom, int pad_left, int pad_right,
377 int mode, Tensor& Y) {
378 10 BROTENSOR_CHECK_SPATIAL_2D("unfold2d_forward", X, N, C, H, W);
379 10 const auto& v = detail::dispatch(X, Y);
380
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.unfold2d_forward) detail::throw_not_implemented("unfold2d_forward", X.device);
381 10 detail::adopt_output(Y, X.device);
382 20 v.unfold2d_forward(X, N, C, H, W, kH, kW, stride_h, stride_w,
383 10 pad_top, pad_bottom, pad_left, pad_right, mode, Y);
384 10 }
385
386 10 void l2_normalize_nchw_forward(const Tensor& X, int N, int C, int H, int W,
387 float eps, Tensor& Y) {
388 10 BROTENSOR_CHECK_SPATIAL_2D("l2_normalize_nchw_forward", X, N, C, H, W);
389 10 const auto& v = detail::dispatch(X, Y);
390
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.l2_normalize_nchw_forward)
391 ✗ detail::throw_not_implemented("l2_normalize_nchw_forward", X.device);
392 10 detail::adopt_output(Y, X.device);
393 10 v.l2_normalize_nchw_forward(X, N, C, H, W, eps, Y);
394 10 }
395
396 10 void convex_upsample_forward(const Tensor& X, const Tensor& Mask,
397 int N, int C, int H, int W, int scale, Tensor& Y) {
398 10 BROTENSOR_CHECK_SPATIAL_2D("convex_upsample_forward", X, N, C, H, W);
399 10 const auto& v = detail::dispatch(X, Mask, Y);
400
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.convex_upsample_forward)
401 ✗ detail::throw_not_implemented("convex_upsample_forward", X.device);
402 10 detail::adopt_output(Y, X.device);
403 10 v.convex_upsample_forward(X, Mask, N, C, H, W, scale, Y);
404 10 }
405
406 63 void top_k_rows(const Tensor& X, int k, Tensor& Vals, Tensor& Idx) {
407 63 const auto& v = detail::dispatch(X, Vals, Idx);
408
1/2
✓ Branch 0 taken 63 times.
✗ Branch 1 not taken.
 63 if (!v.top_k_rows) detail::throw_not_implemented("top_k_rows", X.device);
409 63 detail::adopt_output(Vals, X.device);
410 63 detail::adopt_output(Idx, X.device);
411 63 v.top_k_rows(X, k, Vals, Idx);
412 63 }
413
414 46 void adaptive_avg_pool2d_forward(const Tensor& X, int N, int C, int H, int W,
415 int H_out, int W_out, Tensor& Y) {
416 46 BROTENSOR_CHECK_SPATIAL_2D("adaptive_avg_pool2d_forward", X, N, C, H, W);
417 46 const auto& v = detail::dispatch(X, Y);
418
1/2
✓ Branch 0 taken 46 times.
✗ Branch 1 not taken.
 46 if (!v.adaptive_avg_pool2d_forward)
419 ✗ detail::throw_not_implemented("adaptive_avg_pool2d_forward", X.device);
420 46 detail::adopt_output(Y, X.device);
421 46 v.adaptive_avg_pool2d_forward(X, N, C, H, W, H_out, W_out, Y);
422 46 }
423 8 void adaptive_avg_pool2d_backward(const Tensor& dY, int N, int C, int H, int W,
424 int H_out, int W_out, Tensor& dX) {
425
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (N <= 0 || C <= 0 || H <= 0 || W <= 0) {
426 ✗ detail::throw_spatial_2d_error("adaptive_avg_pool2d_backward", "dY",
427 ✗ N, C, H, W, dY.rows, dY.cols);
428 }
429 8 BROTENSOR_CHECK_SPATIAL_2D("adaptive_avg_pool2d_backward", dY, N, C, H_out, W_out);
430 8 const auto& v = detail::dispatch(dY, dX);
431
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.adaptive_avg_pool2d_backward)
432 ✗ detail::throw_not_implemented("adaptive_avg_pool2d_backward", dY.device);
433 8 detail::adopt_output(dX, dY.device);
434 8 v.adaptive_avg_pool2d_backward(dY, N, C, H, W, H_out, W_out, dX);
435 8 }
436
437 23 void max_pool2d_forward(const Tensor& X, int N, int C, int H, int W,
438 int kH, int kW, int stride_h, int stride_w,
439 int pad_h, int pad_w, Tensor& Y, Tensor& Idx) {
440 23 BROTENSOR_CHECK_SPATIAL_2D("max_pool2d_forward", X, N, C, H, W);
441 23 const auto& v = detail::dispatch(X, Y, Idx);
442
1/2
✓ Branch 0 taken 23 times.
✗ Branch 1 not taken.
 23 if (!v.max_pool2d_forward)
443 ✗ detail::throw_not_implemented("max_pool2d_forward", X.device);
444 23 detail::adopt_output(Y, X.device);
445 23 detail::adopt_output(Idx, X.device);
446 46 v.max_pool2d_forward(X, N, C, H, W, kH, kW, stride_h, stride_w,
447 23 pad_h, pad_w, Y, Idx);
448 23 }
449 10 void max_pool2d_backward(const Tensor& dY, const Tensor& Idx,
450 int N, int C, int H, int W, int H_out, int W_out,
451 Tensor& dX) {
452
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (N <= 0 || C <= 0 || H <= 0 || W <= 0) {
453 ✗ detail::throw_spatial_2d_error("max_pool2d_backward", "dY",
454 ✗ N, C, H, W, dY.rows, dY.cols);
455 }
456 10 BROTENSOR_CHECK_SPATIAL_2D("max_pool2d_backward", dY, N, C, H_out, W_out);
457 10 const auto& v = detail::dispatch(dY, Idx, dX);
458
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.max_pool2d_backward)
459 ✗ detail::throw_not_implemented("max_pool2d_backward", dY.device);
460 10 detail::adopt_output(dX, dY.device);
461 10 v.max_pool2d_backward(dY, Idx, N, C, H, W, H_out, W_out, dX);
462 10 }
463
464 18 void gather_rows(const Tensor& X, const Tensor& Idx, Tensor& Y) {
465 18 const auto& v = detail::dispatch(X, Idx, Y);
466
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 if (!v.gather_rows) detail::throw_not_implemented("gather_rows", X.device);
467 18 detail::adopt_output(Y, X.device);
468 18 v.gather_rows(X, Idx, Y);
469 18 }
470 11 void scatter_rows_add(const Tensor& dY, const Tensor& Idx, int R, Tensor& dX) {
471 11 const auto& v = detail::dispatch(dY, Idx, dX);
472
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.scatter_rows_add)
473 ✗ detail::throw_not_implemented("scatter_rows_add", dY.device);
474 11 detail::adopt_output(dX, dY.device);
475 11 v.scatter_rows_add(dY, Idx, R, dX);
476 11 }
477 8 void scatter_rows(const Tensor& Y, const Tensor& Idx, Tensor& X) {
478 8 const auto& v = detail::dispatch(Y, Idx, X);
479
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.scatter_rows) detail::throw_not_implemented("scatter_rows", Y.device);
480 8 v.scatter_rows(Y, Idx, X);
481 8 }
482
483 90 void conv_transpose2d_forward(const Tensor& X, const Tensor& Wt,
484 const Tensor* bias,
485 int N, int C_in, int H, int W,
486 int C_out, int kH, int kW,
487 int stride_h, int stride_w,
488 int pad_h, int pad_w,
489 int output_padding_h, int output_padding_w,
490 int dil_h, int dil_w, int groups,
491 Tensor& Y) {
492 90 BROTENSOR_CHECK_SPATIAL_2D("conv_transpose2d_forward", X, N, C_in, H, W);
493 90 const auto& v = detail::dispatch(X, Wt, Y);
494
1/2
✓ Branch 0 taken 90 times.
✗ Branch 1 not taken.
 90 if (!v.conv_transpose2d_forward)
495 ✗ detail::throw_not_implemented("conv_transpose2d_forward", X.device);
496 90 detail::adopt_output(Y, X.device);
497 180 v.conv_transpose2d_forward(X, Wt, bias, N, C_in, H, W, C_out, kH, kW,
498 90 stride_h, stride_w, pad_h, pad_w,
499 90 output_padding_h, output_padding_w,
500 90 dil_h, dil_w, groups, Y);
501 90 }
502 10 void conv_transpose2d_backward_input(const Tensor& Wt, const Tensor& dY,
503 int N, int C_in, int H, int W,
504 int C_out, int kH, int kW,
505 int stride_h, int stride_w,
506 int pad_h, int pad_w,
507 int output_padding_h, int output_padding_w,
508 int dil_h, int dil_w, int groups,
509 Tensor& dX) {
510
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (N <= 0 || C_in <= 0 || H <= 0 || W <= 0 || C_out <= 0) {
511 ✗ detail::throw_spatial_2d_error("conv_transpose2d_backward_input", "dY",
512 ✗ N, C_in, H, W, dY.rows, dY.cols);
513 }
514 10 const auto& v = detail::dispatch(Wt, dY, dX);
515
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 if (!v.conv_transpose2d_backward_input)
516 ✗ detail::throw_not_implemented("conv_transpose2d_backward_input", Wt.device);
517 10 detail::adopt_output(dX, Wt.device);
518 20 v.conv_transpose2d_backward_input(Wt, dY, N, C_in, H, W, C_out, kH, kW,
519 10 stride_h, stride_w, pad_h, pad_w,
520 10 output_padding_h, output_padding_w,
521 10 dil_h, dil_w, groups, dX);
522 10 }
523 9 void conv_transpose2d_backward_weight(const Tensor& X, const Tensor& dY,
524 int N, int C_in, int H, int W,
525 int C_out, int kH, int kW,
526 int stride_h, int stride_w,
527 int pad_h, int pad_w,
528 int output_padding_h, int output_padding_w,
529 int dil_h, int dil_w, int groups,
530 Tensor& dWt) {
531 9 BROTENSOR_CHECK_SPATIAL_2D("conv_transpose2d_backward_weight", X, N, C_in, H, W);
532 9 const auto& v = detail::dispatch(X, dY, dWt);
533
1/2
✓ Branch 0 taken 9 times.
✗ Branch 1 not taken.
 9 if (!v.conv_transpose2d_backward_weight)
534 ✗ detail::throw_not_implemented("conv_transpose2d_backward_weight", X.device);
535 9 detail::adopt_output(dWt, X.device);
536 18 v.conv_transpose2d_backward_weight(X, dY, N, C_in, H, W, C_out, kH, kW,
537 9 stride_h, stride_w, pad_h, pad_w,
538 9 output_padding_h, output_padding_w,
539 9 dil_h, dil_w, groups, dWt);
540 9 }
541 7 void conv_transpose2d_backward_bias(const Tensor& dY, int N, int C_out,
542 int H_out, int W_out, Tensor& dB) {
543 7 BROTENSOR_CHECK_SPATIAL_2D("conv_transpose2d_backward_bias", dY, N, C_out, H_out, W_out);
544 7 const auto& v = detail::dispatch(dY, dB);
545
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (!v.conv_transpose2d_backward_bias)
546 ✗ detail::throw_not_implemented("conv_transpose2d_backward_bias", dY.device);
547 7 detail::adopt_output(dB, dY.device);
548 7 v.conv_transpose2d_backward_bias(dY, N, C_out, H_out, W_out, dB);
549 7 }
550
551 18 void window_partition_forward(const Tensor& X, int N, int C, int H, int W,
552 int window, Tensor& Y) {
553 18 BROTENSOR_CHECK_SPATIAL_2D("window_partition_forward", X, N, C, H, W);
554 18 const auto& v = detail::dispatch(X, Y);
555
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 if (!v.window_partition_forward)
556 ✗ detail::throw_not_implemented("window_partition_forward", X.device);
557 18 detail::adopt_output(Y, X.device);
558 18 v.window_partition_forward(X, N, C, H, W, window, Y);
559 18 }
560 12 void window_reverse_forward(const Tensor& X, int N, int C, int H, int W,
561 int window, Tensor& Y) {
562
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (N <= 0 || C <= 0 || H <= 0 || W <= 0) {
563 ✗ detail::throw_spatial_2d_error("window_reverse_forward", "X",
564 ✗ N, C, H, W, X.rows, X.cols);
565 }
566
3/6
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 12 times.
 12 if (window > 0 && H % window == 0 && W % window == 0) {
567 12 BROTENSOR_CHECK_SPATIAL_2D("window_reverse_forward", X,
568 N * (H / window) * (W / window), C, window, window);
569 12 }
570 12 const auto& v = detail::dispatch(X, Y);
571
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 if (!v.window_reverse_forward)
572 ✗ detail::throw_not_implemented("window_reverse_forward", X.device);
573 12 detail::adopt_output(Y, X.device);
574 12 v.window_reverse_forward(X, N, C, H, W, window, Y);
575 12 }
576
577 // ─── NCHW <-> sequence ─────────────────────────────────────────────────────
578
579 34 void nchw_to_sequence(const Tensor& X, int N, int C, int H, int W, Tensor& Y) {
580 34 BROTENSOR_CHECK_SPATIAL_2D("nchw_to_sequence", X, N, C, H, W);
581 34 const auto& v = detail::dispatch(X, Y);
582
1/2
✓ Branch 0 taken 34 times.
✗ Branch 1 not taken.
 34 if (!v.nchw_to_sequence) detail::throw_not_implemented("nchw_to_sequence", X.device);
583 34 detail::adopt_output(Y, X.device);
584 34 v.nchw_to_sequence(X, N, C, H, W, Y);
585 34 }
586
587 27 void sequence_to_nchw(const Tensor& X, int N, int C, int H, int W, Tensor& Y) {
588
1/2
✓ Branch 0 taken 27 times.
✗ Branch 1 not taken.
 27 if (N <= 0 || C <= 0 || H <= 0 || W <= 0) {
589 ✗ detail::throw_spatial_2d_error("sequence_to_nchw", "X",
590 ✗ N, C, H, W, X.rows, X.cols);
591 }
592 27 BROTENSOR_CHECK_SPATIAL_1D("sequence_to_nchw", X, N * H * W, C, 1);
593 27 const auto& v = detail::dispatch(X, Y);
594
1/2
✓ Branch 0 taken 27 times.
✗ Branch 1 not taken.
 27 if (!v.sequence_to_nchw) detail::throw_not_implemented("sequence_to_nchw", X.device);
595 27 detail::adopt_output(Y, X.device);
596 27 v.sequence_to_nchw(X, N, C, H, W, Y);
597 27 }
598
599 38 void spatial_merge_2x2_forward(const Tensor& X, int N, int C, int H, int W,
600 bool channel_major, Tensor& Y) {
601 38 BROTENSOR_CHECK_SPATIAL_2D("spatial_merge_2x2_forward", X, N, C, H, W);
602 38 const auto& v = detail::dispatch(X, Y);
603
1/2
✓ Branch 0 taken 38 times.
✗ Branch 1 not taken.
 38 if (!v.spatial_merge_2x2_forward)
604 ✗ detail::throw_not_implemented("spatial_merge_2x2_forward", X.device);
605 38 detail::adopt_output(Y, X.device);
606 38 v.spatial_merge_2x2_forward(X, N, C, H, W, channel_major, Y);
607 38 }
608
609 52 void pixel_shuffle_upsample_2x_forward(const Tensor& X, int N, int C_in,
610 int H, int W, int C_out, Tensor& Y) {
611
7/8
✓ Branch 0 taken 51 times.
✓ Branch 1 taken 1 time.
✓ Branch 2 taken 50 times.
✓ Branch 3 taken 1 time.
✓ Branch 4 taken 48 times.
✓ Branch 5 taken 2 times.
✗ Branch 6 not taken.
✓ Branch 7 taken 48 times.
 52 if (N > 0 && C_in > 0 && H > 0 && W > 0) {
612 48 BROTENSOR_CHECK_SPATIAL_2D("pixel_shuffle_upsample_2x_forward", X, N, C_in, H, W);
613 48 }
614 52 const auto& v = detail::dispatch(X, Y);
615
1/2
✓ Branch 0 taken 52 times.
✗ Branch 1 not taken.
 52 if (!v.pixel_shuffle_upsample_2x_forward)
616 ✗ detail::throw_not_implemented("pixel_shuffle_upsample_2x_forward", X.device);
617 52 detail::adopt_output(Y, X.device);
618 52 v.pixel_shuffle_upsample_2x_forward(X, N, C_in, H, W, C_out, Y);
619 52 }
620
621 16 void patch_unpack_forward(const Tensor& tokens, int hp, int wp, int P,
622 int C_total, int C_keep, bool channel_major,
623 Tensor& Y) {
624 16 const auto& v = detail::dispatch(tokens, Y);
625
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 if (!v.patch_unpack_forward)
626 ✗ detail::throw_not_implemented("patch_unpack_forward", tokens.device);
627 16 detail::adopt_output(Y, tokens.device);
628 16 v.patch_unpack_forward(tokens, hp, wp, P, C_total, C_keep, channel_major, Y);
629 16 }
630
631 // ─── ResBlock ──────────────────────────────────────────────────────────────
632
633 26 void resblock_forward(const Tensor& X,
634 const Tensor& gamma1, const Tensor& beta1,
635 const Tensor& W1, const Tensor* b1,
636 const Tensor* t_emb_shift,
637 const Tensor& gamma2, const Tensor& beta2,
638 const Tensor& W2, const Tensor* b2,
639 const Tensor* Wskip, const Tensor* bskip,
640 int N, int C_in, int C_out, int H, int W,
641 int num_groups, float eps, Tensor& Y) {
642 26 BROTENSOR_CHECK_SPATIAL_2D("resblock_forward", X, N, C_in, H, W);
643 26 const auto& v = detail::dispatch_with_opts(
644 104 X, gamma1, {&beta1, &W1, b1, t_emb_shift, &gamma2, &beta2, &W2, b2,
645 78 Wskip, bskip, &Y});
646
1/2
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
 26 if (!v.resblock_forward) detail::throw_not_implemented("resblock_forward", X.device);
647 26 detail::adopt_output(Y, X.device);
648 52 v.resblock_forward(X, gamma1, beta1, W1, b1, t_emb_shift,
649 26 gamma2, beta2, W2, b2, Wskip, bskip,
650 26 N, C_in, C_out, H, W, num_groups, eps, Y);
651 26 }
652
653 7 void resblock_forward_int8w_fp16(const Tensor& X,
654 const Tensor& gamma1, const Tensor& beta1,
655 const Tensor& W1_int8, const Tensor& s1,
656 const Tensor* b1,
657 const Tensor* t_emb_shift,
658 const Tensor& gamma2, const Tensor& beta2,
659 const Tensor& W2_int8, const Tensor& s2,
660 const Tensor* b2,
661 const Tensor* Wskip_int8, const Tensor* sskip,
662 const Tensor* bskip,
663 int N, int C_in, int C_out, int H, int W,
664 int num_groups, float eps, Tensor& Y) {
665 7 BROTENSOR_CHECK_SPATIAL_2D("resblock_forward_int8w_fp16", X, N, C_in, H, W);
666 7 const auto& v = detail::dispatch_with_opts(
667 70 X, gamma1, {&beta1, &W1_int8, &s1, b1, t_emb_shift,
668 35 &gamma2, &beta2, &W2_int8, &s2, b2,
669 28 Wskip_int8, sskip, bskip, &Y});
670
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (!v.resblock_forward_int8w_fp16)
671 ✗ detail::throw_not_implemented("resblock_forward_int8w_fp16", X.device);
672 7 detail::adopt_output(Y, X.device);
673 14 v.resblock_forward_int8w_fp16(X, gamma1, beta1, W1_int8, s1, b1, t_emb_shift,
674 7 gamma2, beta2, W2_int8, s2, b2,
675 7 Wskip_int8, sskip, bskip,
676 7 N, C_in, C_out, H, W, num_groups, eps, Y);
677 7 }
678
679 20 void resblock_backward(const Tensor& X,
680 const Tensor& gamma1, const Tensor& beta1,
681 const Tensor& W1, const Tensor* b1,
682 const Tensor* t_emb_shift,
683 const Tensor& gamma2, const Tensor& beta2,
684 const Tensor& W2, const Tensor* b2,
685 const Tensor* Wskip, const Tensor* bskip,
686 int N, int C_in, int C_out, int H, int W,
687 int num_groups, float eps,
688 const Tensor& dY,
689 Tensor& dX,
690 Tensor& dGamma1, Tensor& dBeta1,
691 Tensor& dW1, Tensor* db1,
692 Tensor* dt_emb_shift,
693 Tensor& dGamma2, Tensor& dBeta2,
694 Tensor& dW2, Tensor* db2,
695 Tensor* dWskip, Tensor* dbskip) {
696 20 BROTENSOR_CHECK_SPATIAL_2D("resblock_backward", X, N, C_in, H, W);
697 20 BROTENSOR_CHECK_SPATIAL_2D("resblock_backward", dY, N, C_out, H, W);
698 20 const auto& v = detail::dispatch_with_opts(
699 320 X, gamma1, {&beta1, &W1, b1, t_emb_shift, &gamma2, &beta2, &W2, b2,
700 160 Wskip, bskip, &dY, &dX, &dGamma1, &dBeta1, &dW1, db1,
701 140 dt_emb_shift, &dGamma2, &dBeta2, &dW2, db2, dWskip, dbskip});
702
1/2
✓ Branch 0 taken 20 times.
✗ Branch 1 not taken.
 20 if (!v.resblock_backward) detail::throw_not_implemented("resblock_backward", X.device);
703 20 detail::adopt_output(dX, X.device);
704 20 detail::adopt_output(dGamma1, X.device);
705 20 detail::adopt_output(dBeta1, X.device);
706 20 detail::adopt_output(dW1, X.device);
707
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 18 times.
 20 if (db1) detail::adopt_output(*db1, X.device);
708
2/2
✓ Branch 0 taken 10 times.
✓ Branch 1 taken 10 times.
 20 if (dt_emb_shift) detail::adopt_output(*dt_emb_shift, X.device);
709 20 detail::adopt_output(dGamma2, X.device);
710 20 detail::adopt_output(dBeta2, X.device);
711 20 detail::adopt_output(dW2, X.device);
712
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 18 times.
 20 if (db2) detail::adopt_output(*db2, X.device);
713
2/2
✓ Branch 0 taken 11 times.
✓ Branch 1 taken 9 times.
 20 if (dWskip) detail::adopt_output(*dWskip, X.device);
714
2/2
✓ Branch 0 taken 11 times.
✓ Branch 1 taken 9 times.
 20 if (dbskip) detail::adopt_output(*dbskip, X.device);
715 40 v.resblock_backward(X, gamma1, beta1, W1, b1, t_emb_shift,
716 20 gamma2, beta2, W2, b2, Wskip, bskip,
717 20 N, C_in, C_out, H, W, num_groups, eps,
718 20 dY, dX, dGamma1, dBeta1, dW1, db1, dt_emb_shift,
719 20 dGamma2, dBeta2, dW2, db2, dWskip, dbskip);
720 20 }
721
722 // ─── BatchNorm ─────────────────────────────────────────────────────────────
723
724 11 void batch_norm_forward(const Tensor& X,
725 const Tensor& gamma, const Tensor& beta,
726 Tensor& running_mean, Tensor& running_var,
727 int N, int C, int H, int W,
728 float eps, float momentum,
729 Tensor& Y,
730 Tensor& saved_mean, Tensor& saved_rstd) {
731 11 BROTENSOR_CHECK_SPATIAL_2D("batch_norm_forward", X, N, C, H, W);
732 22 const auto& v = detail::dispatch(X, gamma, beta, running_mean, running_var,
733 11 Y, saved_mean, saved_rstd);
734
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!v.batch_norm_forward)
735 ✗ detail::throw_not_implemented("batch_norm_forward", X.device);
736 11 detail::adopt_output(running_mean, X.device);
737 11 detail::adopt_output(running_var, X.device);
738 11 detail::adopt_output(Y, X.device);
739 11 detail::adopt_output(saved_mean, X.device);
740 11 detail::adopt_output(saved_rstd, X.device);
741 22 v.batch_norm_forward(X, gamma, beta, running_mean, running_var,
742 11 N, C, H, W, eps, momentum,
743 11 Y, saved_mean, saved_rstd);
744 11 }
745
746 17 void batch_norm_inference(const Tensor& X,
747 const Tensor& gamma, const Tensor& beta,
748 const Tensor& running_mean,
749 const Tensor& running_var,
750 int N, int C, int H, int W,
751 float eps,
752 Tensor& Y) {
753 17 BROTENSOR_CHECK_SPATIAL_2D("batch_norm_inference", X, N, C, H, W);
754 17 const auto& v = detail::dispatch(X, gamma, beta, running_mean, running_var, Y);
755
1/2
✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
 17 if (!v.batch_norm_inference)
756 ✗ detail::throw_not_implemented("batch_norm_inference", X.device);
757 17 detail::adopt_output(Y, X.device);
758 34 v.batch_norm_inference(X, gamma, beta, running_mean, running_var,
759 17 N, C, H, W, eps, Y);
760 17 }
761
762 7 void batch_norm_backward(const Tensor& X,
763 const Tensor& gamma,
764 const Tensor& saved_mean,
765 const Tensor& saved_rstd,
766 const Tensor& dY,
767 int N, int C, int H, int W,
768 Tensor& dX,
769 Tensor& dGamma, Tensor& dBeta) {
770 7 BROTENSOR_CHECK_SPATIAL_2D("batch_norm_backward", X, N, C, H, W);
771 7 BROTENSOR_CHECK_SPATIAL_2D("batch_norm_backward", dY, N, C, H, W);
772 14 const auto& v = detail::dispatch(X, gamma, saved_mean, saved_rstd, dY,
773 7 dX, dGamma, dBeta);
774
1/2
✓ Branch 0 taken 7 times.
✗ Branch 1 not taken.
 7 if (!v.batch_norm_backward)
775 ✗ detail::throw_not_implemented("batch_norm_backward", X.device);
776 7 detail::adopt_output(dX, X.device);
777 7 detail::adopt_output(dGamma, X.device);
778 7 detail::adopt_output(dBeta, X.device);
779 14 v.batch_norm_backward(X, gamma, saved_mean, saved_rstd, dY,
780 7 N, C, H, W, dX, dGamma, dBeta);
781 7 }
782
783 // ─── Image preprocessing helpers ───────────────────────────────────────────
784
785 8 void image_normalize(const Tensor& X,
786 const Tensor& mean, const Tensor& std_,
787 int N, int C, int H, int W,
788 Tensor& Y) {
789 8 BROTENSOR_CHECK_SPATIAL_2D("image_normalize", X, N, C, H, W);
790 8 const auto& v = detail::dispatch(X, mean, std_, Y);
791
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!v.image_normalize)
792 ✗ detail::throw_not_implemented("image_normalize", X.device);
793 8 detail::adopt_output(Y, X.device);
794 8 v.image_normalize(X, mean, std_, N, C, H, W, Y);
795 8 }
796
797 17 void image_u8_to_f32_nhwc_to_nchw(const uint8_t* src,
798 int N, int H, int W, int C,
799 float scale, float bias,
800 Tensor& Y) {
801 17 const auto& v = detail::dispatch(Y);
802
1/2
✓ Branch 0 taken 17 times.
✗ Branch 1 not taken.
 17 if (!v.image_u8_to_f32_nhwc_to_nchw)
803 ✗ detail::throw_not_implemented("image_u8_to_f32_nhwc_to_nchw", Y.device);
804 17 detail::adopt_output(Y, Y.device);
805 17 v.image_u8_to_f32_nhwc_to_nchw(src, N, H, W, C, scale, bias, Y);
806 17 }
807
808 // ─── StyleGAN3 bias_act / upfirdn2d / modulated_conv2d ─────────────────────
809
810 65 void bias_act_forward(const Tensor& X, const Tensor* b,
811 int N, int C, int HW, int act, float alpha,
812 float gain, float clamp, Tensor& Y) {
813 65 BROTENSOR_CHECK_SPATIAL_1D("bias_act_forward", X, N, C, HW);
814 65 const auto& v = detail::dispatch_with_opts(X, {b, &Y});
815
1/2
✓ Branch 0 taken 65 times.
✗ Branch 1 not taken.
 65 if (!v.bias_act_forward) detail::throw_not_implemented("bias_act_forward", X.device);
816 65 detail::adopt_output(Y, X.device);
817 65 v.bias_act_forward(X, b, N, C, HW, act, alpha, gain, clamp, Y);
818 65 }
819
820 68 void bias_act_backward(const Tensor& dY, const Tensor& X, const Tensor* b,
821 int N, int C, int HW, int act, float alpha,
822 float gain, float clamp, Tensor& dX, Tensor* dB) {
823 68 BROTENSOR_CHECK_SPATIAL_1D("bias_act_backward", dY, N, C, HW);
824 68 BROTENSOR_CHECK_SPATIAL_1D("bias_act_backward", X, N, C, HW);
825 68 const auto& v = detail::dispatch_with_opts(dY, X, {b, dB, &dX});
826
1/2
✓ Branch 0 taken 68 times.
✗ Branch 1 not taken.
 68 if (!v.bias_act_backward) detail::throw_not_implemented("bias_act_backward", dY.device);
827 68 detail::adopt_output(dX, dY.device);
828
2/2
✓ Branch 0 taken 36 times.
✓ Branch 1 taken 32 times.
 68 if (dB) detail::adopt_output(*dB, dY.device);
829 68 v.bias_act_backward(dY, X, b, N, C, HW, act, alpha, gain, clamp, dX, dB);
830 68 }
831
832 50 void upfirdn2d_forward(const Tensor& X, const Tensor& f,
833 int N, int C, int H, int Wd, int fH, int fW,
834 int up_x, int up_y, int down_x, int down_y,
835 int pad_x0, int pad_x1, int pad_y0, int pad_y1,
836 bool flip_filter, float gain, Tensor& Y) {
837 50 BROTENSOR_CHECK_SPATIAL_2D("upfirdn2d_forward", X, N, C, H, Wd);
838 50 const auto& v = detail::dispatch(X, f, Y);
839
1/2
✓ Branch 0 taken 50 times.
✗ Branch 1 not taken.
 50 if (!v.upfirdn2d_forward) detail::throw_not_implemented("upfirdn2d_forward", X.device);
840 50 detail::adopt_output(Y, X.device);
841 100 v.upfirdn2d_forward(X, f, N, C, H, Wd, fH, fW, up_x, up_y, down_x, down_y,
842 50 pad_x0, pad_x1, pad_y0, pad_y1, flip_filter, gain, Y);
843 50 }
844
845 47 void upfirdn2d_backward(const Tensor& dY, const Tensor& f,
846 int N, int C, int H, int Wd, int fH, int fW,
847 int up_x, int up_y, int down_x, int down_y,
848 int pad_x0, int pad_x1, int pad_y0, int pad_y1,
849 bool flip_filter, float gain, Tensor& dX) {
850
1/2
✓ Branch 0 taken 47 times.
✗ Branch 1 not taken.
 47 if (N <= 0 || C <= 0 || H <= 0 || Wd <= 0) {
851 ✗ detail::throw_spatial_2d_error("upfirdn2d_backward", "dY",
852 ✗ N, C, H, Wd, dY.rows, dY.cols);
853 }
854 47 const auto& v = detail::dispatch(dY, f, dX);
855
1/2
✓ Branch 0 taken 47 times.
✗ Branch 1 not taken.
 47 if (!v.upfirdn2d_backward) detail::throw_not_implemented("upfirdn2d_backward", dY.device);
856 47 detail::adopt_output(dX, dY.device);
857 94 v.upfirdn2d_backward(dY, f, N, C, H, Wd, fH, fW, up_x, up_y, down_x, down_y,
858 47 pad_x0, pad_x1, pad_y0, pad_y1, flip_filter, gain, dX);
859 47 }
860
861 32 void modulated_conv2d_forward(const Tensor& X, const Tensor& W, const Tensor& s,
862 int N, int C_in, int H, int Wd,
863 int C_out, int kH, int kW,
864 int pad_h, int pad_w,
865 bool demodulate, float eps,
866 Tensor& dcoef, Tensor& Y) {
867 32 BROTENSOR_CHECK_SPATIAL_2D("modulated_conv2d_forward", X, N, C_in, H, Wd);
868 32 const auto& v = detail::dispatch(X, W, s, dcoef, Y);
869
1/2
✓ Branch 0 taken 32 times.
✗ Branch 1 not taken.
 32 if (!v.modulated_conv2d_forward)
870 ✗ detail::throw_not_implemented("modulated_conv2d_forward", X.device);
871 32 detail::adopt_output(dcoef, X.device);
872 32 detail::adopt_output(Y, X.device);
873 64 v.modulated_conv2d_forward(X, W, s, N, C_in, H, Wd, C_out, kH, kW,
874 32 pad_h, pad_w, demodulate, eps, dcoef, Y);
875 32 }
876
877 42 void modulated_conv2d_backward(const Tensor& X, const Tensor& W, const Tensor& s,
878 const Tensor& dcoef, const Tensor& dY,
879 int N, int C_in, int H, int Wd,
880 int C_out, int kH, int kW,
881 int pad_h, int pad_w, bool demodulate, float eps,
882 Tensor& dX, Tensor& dW, Tensor& ds) {
883 42 BROTENSOR_CHECK_SPATIAL_2D("modulated_conv2d_backward", X, N, C_in, H, Wd);
884 42 BROTENSOR_CHECK_SPATIAL_2D("modulated_conv2d_backward", dY, N, C_out, H, Wd);
885 42 const auto& v = detail::dispatch(X, W, s, dcoef, dY, dX, dW, ds);
886
1/2
✓ Branch 0 taken 42 times.
✗ Branch 1 not taken.
 42 if (!v.modulated_conv2d_backward)
887 ✗ detail::throw_not_implemented("modulated_conv2d_backward", X.device);
888 42 detail::adopt_output(dX, X.device);
889 42 detail::adopt_output(dW, X.device);
890 42 detail::adopt_output(ds, X.device);
891 84 v.modulated_conv2d_backward(X, W, s, dcoef, dY, N, C_in, H, Wd, C_out, kH, kW,
892 42 pad_h, pad_w, demodulate, eps, dX, dW, ds);
893 42 }
894
895 } // namespace brotensor
896