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 | |||
| 24 | #include <stdexcept> | ||
| 25 | #include <string> | ||
| 26 | |||
| 27 | namespace brotensor::detail::cpu { | ||
| 28 | |||
| 29 | namespace { | ||
| 30 | |||
| 31 | 76 | inline void check_fp32(const ::brotensor::Tensor& t, | |
| 32 | const char* op, const char* name) { | ||
| 33 |
1/2✓ Branch 0 taken 76 times.
✗ Branch 1 not taken.
|
76 | if (t.dtype != Dtype::FP32) { |
| 34 | ✗ | throw std::runtime_error(std::string("brotensor: ") + op + ": " + | |
| 35 | ✗ | name + " must be FP32 (CPU backend is " | |
| 36 | "FP32-only)"); | ||
| 37 | } | ||
| 38 | 76 | } | |
| 39 | |||
| 40 | // Shared engine. `In` is (N, C*Hin*Win); resizes `Out` to (N, C*Hout*Wout). | ||
| 41 | 38 | void upfirdn2d_run(const ::brotensor::Tensor& In, int N, int C, int Hin, int Win, | |
| 42 | const ::brotensor::Tensor& f, int fH, int fW, | ||
| 43 | int up_x, int up_y, int down_x, int down_y, | ||
| 44 | int px0, int px1, int py0, int py1, | ||
| 45 | bool flip_filter, float gain, | ||
| 46 | ::brotensor::Tensor& Out, const char* op) { | ||
| 47 | 38 | check_fp32(In, op, "input"); | |
| 48 | 38 | check_fp32(f, op, "f"); | |
| 49 |
1/2✓ Branch 0 taken 38 times.
✗ Branch 1 not taken.
|
38 | if (up_x < 1 || up_y < 1 || down_x < 1 || down_y < 1) { |
| 50 | ✗ | throw std::runtime_error(std::string("brotensor: ") + op + | |
| 51 | ": up/down factors must be >= 1"); | ||
| 52 | } | ||
| 53 |
1/2✓ Branch 0 taken 38 times.
✗ Branch 1 not taken.
|
38 | if (In.rows != N || In.cols != C * Hin * Win) { |
| 54 | ✗ | throw std::runtime_error(std::string("brotensor: ") + op + | |
| 55 | ": input shape mismatch"); | ||
| 56 | } | ||
| 57 |
1/2✓ Branch 0 taken 38 times.
✗ Branch 1 not taken.
|
38 | if (f.rows != fH || f.cols != fW) { |
| 58 | ✗ | throw std::runtime_error(std::string("brotensor: ") + op + | |
| 59 | ": filter shape mismatch"); | ||
| 60 | } | ||
| 61 | 38 | const int Hu = Hin * up_y, Wu = Win * up_x; | |
| 62 | 38 | const int Hp = Hu + py0 + py1, Wp = Wu + px0 + px1; | |
| 63 |
1/2✓ Branch 0 taken 38 times.
✗ Branch 1 not taken.
|
38 | if (Hp < fH || Wp < fW) { |
| 64 | ✗ | throw std::runtime_error(std::string("brotensor: ") + op + | |
| 65 | ": padded input smaller than filter"); | ||
| 66 | } | ||
| 67 | 38 | const int Hc = Hp - fH + 1, Wc = Wp - fW + 1; | |
| 68 | 38 | const int Hout = (Hc - 1) / down_y + 1; | |
| 69 | 38 | const int Wout = (Wc - 1) / down_x + 1; | |
| 70 | 38 | const int out_cols = C * Hout * Wout; | |
| 71 |
1/6✗ Branch 0 not taken.
✓ Branch 1 taken 38 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
38 | if (Out.rows != N || Out.cols != out_cols || Out.dtype != Dtype::FP32) { |
| 72 | 38 | Out.resize(N, out_cols, Dtype::FP32); | |
| 73 | 38 | } | |
| 74 |
2/4✓ Branch 0 taken 38 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 38 times.
|
38 | if (N == 0 || out_cols == 0) return; |
| 75 | |||
| 76 | 38 | const float* Ip = In.host_f32(); | |
| 77 | 38 | const float* Fp = f.host_f32(); | |
| 78 | 38 | float* Op = Out.host_f32_mut(); | |
| 79 | |||
| 80 |
2/2✓ Branch 0 taken 38 times.
✓ Branch 1 taken 58 times.
|
96 | for (int n = 0; n < N; ++n) { |
| 81 |
2/2✓ Branch 0 taken 134 times.
✓ Branch 1 taken 58 times.
|
192 | for (int c = 0; c < C; ++c) { |
| 82 | 134 | const size_t in_base = (static_cast<size_t>(n) * C + c) * Hin * Win; | |
| 83 | 134 | const size_t out_base = (static_cast<size_t>(n) * C + c) * Hout * Wout; | |
| 84 |
2/2✓ Branch 0 taken 1194 times.
✓ Branch 1 taken 134 times.
|
1328 | for (int oh = 0; oh < Hout; ++oh) { |
| 85 | 1194 | const int py_base = oh * down_y; | |
| 86 |
2/2✓ Branch 0 taken 13154 times.
✓ Branch 1 taken 1194 times.
|
14348 | for (int ow = 0; ow < Wout; ++ow) { |
| 87 | 13154 | const int px_base = ow * down_x; | |
| 88 | 13154 | float acc = 0.0f; | |
| 89 |
2/2✓ Branch 0 taken 69884 times.
✓ Branch 1 taken 13154 times.
|
83038 | for (int kh = 0; kh < fH; ++kh) { |
| 90 | 69884 | const int uy = py_base + kh - py0; | |
| 91 |
6/6✓ Branch 0 taken 62221 times.
✓ Branch 1 taken 7663 times.
✓ Branch 2 taken 55794 times.
✓ Branch 3 taken 6427 times.
✓ Branch 4 taken 23106 times.
✓ Branch 5 taken 32688 times.
|
69884 | if (uy < 0 || uy >= Hu || (uy % up_y) != 0) continue; |
| 92 | 32688 | const int iy = uy / up_y; | |
| 93 |
2/2✓ Branch 0 taken 12654 times.
✓ Branch 1 taken 20034 times.
|
32688 | const int frow = flip_filter ? kh : (fH - 1 - kh); |
| 94 |
2/2✓ Branch 0 taken 210750 times.
✓ Branch 1 taken 32688 times.
|
243438 | for (int kw = 0; kw < fW; ++kw) { |
| 95 | 210750 | const int ux = px_base + kw - px0; | |
| 96 |
6/6✓ Branch 0 taken 188659 times.
✓ Branch 1 taken 22091 times.
✓ Branch 2 taken 169602 times.
✓ Branch 3 taken 19057 times.
✓ Branch 4 taken 58662 times.
✓ Branch 5 taken 110940 times.
|
210750 | if (ux < 0 || ux >= Wu || (ux % up_x) != 0) continue; |
| 97 | 110940 | const int ix = ux / up_x; | |
| 98 |
2/2✓ Branch 0 taken 50424 times.
✓ Branch 1 taken 60516 times.
|
110940 | const int fcol = flip_filter ? kw : (fW - 1 - kw); |
| 99 | 221880 | acc += Ip[in_base + static_cast<size_t>(iy) * Win + ix] * | |
| 100 | 110940 | Fp[static_cast<size_t>(frow) * fW + fcol]; | |
| 101 | 110940 | } | |
| 102 | 32688 | } | |
| 103 | 13154 | Op[out_base + static_cast<size_t>(oh) * Wout + ow] = acc * gain; | |
| 104 | 13154 | } | |
| 105 | 1194 | } | |
| 106 | 134 | } | |
| 107 | 58 | } | |
| 108 | 38 | } | |
| 109 | |||
| 110 | // Forward output height/width for the given params (shared by the public | ||
| 111 | // forward and the backward's padding recompute). | ||
| 112 | 36 | inline int out_dim(int in, int up, int down, int pad0, int pad1, int fdim) { | |
| 113 | 36 | return (in * up + pad0 + pad1 - fdim) / down + 1; | |
| 114 | } | ||
| 115 | |||
| 116 | } // namespace | ||
| 117 | |||
| 118 | 20 | void upfirdn2d_forward(const ::brotensor::Tensor& X, const ::brotensor::Tensor& f, | |
| 119 | int N, int C, int H, int Wd, int fH, int fW, | ||
| 120 | int up_x, int up_y, int down_x, int down_y, | ||
| 121 | int pad_x0, int pad_x1, int pad_y0, int pad_y1, | ||
| 122 | bool flip_filter, float gain, ::brotensor::Tensor& Y) { | ||
| 123 | 40 | upfirdn2d_run(X, N, C, H, Wd, f, fH, fW, | |
| 124 | 20 | up_x, up_y, down_x, down_y, | |
| 125 | 20 | pad_x0, pad_x1, pad_y0, pad_y1, | |
| 126 | 20 | flip_filter, gain, Y, "upfirdn2d_forward"); | |
| 127 | 20 | } | |
| 128 | |||
| 129 | 18 | void upfirdn2d_backward(const ::brotensor::Tensor& dY, const ::brotensor::Tensor& f, | |
| 130 | int N, int C, int H, int Wd, int fH, int fW, | ||
| 131 | int up_x, int up_y, int down_x, int down_y, | ||
| 132 | int pad_x0, int pad_x1, int pad_y0, int pad_y1, | ||
| 133 | bool flip_filter, float gain, ::brotensor::Tensor& dX) { | ||
| 134 | // Forward output dims (dY's spatial extent). | ||
| 135 | 18 | const int Hout = out_dim(H, up_y, down_y, pad_y0, pad_y1, fH); | |
| 136 | 18 | const int Wout = out_dim(Wd, up_x, down_x, pad_x0, pad_x1, fW); | |
| 137 | // Backward padding (NVlabs _upfirdn2d_cuda): swap up<->down, flip the flip. | ||
| 138 | 18 | const int p_x0 = fW - pad_x0 - 1; | |
| 139 | 18 | const int p_x1 = Wd * up_x - Wout * down_x + pad_x0 - up_x + 1; | |
| 140 | 18 | const int p_y0 = fH - pad_y0 - 1; | |
| 141 | 18 | const int p_y1 = H * up_y - Hout * down_y + pad_y0 - up_y + 1; | |
| 142 | 36 | upfirdn2d_run(dY, N, C, Hout, Wout, f, fH, fW, | |
| 143 | 18 | /*up=*/down_x, down_y, /*down=*/up_x, up_y, | |
| 144 | 18 | p_x0, p_x1, p_y0, p_y1, | |
| 145 | 18 | !flip_filter, gain, dX, "upfirdn2d_backward"); | |
| 146 |
1/2✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
|
18 | if (dX.rows != N || dX.cols != C * H * Wd) { |
| 147 | ✗ | throw std::runtime_error("upfirdn2d_backward: internal dX shape " | |
| 148 | "mismatch (param inconsistency)"); | ||
| 149 | } | ||
| 150 | 18 | } | |
| 151 | |||
| 152 | } // namespace brotensor::detail::cpu | ||
| 153 |