src/cpu/upfirdn2d.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // ─── CPU upfirdn2d (StyleGAN3-R) ──────────────────────────────────────────── | ||
| 2 | // | ||
| 3 | // Upsample (zero-insert) → pad/crop → 2D FIR correlation → downsample → gain. | ||
| 4 | // FP32 reference mirroring NVlabs `_upfirdn2d_ref` (general non-separable 2D | ||
| 5 | // path — the one config-R's radial filters need). The filter is a constant | ||
| 6 | // shared across channels (depthwise); there is no gradient to it. | ||
| 7 | // | ||
| 8 | // All of forward and backward funnel through `upfirdn2d_run`: the op is linear | ||
| 9 | // in X, so the backward is itself a forward call with up/down swapped, the | ||
| 10 | // filter flip inverted, and padding recomputed (mirrors `_upfirdn2d_cuda`). | ||
| 11 | // | ||
| 12 | // Index relation (per output pixel oh,ow and filter tap kh,kw), where the | ||
| 13 | // downsample picks conv row oh*down_y: | ||
| 14 | // padded row py = oh*down_y + kh ; upsampled row uy = py - pad_y0 | ||
| 15 | // contributes iff 0<=uy<H*up_y and uy % up_y == 0, reading X[iy=uy/up_y]. | ||
| 16 | // Effective filter weight = flip_filter ? f[kh,kw] : f[fH-1-kh, fW-1-kw] | ||
| 17 | // (the flip makes the default flip_filter=false a true convolution). | ||
| 18 | // | ||
| 19 | // Output dims: H_out = (H*up_y + pad_y0 + pad_y1 - fH)/down_y + 1, i.e. | ||
| 20 | // ceil(conv_valid_len / down_y) — equal to the slicing x[::down_y] count. | ||
| 21 | |||
| 22 | #include <brotensor/tensor.h> | ||
| 23 | #include <brotensor/detail/cpu/thread_pool.h> | ||
| 24 | |||
| 25 | #include <algorithm> | ||
| 26 | #include <stdexcept> | ||
| 27 | #include <string> | ||
| 28 | |||
| 29 | namespace brotensor::detail::cpu { | ||
| 30 | |||
| 31 | namespace { | ||
| 32 | |||
| 33 | 124 | inline void check_fp32(const ::brotensor::Tensor& t, | |
| 34 | const char* op, const char* name) { | ||
| 35 |
1/2✓ Branch 0 taken 124 times.
✗ Branch 1 not taken.
|
124 | if (t.dtype != Dtype::FP32) { |
| 36 | ✗ | throw std::runtime_error(std::string("brotensor: ") + op + ": " + | |
| 37 | ✗ | name + " must be FP32 (CPU backend is " | |
| 38 | "FP32-only)"); | ||
| 39 | } | ||
| 40 | 124 | } | |
| 41 | |||
| 42 | // Shared engine. `In` is (N, C*Hin*Win); resizes `Out` to (N, C*Hout*Wout). | ||
| 43 | 62 | void upfirdn2d_run(const ::brotensor::Tensor& In, int N, int C, int Hin, int Win, | |
| 44 | const ::brotensor::Tensor& f, int fH, int fW, | ||
| 45 | int up_x, int up_y, int down_x, int down_y, | ||
| 46 | int px0, int px1, int py0, int py1, | ||
| 47 | bool flip_filter, float gain, | ||
| 48 | ::brotensor::Tensor& Out, const char* op) { | ||
| 49 | 62 | check_fp32(In, op, "input"); | |
| 50 | 62 | check_fp32(f, op, "f"); | |
| 51 |
1/2✓ Branch 0 taken 62 times.
✗ Branch 1 not taken.
|
62 | if (up_x < 1 || up_y < 1 || down_x < 1 || down_y < 1) { |
| 52 | ✗ | throw std::runtime_error(std::string("brotensor: ") + op + | |
| 53 | ": up/down factors must be >= 1"); | ||
| 54 | } | ||
| 55 |
1/2✓ Branch 0 taken 62 times.
✗ Branch 1 not taken.
|
62 | if (In.rows != N || In.cols != C * Hin * Win) { |
| 56 | ✗ | throw std::runtime_error(std::string("brotensor: ") + op + | |
| 57 | ": input shape mismatch"); | ||
| 58 | } | ||
| 59 |
1/2✓ Branch 0 taken 62 times.
✗ Branch 1 not taken.
|
62 | if (f.rows != fH || f.cols != fW) { |
| 60 | ✗ | throw std::runtime_error(std::string("brotensor: ") + op + | |
| 61 | ": filter shape mismatch"); | ||
| 62 | } | ||
| 63 | 62 | const int Hu = Hin * up_y, Wu = Win * up_x; | |
| 64 | 62 | const int Hp = Hu + py0 + py1, Wp = Wu + px0 + px1; | |
| 65 |
1/2✓ Branch 0 taken 62 times.
✗ Branch 1 not taken.
|
62 | if (Hp < fH || Wp < fW) { |
| 66 | ✗ | throw std::runtime_error(std::string("brotensor: ") + op + | |
| 67 | ": padded input smaller than filter"); | ||
| 68 | } | ||
| 69 | 62 | const int Hc = Hp - fH + 1, Wc = Wp - fW + 1; | |
| 70 | 62 | const int Hout = (Hc - 1) / down_y + 1; | |
| 71 | 62 | const int Wout = (Wc - 1) / down_x + 1; | |
| 72 | 62 | const int out_cols = C * Hout * Wout; | |
| 73 |
1/6✗ Branch 0 not taken.
✓ Branch 1 taken 62 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
62 | if (Out.rows != N || Out.cols != out_cols || Out.dtype != Dtype::FP32) { |
| 74 | 62 | Out.resize(N, out_cols, Dtype::FP32); | |
| 75 | 62 | } | |
| 76 |
2/4✓ Branch 0 taken 62 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 62 times.
|
62 | if (N == 0 || out_cols == 0) return; |
| 77 | |||
| 78 | 62 | const float* Ip = In.host_f32(); | |
| 79 | 62 | const float* Fp = f.host_f32(); | |
| 80 | 62 | float* Op = Out.host_f32_mut(); | |
| 81 | |||
| 82 | // One (n, c) plane per task. Only taps landing on a real (non-zero- | ||
| 83 | // inserted, in-range) input sample contribute, so kh/kw start at the first | ||
| 84 | // such tap and step by the up factor: the same taps in the same order as a | ||
| 85 | // full scan (bit-identical), up_x*up_y fewer iterations. | ||
| 86 |
1/2✓ Branch 0 taken 62 times.
✗ Branch 1 not taken.
|
257 | parallel_for(static_cast<std::size_t>(N) * C, [&](std::size_t nc) { |
| 87 | 195 | const size_t in_base = nc * Hin * Win; | |
| 88 | 195 | const size_t out_base = nc * Hout * Wout; | |
| 89 |
2/2✓ Branch 0 taken 2646 times.
✓ Branch 1 taken 195 times.
|
2841 | for (int oh = 0; oh < Hout; ++oh) { |
| 90 | 2646 | const int by = oh * down_y - py0; | |
| 91 | 2646 | int kh0 = std::max(0, -by); | |
| 92 |
2/2✓ Branch 0 taken 1942 times.
✓ Branch 1 taken 704 times.
|
2646 | if (const int r = (by + kh0) % up_y) kh0 += up_y - r; |
| 93 | 2646 | const int kh1 = std::min(fH, Hu - by); | |
| 94 |
2/2✓ Branch 0 taken 49395 times.
✓ Branch 1 taken 2646 times.
|
52041 | for (int ow = 0; ow < Wout; ++ow) { |
| 95 | 49395 | const int bx = ow * down_x - px0; | |
| 96 | 49395 | int kw0 = std::max(0, -bx); | |
| 97 |
2/2✓ Branch 0 taken 32252 times.
✓ Branch 1 taken 17143 times.
|
49395 | if (const int r = (bx + kw0) % up_x) kw0 += up_x - r; |
| 98 | 49395 | const int kw1 = std::min(fW, Wu - bx); | |
| 99 | 49395 | float acc = 0.0f; | |
| 100 |
2/2✓ Branch 0 taken 158015 times.
✓ Branch 1 taken 49395 times.
|
207410 | for (int kh = kh0; kh < kh1; kh += up_y) { |
| 101 | 158015 | const float* xr = Ip + in_base + static_cast<size_t>((by + kh) / up_y) * Win; | |
| 102 |
2/2✓ Branch 0 taken 73958 times.
✓ Branch 1 taken 84057 times.
|
158015 | const float* fr = Fp + static_cast<size_t>(flip_filter ? kh : (fH - 1 - kh)) * fW; |
| 103 |
2/2✓ Branch 0 taken 779809 times.
✓ Branch 1 taken 158015 times.
|
937824 | for (int kw = kw0; kw < kw1; kw += up_x) { |
| 104 |
2/2✓ Branch 0 taken 374352 times.
✓ Branch 1 taken 405457 times.
|
779809 | acc += xr[(bx + kw) / up_x] * fr[flip_filter ? kw : (fW - 1 - kw)]; |
| 105 | 779809 | } | |
| 106 | 158015 | } | |
| 107 | 49395 | Op[out_base + static_cast<size_t>(oh) * Wout + ow] = acc * gain; | |
| 108 | 49395 | } | |
| 109 | 2646 | } | |
| 110 | 195 | }); | |
| 111 | 62 | } | |
| 112 | |||
| 113 | // Forward output height/width for the given params (shared by the public | ||
| 114 | // forward and the backward's padding recompute). | ||
| 115 | 60 | inline int out_dim(int in, int up, int down, int pad0, int pad1, int fdim) { | |
| 116 | 60 | return (in * up + pad0 + pad1 - fdim) / down + 1; | |
| 117 | } | ||
| 118 | |||
| 119 | } // namespace | ||
| 120 | |||
| 121 | 32 | void upfirdn2d_forward(const ::brotensor::Tensor& X, const ::brotensor::Tensor& f, | |
| 122 | int N, int C, int H, int Wd, int fH, int fW, | ||
| 123 | int up_x, int up_y, int down_x, int down_y, | ||
| 124 | int pad_x0, int pad_x1, int pad_y0, int pad_y1, | ||
| 125 | bool flip_filter, float gain, ::brotensor::Tensor& Y) { | ||
| 126 | 64 | upfirdn2d_run(X, N, C, H, Wd, f, fH, fW, | |
| 127 | 32 | up_x, up_y, down_x, down_y, | |
| 128 | 32 | pad_x0, pad_x1, pad_y0, pad_y1, | |
| 129 | 32 | flip_filter, gain, Y, "upfirdn2d_forward"); | |
| 130 | 32 | } | |
| 131 | |||
| 132 | 30 | void upfirdn2d_backward(const ::brotensor::Tensor& dY, const ::brotensor::Tensor& f, | |
| 133 | int N, int C, int H, int Wd, int fH, int fW, | ||
| 134 | int up_x, int up_y, int down_x, int down_y, | ||
| 135 | int pad_x0, int pad_x1, int pad_y0, int pad_y1, | ||
| 136 | bool flip_filter, float gain, ::brotensor::Tensor& dX) { | ||
| 137 | // Forward output dims (dY's spatial extent). | ||
| 138 | 30 | const int Hout = out_dim(H, up_y, down_y, pad_y0, pad_y1, fH); | |
| 139 | 30 | const int Wout = out_dim(Wd, up_x, down_x, pad_x0, pad_x1, fW); | |
| 140 | // Backward padding (NVlabs _upfirdn2d_cuda): swap up<->down, flip the flip. | ||
| 141 | 30 | const int p_x0 = fW - pad_x0 - 1; | |
| 142 | 30 | const int p_x1 = Wd * up_x - Wout * down_x + pad_x0 - up_x + 1; | |
| 143 | 30 | const int p_y0 = fH - pad_y0 - 1; | |
| 144 | 30 | const int p_y1 = H * up_y - Hout * down_y + pad_y0 - up_y + 1; | |
| 145 | 60 | upfirdn2d_run(dY, N, C, Hout, Wout, f, fH, fW, | |
| 146 | 30 | /*up=*/down_x, down_y, /*down=*/up_x, up_y, | |
| 147 | 30 | p_x0, p_x1, p_y0, p_y1, | |
| 148 | 30 | !flip_filter, gain, dX, "upfirdn2d_backward"); | |
| 149 |
1/2✓ Branch 0 taken 30 times.
✗ Branch 1 not taken.
|
30 | if (dX.rows != N || dX.cols != C * H * Wd) { |
| 150 | ✗ | throw std::runtime_error("upfirdn2d_backward: internal dX shape " | |
| 151 | "mismatch (param inconsistency)"); | ||
| 152 | } | ||
| 153 | 30 | } | |
| 154 | |||
| 155 | } // namespace brotensor::detail::cpu | ||
| 156 |