src/cpu/convex_upsample.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // ─── CPU convex (mask-based) upsample, NCHW ───────────────────────────────── | ||
| 2 | // | ||
| 3 | // FP32 scalar host implementation of convex_upsample_forward — the RAFT-style | ||
| 4 | // "learned convex combination" upsampler used by optical-flow, stereo, and | ||
| 5 | // surface-normal refinement (DSINE up_prob_head). Each low-res pixel is | ||
| 6 | // expanded to a scale×scale block; every fine pixel in the block is a softmax- | ||
| 7 | // weighted blend of the 3×3 low-res neighborhood around its source pixel: | ||
| 8 | // | ||
| 9 | // Y[n, c, k*y+sy, k*x+sx] = sum_{m=0..8} W[n, m, sy, sx, y, x] * X[n, c, ny, nx] | ||
| 10 | // W = softmax over the 9 neighbors m of Mask[n, m, sy, sx, y, x] | ||
| 11 | // neighbor m: ny = clamp(y - 1 + m/3), nx = clamp(x - 1 + m%3) (replicate pad) | ||
| 12 | // | ||
| 13 | // Mask layout (matches torch view (N, 9, k, k, H, W)): | ||
| 14 | // flat channel = ((m*k + sy)*k + sx), so | ||
| 15 | // Mask[n, m, sy, sx, y, x] = mask[n, (m*k*k + sy*k + sx)*HW + (y*W + x)]. | ||
| 16 | // k = scale; the 9 axis is the 3×3 spatial neighborhood (m = my*3 + mx). | ||
| 17 | // | ||
| 18 | // X: (N, C*H*W). | ||
| 19 | // Mask: (N, 9*k*k*H*W). | ||
| 20 | // Y: (N, C*(k*H)*(k*W)), resized + dtype-set to X. | ||
| 21 | // | ||
| 22 | // Softmax accumulates in double. NOTE: the softmax over the 9 neighbors is the | ||
| 23 | // same for every channel, but we recompute it per (channel, fine-pixel) for | ||
| 24 | // implementation simplicity — fine at the C/scale this targets (DSINE C=3, | ||
| 25 | // k=8). Y OVERWRITTEN. Inference-only: no backward. | ||
| 26 | |||
| 27 | #include <brotensor/tensor.h> | ||
| 28 | |||
| 29 | #include <algorithm> | ||
| 30 | #include <cmath> | ||
| 31 | #include <stdexcept> | ||
| 32 | #include <string> | ||
| 33 | |||
| 34 | namespace brotensor::detail::cpu { | ||
| 35 | |||
| 36 | namespace { | ||
| 37 | |||
| 38 | ✗ | [[noreturn]] inline void fail(const char* op, const std::string& reason) { | |
| 39 | ✗ | throw std::runtime_error(std::string("brotensor: ") + op + ": " + reason); | |
| 40 | ✗ | } | |
| 41 | |||
| 42 | 10 | inline void check_fp32(const ::brotensor::Tensor& t, | |
| 43 | const char* op, const char* name) { | ||
| 44 |
1/2✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
|
10 | if (t.dtype != Dtype::FP32) |
| 45 | ✗ | fail(op, std::string(name) + " must be FP32 (CPU backend is FP32-only)"); | |
| 46 | 10 | } | |
| 47 | |||
| 48 |
4/4✓ Branch 0 taken 10962 times.
✓ Branch 1 taken 195048 times.
✓ Branch 2 taken 10962 times.
✓ Branch 3 taken 184086 times.
|
206010 | inline int clampi(int v, int lo, int hi) { return v < lo ? lo : (v > hi ? hi : v); } |
| 49 | |||
| 50 | } // namespace | ||
| 51 | |||
| 52 | 5 | void convex_upsample_forward(const ::brotensor::Tensor& X, | |
| 53 | const ::brotensor::Tensor& Mask, | ||
| 54 | int N, int C, int H, int W, int scale, | ||
| 55 | ::brotensor::Tensor& Y) { | ||
| 56 | 5 | const char* op = "convex_upsample_forward"; | |
| 57 | 5 | check_fp32(X, op, "X"); | |
| 58 | 5 | check_fp32(Mask, op, "Mask"); | |
| 59 |
1/4✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
5 | if (N < 0 || C < 1 || H < 1 || W < 1) fail(op, "C/H/W must be >=1 and N >=0"); |
| 60 |
1/4✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
5 | if (scale < 1) fail(op, "scale must be >=1"); |
| 61 | 5 | const int HW = H * W; | |
| 62 | 5 | const int kk = scale * scale; | |
| 63 |
1/4✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
|
5 | if (X.rows != N || X.cols != C * HW) fail(op, "X shape must be (N, C*H*W)"); |
| 64 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | if (Mask.rows != N || Mask.cols != 9 * kk * HW) |
| 65 | ✗ | fail(op, "Mask shape must be (N, 9*scale*scale*H*W)"); | |
| 66 | 5 | const int oH = scale * H, oW = scale * W, oHW = oH * oW; | |
| 67 |
1/6✗ Branch 0 not taken.
✓ Branch 1 taken 5 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
|
5 | if (Y.rows != N || Y.cols != C * oHW || Y.dtype != Dtype::FP32) |
| 68 | 5 | Y.resize(N, C * oHW, Dtype::FP32); | |
| 69 |
1/2✓ Branch 0 taken 5 times.
✗ Branch 1 not taken.
|
5 | if (N == 0) return; |
| 70 | |||
| 71 | 5 | const float* Xp = X.host_f32(); | |
| 72 | 5 | const float* Mp = Mask.host_f32(); | |
| 73 | 5 | float* Yp = Y.host_f32_mut(); | |
| 74 | |||
| 75 | double w[9]; | ||
| 76 |
2/2✓ Branch 0 taken 6 times.
✓ Branch 1 taken 5 times.
|
11 | for (int n = 0; n < N; ++n) { |
| 77 | 6 | const float* x_img = Xp + static_cast<long>(n) * C * HW; | |
| 78 | 6 | const float* m_img = Mp + static_cast<long>(n) * 9 * kk * HW; | |
| 79 | 6 | float* y_img = Yp + static_cast<long>(n) * C * oHW; | |
| 80 |
2/2✓ Branch 0 taken 27 times.
✓ Branch 1 taken 6 times.
|
33 | for (int y = 0; y < H; ++y) { |
| 81 |
2/2✓ Branch 0 taken 171 times.
✓ Branch 1 taken 27 times.
|
198 | for (int x = 0; x < W; ++x) { |
| 82 | 171 | const int pix = y * W + x; | |
| 83 |
2/2✓ Branch 0 taken 679 times.
✓ Branch 1 taken 171 times.
|
850 | for (int sy = 0; sy < scale; ++sy) { |
| 84 |
2/2✓ Branch 0 taken 4005 times.
✓ Branch 1 taken 679 times.
|
4684 | for (int sx = 0; sx < scale; ++sx) { |
| 85 | 4005 | const int sub = sy * scale + sx; | |
| 86 | // softmax over the 9 neighbors m | ||
| 87 | 4005 | double mx = -1e300; | |
| 88 |
2/2✓ Branch 0 taken 36045 times.
✓ Branch 1 taken 4005 times.
|
40050 | for (int m = 0; m < 9; ++m) { |
| 89 | 36045 | const double v = m_img[(static_cast<long>(m) * kk + sub) * HW + pix]; | |
| 90 |
2/2✓ Branch 0 taken 24807 times.
✓ Branch 1 taken 11238 times.
|
36045 | if (v > mx) mx = v; |
| 91 | 36045 | } | |
| 92 | 4005 | double sum = 0.0; | |
| 93 |
2/2✓ Branch 0 taken 36045 times.
✓ Branch 1 taken 4005 times.
|
40050 | for (int m = 0; m < 9; ++m) { |
| 94 | 36045 | const double e = std::exp( | |
| 95 | 36045 | m_img[(static_cast<long>(m) * kk + sub) * HW + pix] - mx); | |
| 96 | 36045 | w[m] = e; sum += e; | |
| 97 | 36045 | } | |
| 98 | 4005 | const double invs = 1.0 / sum; | |
| 99 |
2/2✓ Branch 0 taken 36045 times.
✓ Branch 1 taken 4005 times.
|
40050 | for (int m = 0; m < 9; ++m) w[m] *= invs; |
| 100 | |||
| 101 | 4005 | const int oy = scale * y + sy, ox = scale * x + sx; | |
| 102 | 4005 | const long opix = static_cast<long>(oy) * oW + ox; | |
| 103 |
2/2✓ Branch 0 taken 11445 times.
✓ Branch 1 taken 4005 times.
|
15450 | for (int c = 0; c < C; ++c) { |
| 104 | 11445 | const float* xc = x_img + static_cast<long>(c) * HW; | |
| 105 | 11445 | double acc = 0.0; | |
| 106 |
2/2✓ Branch 0 taken 103005 times.
✓ Branch 1 taken 11445 times.
|
114450 | for (int m = 0; m < 9; ++m) { |
| 107 | 103005 | const int ny = clampi(y - 1 + m / 3, 0, H - 1); | |
| 108 | 103005 | const int nx = clampi(x - 1 + m % 3, 0, W - 1); | |
| 109 | 103005 | acc += w[m] * xc[static_cast<long>(ny) * W + nx]; | |
| 110 | 103005 | } | |
| 111 | 11445 | y_img[static_cast<long>(c) * oHW + opix] = static_cast<float>(acc); | |
| 112 | 11445 | } | |
| 113 | 4005 | } | |
| 114 | 679 | } | |
| 115 | 171 | } | |
| 116 | 27 | } | |
| 117 | 6 | } | |
| 118 | 5 | } | |
| 119 | |||
| 120 | } // namespace brotensor::detail::cpu | ||
| 121 |