code/tt_moge/host_ext/moge_host.cpp

// SPDX-License-Identifier: Apache-2.0
// MoGe-2 host pre/post-processing kernels (fp32, multithreaded, no torch dependency).
// Built on first use by tt_moge/tt/host_ext.py (g++ -O3 -march=native) and loaded with ctypes.
//
// moge_tail:  the fused-path host tail of MoGeModel.forward in ONE pass over the full-resolution
//             output: bf16 channel-major low-res device output (o-major phase rows, see
//             fused_host.chanmajor_to_planar) -> bilinear upsample (align_corners=False, no
//             antialias; torch's index/weight rules and its width-then-height evaluation order)
//             -> exp remap of the points, L2 normalize of the normals, sigmoid of the mask,
//             written straight into channels-last fp32 outputs [H, W, 3] / [H, W, 3] / [H, W].
// moge_input: the encoder's host input: separable resize with precomputed (torch-extracted)
//             weights, (x - mean) / std, 14x14 patch rows (c, kh, kw) -> bf16 (round to nearest
//             even) [rows_total, k_pad], zero padding rows/cols.
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <limits>
#if defined(__F16C__) || defined(__AVX512F__) || defined(__SSE4_1__)
#include <immintrin.h>
#endif
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <mutex>
#include <memory>
#include <thread>
#include <vector>

namespace {

inline float bf16_to_f32(uint16_t v) {
    uint32_t u = static_cast<uint32_t>(v) << 16;
    float f;
    std::memcpy(&f, &u, 4);
    return f;
}

inline uint16_t f32_to_bf16(float f) {  // round to nearest even (torch's .to(bfloat16)); no NaN inputs
    uint32_t u;
    std::memcpy(&u, &f, 4);
    u += 0x7FFFu + ((u >> 16) & 1u);
    return static_cast<uint16_t>(u >> 16);
}


#if defined(__AVX512F__)
#include <immintrin.h>
#define MOGE_AVX512 1
// expf, Cephes polynomial (~1 ulp in fp32), input clamped to [-88.37, 88.37]
inline __m512 exp16(__m512 x) {
    x = _mm512_min_ps(_mm512_max_ps(x, _mm512_set1_ps(-88.3762626647949f)), _mm512_set1_ps(88.3762626647949f));
    __m512 fx = _mm512_roundscale_ps(_mm512_mul_ps(x, _mm512_set1_ps(1.44269504088896341f)),
                                     _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
    x = _mm512_fnmadd_ps(fx, _mm512_set1_ps(0.693359375f), x);
    x = _mm512_fnmadd_ps(fx, _mm512_set1_ps(-2.12194440e-4f), x);
    __m512 y = _mm512_set1_ps(1.9875691500E-4f);
    y = _mm512_fmadd_ps(y, x, _mm512_set1_ps(1.3981999507E-3f));
    y = _mm512_fmadd_ps(y, x, _mm512_set1_ps(8.3334519073E-3f));
    y = _mm512_fmadd_ps(y, x, _mm512_set1_ps(4.1665795894E-2f));
    y = _mm512_fmadd_ps(y, x, _mm512_set1_ps(1.6666665459E-1f));
    y = _mm512_fmadd_ps(y, x, _mm512_set1_ps(5.0000001201E-1f));
    y = _mm512_fmadd_ps(y, _mm512_mul_ps(x, x), _mm512_add_ps(x, _mm512_set1_ps(1.f)));
    return _mm512_scalef_ps(y, fx);
}
struct Interleave3 {   // out[3i + k] = v_k[i] for 16 pixels -> 3 vectors
    __m512i ab[3], c[3];
    __mmask16 mc[3];
    Interleave3() {
        alignas(64) int iab[3][16], ic[3][16];
        for (int k = 0; k < 3; ++k) {
            uint32_t m = 0;
            for (int j = 0; j < 16; ++j) {
                int g = 16 * k + j, i = g / 3, ch = g % 3;
                iab[k][j] = (ch == 1) ? 16 + i : i;
                ic[k][j] = i;
                if (ch == 2) m |= (1u << j);
            }
            ab[k] = _mm512_load_si512(iab[k]);
            c[k] = _mm512_load_si512(ic[k]);
            mc[k] = static_cast<__mmask16>(m);
        }
    }
    inline void store(float* dst, __m512 a, __m512 b, __m512 cc, bool nt) const {
        for (int k = 0; k < 3; ++k) {
            __m512 t = _mm512_permutex2var_ps(a, ab[k], b);
            t = _mm512_mask_permutexvar_ps(t, mc[k], c[k], cc);
            if (nt) _mm512_stream_ps(dst + 16 * k, t);   // full-resolution outputs: no read-for-ownership
            else _mm512_storeu_ps(dst + 16 * k, t);
        }
    }
};
#endif

// torch: area_pixel_compute_source_index(scale, i, align_corners=false) + guard_index_and_lambda
inline void lin_index(int i, int in_size, float scale, int& i0, int& i1, float& w0, float& w1) {
    float src = std::fma(scale, static_cast<float>(i) + 0.5f, -0.5f);   // torch's build contracts this to an fma
    if (src < 0.f) src = 0.f;
    int idx = std::min(static_cast<int>(std::floor(src)), in_size - 1);
    float lam = std::min(std::max(src - static_cast<float>(idx), 0.f), 1.f);
    i0 = idx;
    i1 = idx + ((idx < in_size - 1) ? 1 : 0);
    w0 = 1.f - lam;
    w1 = lam;
}

// run_parallel(n, fn): fn(0..n-1), fn(0) on the caller, one fresh std::thread per other share. (A
// persistent condvar pool was tried in round 5: ~5x lower empty-call latency, but the memory-bound input
// prep ran 0.2-0.7 ms SLOWER after the usual idle gap between calls (cold workers), so: not used.)
template <class F>
void run_parallel(int nthreads, F&& fn) {
    if (nthreads <= 1) { fn(0); return; }
    std::vector<std::thread> th;
    th.reserve(nthreads - 1);
    for (int t = 1; t < nthreads; ++t) th.emplace_back([&fn, t] { fn(t); });
    fn(0);
    for (auto& x : th) x.join();
}

template <class F>
void parallel_rows(int n, int nthreads, F&& fn) {
    nthreads = std::max(1, std::min(nthreads, n));
    if (nthreads == 1) {
        fn(0, n);
        return;
    }
    const int chunk = (n + nthreads - 1) / nthreads;
    run_parallel(nthreads, [&](int t) {
        const int a = t * chunk, b = std::min(n, a + chunk);
        if (a < b) fn(a, b);
    });
}

}  // namespace

extern "C" {

// y      : bf16 [rows, ld]; row (o*4 + dy*2 + dx), column (Y*W3 + X) = low-res pixel (2Y+dy, 2X+dx)
//          of channel o (low-res image is [2*H3, 2*W3]).
// c_pts/c_nrm/c_msk : channel index o of points (3 ch) / normal (3 ch) / mask (1 ch), -1 = absent.
// pts/nrm: fp32 [H, W, 3] contiguous; msk: fp32 [H, W]; rows [y_begin, y_end) are written.
// Banded input: the low-res rows Y live in nbands buffers of band_Y rows each (Y-major bands,
// read back one by one); band b holds rows [b*band_Y, (b+1)*band_Y) as bf16 [rows, ld], ld = band_Y*W3.
// One band (band_Y = H3) is the plain [rows, H3*W3] layout.
int moge_tail_banded(const uint16_t* const* bands, int nbands, int band_Y, int64_t ld, int H3, int W3,
                     int c_pts, int c_nrm, int c_msk, int H, int W, float* pts, float* nrm, float* msk,
                     int y_begin, int y_end, int nthreads) {
    (void)nbands;
    const int Hl = 2 * H3, Wl = 2 * W3;
    const float sh = static_cast<float>(Hl) / static_cast<float>(H);
    const float sw = static_cast<float>(Wl) / static_cast<float>(W);
    std::vector<int> xi0(W), xi1(W);
    std::vector<float> xw0(W), xw1(W);
    for (int x = 0; x < W; ++x) lin_index(x, Wl, sw, xi0[x], xi1[x], xw0[x], xw1[x]);
    // per 16-output block: window start (or -1 -> gather) and the in-window tap indices
    std::vector<int> pbase(W / 16 + 1, -1), prel0(W), prel1(W);
    for (int b0 = 0; b0 + 16 <= W; b0 += 16) {
        int lo = xi0[b0], hi = 0;
        for (int j = 0; j < 16; ++j) hi = std::max(hi, xi1[b0 + j]);
        if (hi - lo < 16 && lo + 16 <= Wl + 32) {   // dec has 32 floats of slack past Wl
            pbase[b0 / 16] = lo;
            for (int j = 0; j < 16; ++j) { prel0[b0 + j] = xi0[b0 + j] - lo; prel1[b0 + j] = xi1[b0 + j] - lo; }
        }
    }
    int chans[7];
    int nch = 0;
    if (c_pts >= 0) for (int k = 0; k < 3; ++k) chans[nch++] = c_pts + k;
    if (c_nrm >= 0) for (int k = 0; k < 3; ++k) chans[nch++] = c_nrm + k;
    if (c_msk >= 0) chans[nch++] = c_msk;
    if (y_end > H) y_end = H;
    const int n = y_end - y_begin;
    if (n <= 0) return 0;

    parallel_rows(n, nthreads, [&](int a, int b) {
        // per thread: decoded low-res row (Wl) and two horizontally interpolated rows (W) per channel
        std::vector<float> dec(Wl + 32);
#ifdef MOGE_AVX512
        alignas(64) int lo[16], hi[16];
        for (int j = 0; j < 8; ++j) { lo[2 * j] = j; lo[2 * j + 1] = 16 + j; hi[2 * j] = 8 + j; hi[2 * j + 1] = 24 + j; }
        const __m512i ilo = _mm512_load_si512(lo), ihi = _mm512_load_si512(hi);
        const Interleave3 il3;
#endif
        std::vector<float> hrow[2];
        hrow[0].resize(static_cast<size_t>(nch) * W);
        hrow[1].resize(static_cast<size_t>(nch) * W);
        int cached[2] = {-1, -1};
        auto hinterp = [&](int r, float* dst) {   // low-res row r of every channel -> width W
            const int Y = r >> 1, dy = r & 1;
            for (int ci = 0; ci < nch; ++ci) {
                const uint16_t* yb = bands[Y / band_Y];
                const int64_t yo = static_cast<int64_t>(Y % band_Y) * W3;
                const uint16_t* r0 = yb + (static_cast<int64_t>(chans[ci]) * 4 + dy * 2 + 0) * ld + yo;
                const uint16_t* r1 = yb + (static_cast<int64_t>(chans[ci]) * 4 + dy * 2 + 1) * ld + yo;
                int X = 0;
                float* d = dst + static_cast<size_t>(ci) * W;
                const float* s = dec.data();
                int x = 0;
#ifdef MOGE_AVX512
                for (; X + 16 <= W3; X += 16) {
                    __m512i a = _mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256(reinterpret_cast<const __m256i*>(r0 + X))), 16);
                    __m512i b = _mm512_slli_epi32(_mm512_cvtepu16_epi32(_mm256_loadu_si256(reinterpret_cast<const __m256i*>(r1 + X))), 16);
                    _mm512_storeu_si512(dec.data() + 2 * X, _mm512_permutex2var_epi32(a, ilo, b));
                    _mm512_storeu_si512(dec.data() + 2 * X + 16, _mm512_permutex2var_epi32(a, ihi, b));
                }
#endif
                for (; X < W3; ++X) {
                    dec[2 * X] = bf16_to_f32(r0[X]);
                    dec[2 * X + 1] = bf16_to_f32(r1[X]);
                }
#ifdef MOGE_AVX512
                for (; x + 16 <= W; x += 16) {
                    __m512 g0, g1;
                    const int blk = x >> 4;
                    if (pbase[blk] >= 0) {     // all 16 taps inside one 16-float window: load + 2 permutes
                        const __m512 win = _mm512_loadu_ps(s + pbase[blk]);
                        g0 = _mm512_permutexvar_ps(_mm512_loadu_si512(prel0.data() + x), win);
                        g1 = _mm512_permutexvar_ps(_mm512_loadu_si512(prel1.data() + x), win);
                    } else {
                        g0 = _mm512_i32gather_ps(_mm512_loadu_si512(xi0.data() + x), s, 4);
                        g1 = _mm512_i32gather_ps(_mm512_loadu_si512(xi1.data() + x), s, 4);
                    }
                    _mm512_storeu_ps(d + x, _mm512_fmadd_ps(g0, _mm512_loadu_ps(xw0.data() + x),
                                                           _mm512_mul_ps(g1, _mm512_loadu_ps(xw1.data() + x))));
                }
#endif
                for (; x < W; ++x) d[x] = std::fma(s[xi0[x]], xw0[x], s[xi1[x]] * xw1[x]);
            }
        };
        auto get = [&](int r) -> const float* {
            for (int k = 0; k < 2; ++k)
                if (cached[k] == r) return hrow[k].data();
            // evict the slot not holding the other needed row: rows are visited in increasing order
            int slot = (cached[0] < cached[1]) ? 0 : 1;
            hinterp(r, hrow[slot].data());
            cached[slot] = r;
            return hrow[slot].data();
        };
        for (int yy = y_begin + a; yy < y_begin + b; ++yy) {
            int h0, h1;
            float wh0, wh1;
            lin_index(yy, Hl, sh, h0, h1, wh0, wh1);
            const float* A = get(h0);
            const float* B = get(h1);
            const size_t row = static_cast<size_t>(yy) * W;
            int x = 0;
#ifdef MOGE_AVX512
            const __m512 v0 = _mm512_set1_ps(wh0), v1 = _mm512_set1_ps(wh1);
            auto al = [](const float* p) { return (reinterpret_cast<uintptr_t>(p) & 63) == 0; };
            const bool nt_p = pts && al(pts + row * 3) && (W % 16 == 0);
            const bool nt_n = nrm && al(nrm + row * 3) && (W % 16 == 0);
            const bool nt_m = msk && al(msk + row) && (W % 16 == 0);
            for (; x + 16 <= W; x += 16) {
                __m512 v[7];
                for (int ci = 0; ci < nch; ++ci) {
                    const size_t off = static_cast<size_t>(ci) * W + x;
                    v[ci] = _mm512_fmadd_ps(_mm512_loadu_ps(A + off), v0, _mm512_mul_ps(_mm512_loadu_ps(B + off), v1));
                }
                int ci = 0;
                if (c_pts >= 0) {
                    __m512 z = exp16(v[2]);
                    il3.store(pts + (row + x) * 3, _mm512_mul_ps(v[0], z), _mm512_mul_ps(v[1], z), z, nt_p);
                    ci = 3;
                }
                if (c_nrm >= 0) {
                    __m512 nx = v[ci], ny = v[ci + 1], nz = v[ci + 2];
                    __m512 s2 = _mm512_add_ps(_mm512_add_ps(_mm512_mul_ps(nx, nx), _mm512_mul_ps(ny, ny)), _mm512_mul_ps(nz, nz));
                    __m512 sn = _mm512_max_ps(_mm512_sqrt_ps(s2), _mm512_set1_ps(1e-12f));
                    il3.store(nrm + (row + x) * 3, _mm512_div_ps(nx, sn), _mm512_div_ps(ny, sn), _mm512_div_ps(nz, sn), nt_n);
                    ci += 3;
                }
                if (c_msk >= 0) {
                    __m512 e = exp16(_mm512_sub_ps(_mm512_setzero_ps(), v[ci]));
                    __m512 sg = _mm512_div_ps(_mm512_set1_ps(1.f), _mm512_add_ps(_mm512_set1_ps(1.f), e));
                    if (nt_m) _mm512_stream_ps(msk + row + x, sg);
                    else _mm512_storeu_ps(msk + row + x, sg);
                }
            }
#endif
            for (; x < W; ++x) {
                float v[7];
                for (int ci = 0; ci < nch; ++ci) {
                    const size_t off = static_cast<size_t>(ci) * W + x;
                    v[ci] = std::fma(A[off], wh0, B[off] * wh1);
                }
                int ci = 0;
                if (c_pts >= 0) {
                    float z = std::exp(v[2]);
                    float* dst = pts + (row + x) * 3;
                    dst[0] = v[0] * z;
                    dst[1] = v[1] * z;
                    dst[2] = z;
                    ci = 3;
                }
                if (c_nrm >= 0) {
                    float s = std::max(std::sqrt(v[ci] * v[ci] + v[ci + 1] * v[ci + 1] + v[ci + 2] * v[ci + 2]), 1e-12f);
                    float* dst = nrm + (row + x) * 3;
                    dst[0] = v[ci] / s;
                    dst[1] = v[ci + 1] / s;
                    dst[2] = v[ci + 2] / s;
                    ci += 3;
                }
                if (c_msk >= 0) msk[row + x] = 1.f / (1.f + std::exp(-v[ci]));
            }
        }
#ifdef MOGE_AVX512
        _mm_sfence();
#endif
    });
    return 0;
}

int moge_tail(const uint16_t* y, int64_t ld, int H3, int W3, int c_pts, int c_nrm, int c_msk,
              int H, int W, float* pts, float* nrm, float* msk, int y_begin, int y_end, int nthreads) {
    const uint16_t* bands[1] = {y};
    return moge_tail_banded(bands, 1, H3, ld, H3, W3, c_pts, c_nrm, c_msk, H, W, pts, nrm, msk, y_begin, y_end, nthreads);
}

// number of leading output rows (of H) whose bilinear source rows all lie in low-res band rows Y < Y_avail
int moge_tail_ready(int H, int H3, int Y_avail) {
    const int Hl = 2 * H3;
    const float sh = static_cast<float>(Hl) / static_cast<float>(H);
    int n = 0;
    for (int yy = 0; yy < H; ++yy) {
        int h0, h1;
        float w0, w1;
        lin_index(yy, Hl, sh, h0, h1, w0, w1);
        if ((h1 >> 1) >= Y_avail) break;
        n = yy + 1;
    }
    return n;
}

// asynchronous banded tail: runs on its own std::thread (which fans out to nthreads), so the caller
// can read back the next band meanwhile; moge_tail_wait joins it.
struct TailJob {
    std::vector<const uint16_t*> bands;
    std::thread th;
};

void* moge_tail_start(const uint16_t* const* bands, int nbands, int band_Y, int64_t ld, int H3, int W3,
                      int c_pts, int c_nrm, int c_msk, int H, int W, float* pts, float* nrm, float* msk,
                      int y_begin, int y_end, int nthreads) {
    auto* job = new TailJob;
    job->bands.assign(bands, bands + nbands);
    const uint16_t* const* bp = job->bands.data();
    job->th = std::thread([=] {
        moge_tail_banded(bp, nbands, band_Y, ld, H3, W3, c_pts, c_nrm, c_msk, H, W, pts, nrm, msk, y_begin, y_end, nthreads);
    });
    return job;
}

int moge_tail_wait(void* handle) {
    auto* job = static_cast<TailJob*>(handle);
    if (job->th.joinable()) job->th.join();
    delete job;
    return 0;
}

// img      : fp32 image, element (c, h, w) at img[c*sc + h*shs + w*sws] (any stride layout), C <= 4.
// ry_*     : vertical taps   per output row y  (Ho rows): start ry_s[y], weights ry_w[y*ky + k], k < ky
// rx_*     : horizontal taps per output col x (Wo cols): start rx_s[x], weights rx_w[x*kx + k], k < kx
//            (zero-weight padding taps allowed; start + k must stay in range)
// Evaluation order: height pass first (contiguous rows), then width (torch does width first: the
// results differ by fp32 rounding only); each output is s0*w0 + unfused accumulation of the taps.
// out      : bf16 [rows_total, k_pad]; patch row (i*bw + j), column (c*p*p + kh*p + kw)
}  // extern "C"
namespace {
// One input-prep job: output rows (patch row i, kh) are independent units; unit (i, 0) also zeroes the
// padding columns of patch row i's tokens, unit (bh-1, 0) the trailing rows. Rows go to the band buffer
// that holds patch row i (bands = consecutive patch-row ranges; one band = the whole-image layout).
// u8 -> float32 exactly as postprocess.fit_to_canvas / moge_fit_canvas: float32(u8) / 255.0f (IEEE division:
// the same value as their lookup table; written as a division so the loops vectorize)
inline float ldf(const float* p, int64_t i) { return p[i]; }
inline float ldf(const uint8_t* p, int64_t i) { return static_cast<float>(p[i]) / 255.0f; }

struct InputPrep {
    const float* img; int C, Hi, Wi; int64_t sc, shs, sws; int Ho, Wo;
    const int* ry_s; const float* ry_w; int ky; const int* rx_s; const float* rx_w; int kx;
    const float* mean; const float* stdv; int patch, rows_total, k_pad;
    // round 12: uint8 source (the served canvas, HWC or any strides; elements = u8 / 255 as float32, bit-identical
    // to reading the float32 canvas that moge_fit_canvas would have produced from it). nullptr = float `img`.
    const uint8_t* img8 = nullptr;
    int bh = 0, bw = 0, K = 0;
    bool hwc = false;
    std::vector<float> wxT;
    // round 12: width pass without gathers. Per 16-output block, its taps lie in a 48-float window from wbase[blk]
    // (planar rows): tap k of lane j = window[widx[(blk * kx + k) * 16 + j]], fetched with two permutex2var + a
    // blend instead of a gather; the multiply-adds are unchanged (bit-identical). wok[blk] = 0 -> gather.
    std::vector<int> wbase, widx;
    std::vector<uint8_t> wok;
    std::vector<int> band_of;            // patch row -> band
    std::vector<uint16_t*> base;         // band buffer shifted so that token row r is base + r * k_pad
    void init(int i_lo, int i_hi, const int* edges, int nb, uint16_t* const* outs) {
        bh = Ho / patch; bw = Wo / patch; K = C * patch * patch;
        wxT.assign(static_cast<size_t>(kx) * Wo, 0.f);
        for (int x = 0; x < Wo; ++x)
            for (int k = 0; k < kx; ++k) wxT[static_cast<size_t>(k) * Wo + x] = rx_w[static_cast<size_t>(x) * kx + k];
        hwc = (sc == 1 && sws == C);
        const int nblk = Wo / 16;
        wbase.assign(nblk, 0);
        wok.assign(nblk, 0);
        widx.assign(static_cast<size_t>(nblk) * kx * 16, 0);
        for (int b = 0; b < nblk; ++b) {
            int lo = rx_s[16 * b], hi = lo;
            for (int j = 0; j < 16; ++j) { lo = std::min(lo, rx_s[16 * b + j]); hi = std::max(hi, rx_s[16 * b + j] + kx - 1); }
            wbase[b] = lo;
            wok[b] = (hi - lo < 48) ? 1 : 0;
            for (int k = 0; k < kx; ++k)
                for (int j = 0; j < 16; ++j) widx[(static_cast<size_t>(b) * kx + k) * 16 + j] = rx_s[16 * b + j] + k - lo;
        }
        band_of.assign(bh, -1);
        base.assign(nb, nullptr);
        for (int k = 0; k < nb; ++k) {
            for (int i = std::max(edges[k], i_lo); i < std::min(edges[k + 1], i_hi); ++i) band_of[i] = k;
            base[k] = outs[k] - static_cast<int64_t>(edges[k]) * bw * k_pad;
        }
    }
    struct Bufs {
        std::vector<float> tmp, il, vrow;
        std::vector<uint16_t> brow;
        // u8 source: a per-thread ring of input rows converted once to planar float32 (u8 / 255.0f); slot of input
        // row r = r % ring_n (ring_n >= ky + 4: the taps of consecutive output rows overlap and advance ~2.4 rows)
        int ring_n = 0;
        std::vector<float> ring;
        std::vector<int> ring_tag;
        explicit Bufs(const InputPrep& P)
            : tmp(static_cast<size_t>(P.Wi) * P.C + 64), il(static_cast<size_t>(P.Wi) * P.C + 64), vrow(P.Wo), brow(P.Wo + 32) {
            if (P.img8) {
                ring_n = P.ky + 4;
                ring.resize(static_cast<size_t>(ring_n) * P.C * P.Wi);
                ring_tag.assign(ring_n, -1);
            }
        }
    };
#ifdef MOGE_AVX512
    // [Wi][3] interleaved -> planar dst[c * Wi + w] (16 pixels per step: 2 permutes per channel)
    void deinterleave3(const float* il, float* dst) const {
        alignas(64) int iab[3][16], ic[3][16];
        __mmask16 mc[3];
        for (int c = 0; c < 3; ++c) {
            mc[c] = 0;
            for (int j = 0; j < 16; ++j) {
                const int p = 3 * j + c;
                iab[c][j] = p & 31;
                ic[c][j] = (p - 32) & 15;
                if (p >= 32) mc[c] |= static_cast<__mmask16>(1u << j);
            }
        }
        const __m512i I0 = _mm512_load_si512(iab[0]), I1 = _mm512_load_si512(iab[1]), I2 = _mm512_load_si512(iab[2]);
        const __m512i J0 = _mm512_load_si512(ic[0]), J1 = _mm512_load_si512(ic[1]), J2 = _mm512_load_si512(ic[2]);
        float* d0 = dst; float* d1 = dst + Wi; float* d2 = dst + 2 * static_cast<size_t>(Wi);
        int w = 0;
        for (; w + 16 <= Wi; w += 16) {
            const float* s = il + 3 * w;
            const __m512 a = _mm512_loadu_ps(s), b = _mm512_loadu_ps(s + 16), c = _mm512_loadu_ps(s + 32);
            _mm512_storeu_ps(d0 + w, _mm512_mask_permutexvar_ps(_mm512_permutex2var_ps(a, I0, b), mc[0], J0, c));
            _mm512_storeu_ps(d1 + w, _mm512_mask_permutexvar_ps(_mm512_permutex2var_ps(a, I1, b), mc[1], J1, c));
            _mm512_storeu_ps(d2 + w, _mm512_mask_permutexvar_ps(_mm512_permutex2var_ps(a, I2, b), mc[2], J2, c));
        }
        for (; w < Wi; ++w) { d0[w] = il[3 * w]; d1[w] = il[3 * w + 1]; d2[w] = il[3 * w + 2]; }
    }
#endif
    // u8 path (HWC source: sc = 1, sws = C): input row r as interleaved float32 [Wi * C] (converted on first use
    // by this thread; contiguous, vectorized)
    const float* ring_row(int r, Bufs& B) const {
        const int slot = r % B.ring_n;
        const int n = Wi * C;
        float* d = B.ring.data() + static_cast<size_t>(slot) * n;
        if (B.ring_tag[slot] != r) {
            const uint8_t* src = img8 + static_cast<int64_t>(r) * shs;
            int j = 0;
#ifdef MOGE_AVX512
            // u8 / 255.0f without the (slow) division: q0 = v * fl(1/255), then one Markstein correction step with
            // explicit FMAs (q1 = q0 + (v - q0 * 255) * fl(1/255)) gives the correctly rounded quotient for all 256
            // values (checked exhaustively by moge_check_u8_div / tests: bit-identical to the division)
            const __m512 c = _mm512_set1_ps(1.0f / 255.0f), k255 = _mm512_set1_ps(255.0f);
            for (; j + 16 <= n; j += 16) {
                const __m512 v = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast<const __m128i*>(src + j))));
                const __m512 q0 = _mm512_mul_ps(v, c);
                const __m512 e = _mm512_fnmadd_ps(q0, k255, v);
                _mm512_storeu_ps(d + j, _mm512_fmadd_ps(e, c, q0));
            }
#endif
            for (; j < n; ++j) d[j] = static_cast<float>(src[j]) / 255.0f;
            B.ring_tag[slot] = r;
        }
        return d;
    }
    // u8 height pass from the ring: the same multiply-adds in the same tap order as vpass on the float32 canvas
    // (bit-identical), then the HWC -> planar transpose of the hwc path
    void hpass_ring(int yo, const float* wv, Bufs& B, float* tmp) const {
        const int s = ry_s[yo];
        const int n = Wi * C;
        const float* rows[64];
        for (int k = 0; k < ky; ++k) rows[k] = ring_row(s + k, B);
        float* il = B.il.data();
        for (int j = 0; j < n; ++j) il[j] = rows[0][j] * wv[0];
        for (int k = 1; k < ky; ++k) {
            const float wk = wv[k];
            const float* rk = rows[k];
            for (int j = 0; j < n; ++j) il[j] += rk[j] * wk;
        }
        (void)tmp;   // result stays interleaved in B.il: the width pass reads it with stride C
    }
    template <class T>
    void vpass(const T* src, int64_t stride_k, const float* wv, int s, int n, float* dst) const {
        const T* r0 = src + static_cast<int64_t>(s) * stride_k;
        for (int w = 0; w < n; ++w) dst[w] = ldf(r0, w) * wv[0];
        for (int k = 1; k < ky; ++k) {
            const float wk = wv[k];
            const T* rk = r0 + static_cast<int64_t>(k) * stride_k;
            for (int w = 0; w < n; ++w) dst[w] += ldf(rk, w) * wk;
        }
    }
    // height pass (contiguous rows, vectorized), planar result tmp[c * Wi + w]
    template <class T>
    void hpass(const T* im, int yo, const float* wv, Bufs& B, float* tmp) const {
        if (hwc) {
            vpass(im, shs, wv, ry_s[yo], Wi * C, B.il.data());
            for (int w = 0; w < Wi; ++w)
                for (int c = 0; c < C; ++c) tmp[static_cast<size_t>(c) * Wi + w] = B.il[static_cast<size_t>(w) * C + c];
        } else if (sws == 1) {
            for (int c = 0; c < C; ++c) vpass(im + c * sc, shs, wv, ry_s[yo], Wi, tmp + static_cast<size_t>(c) * Wi);
        } else {
            for (int c = 0; c < C; ++c) {
                const T* r0 = im + c * sc + static_cast<int64_t>(ry_s[yo]) * shs;
                float* d = tmp + static_cast<size_t>(c) * Wi;
                for (int w = 0; w < Wi; ++w) d[w] = ldf(r0, w * sws) * wv[0];
                for (int k = 1; k < ky; ++k)
                    for (int w = 0; w < Wi; ++w) d[w] += ldf(r0, static_cast<int64_t>(k) * shs + w * sws) * wv[k];
            }
        }
    }
    void unit(int i, int kh, Bufs& B) const {
        uint16_t* out = base[band_of[i]];
        if (kh == 0) {
            for (int r = i * bw; r < (i + 1) * bw; ++r) std::memset(out + static_cast<size_t>(r) * k_pad + K, 0, sizeof(uint16_t) * (k_pad - K));
            if (i == bh - 1)
                std::memset(out + static_cast<size_t>(bh) * bw * k_pad, 0, sizeof(uint16_t) * static_cast<size_t>(rows_total - bh * bw) * k_pad);
        }
        float* tmp = B.tmp.data();
        const int yo = i * patch + kh;
        const float* wv = ry_w + static_cast<size_t>(yo) * ky;
        // height pass first
        // height-pass result interleaved [Wi][C] in B.il (HWC sources; round 12: also float32 HWC, i.e. the bench's
        // permuted HWC image: the width pass reads with stride C instead of transposing; same arithmetic)
        bool il_src = hwc && (img8 ? ky <= 64 : true);
        if (il_src && img8) hpass_ring(yo, wv, B, tmp);
        else if (il_src) vpass(img, shs, wv, ry_s[yo], Wi * C, B.il.data());
        else if (img8) hpass(img8, yo, wv, B, tmp);
        else hpass(img, yo, wv, B, tmp);
#ifdef MOGE_AVX512
        if (il_src && C == 3) {   // vectorized de-interleave [Wi][3] -> planar tmp (pure data movement)
            deinterleave3(B.il.data(), tmp);
            il_src = false;
        }
#endif
        const int cs = il_src ? C : 1;                  // element stride of one channel's row
        float* vrow = B.vrow.data();
        uint16_t* brow = B.brow.data();
        for (int c = 0; c < C; ++c) {
            const float m = mean[c], sd = stdv[c];
            const float* src = il_src ? B.il.data() + c : tmp + static_cast<size_t>(c) * Wi;
            // width pass (gathers inside one row: planar tmp, or the interleaved u8-path row with stride C)
            int x = 0;
#ifdef MOGE_AVX512
            const __m512i vcs = _mm512_set1_epi32(cs);
            const __m512i v16 = _mm512_set1_epi32(16), v32 = _mm512_set1_epi32(32);
            for (; x + 16 <= Wo; x += 16) {
                const int blk = x >> 4;
                if (cs == 1 && wok[blk]) {   // taps from a 48-float window: permutes instead of gathers
                    const float* wsrc = src + wbase[blk];
                    const __m512 A = _mm512_loadu_ps(wsrc), Bv = _mm512_loadu_ps(wsrc + 16), Cv = _mm512_loadu_ps(wsrc + 32);
                    const int* ib = widx.data() + static_cast<size_t>(blk) * kx * 16;
                    auto tap = [&](int k) {
                        const __m512i id = _mm512_loadu_si512(ib + k * 16);
                        const __m512 lo = _mm512_permutex2var_ps(A, id, Bv);
                        const __m512 hi = _mm512_permutex2var_ps(Bv, _mm512_sub_epi32(id, v16), Cv);
                        return _mm512_mask_blend_ps(_mm512_cmpge_epi32_mask(id, v32), lo, hi);
                    };
                    __m512 acc = _mm512_mul_ps(tap(0), _mm512_loadu_ps(wxT.data() + x));
                    for (int k = 1; k < kx; ++k)
                        acc = _mm512_add_ps(acc, _mm512_mul_ps(tap(k), _mm512_loadu_ps(wxT.data() + static_cast<size_t>(k) * Wo + x)));
                    _mm512_storeu_ps(vrow + x, acc);
                    continue;
                }
                __m512i idx = _mm512_mullo_epi32(_mm512_loadu_si512(rx_s + x), vcs);
                __m512 acc = _mm512_mul_ps(_mm512_i32gather_ps(idx, src, 4), _mm512_loadu_ps(wxT.data() + x));
                for (int k = 1; k < kx; ++k) {
                    idx = _mm512_add_epi32(idx, vcs);
                    acc = _mm512_add_ps(acc, _mm512_mul_ps(_mm512_i32gather_ps(idx, src, 4),
                                                           _mm512_loadu_ps(wxT.data() + static_cast<size_t>(k) * Wo + x)));
                }
                _mm512_storeu_ps(vrow + x, acc);
            }
#endif
            for (; x < Wo; ++x) {
                const float* sv = src + static_cast<int64_t>(rx_s[x]) * cs;
                float acc = sv[0] * wxT[x];
                for (int k = 1; k < kx; ++k) acc += sv[static_cast<int64_t>(k) * cs] * wxT[static_cast<size_t>(k) * Wo + x];
                vrow[x] = acc;
            }
            x = 0;
#ifdef MOGE_AVX512
            const __m512 vm = _mm512_set1_ps(m), vsd = _mm512_set1_ps(sd);
            const __m512i rb = _mm512_set1_epi32(0x7FFF), one = _mm512_set1_epi32(1);
            for (; x + 16 <= Wo; x += 16) {   // (v - m) / sd, then round-to-nearest-even to bf16
                __m512i u = _mm512_castps_si512(_mm512_div_ps(_mm512_sub_ps(_mm512_loadu_ps(vrow + x), vm), vsd));
                u = _mm512_add_epi32(u, _mm512_add_epi32(rb, _mm512_and_si512(_mm512_srli_epi32(u, 16), one)));
                _mm256_storeu_si256(reinterpret_cast<__m256i*>(brow + x), _mm512_cvtepi32_epi16(_mm512_srli_epi32(u, 16)));
            }
#endif
            for (; x < Wo; ++x) brow[x] = f32_to_bf16((vrow[x] - m) / sd);
            uint16_t* o0 = out + static_cast<size_t>(i * bw) * k_pad + c * patch * patch + kh * patch;
#ifdef MOGE_AVX512
            if (patch <= 16) {   // one masked 16-lane store per patch (brow has 32 elements of slack)
                const __mmask16 pm = static_cast<__mmask16>((1u << patch) - 1u);
                for (int j = 0; j < bw; ++j)
                    _mm256_mask_storeu_epi16(o0 + static_cast<size_t>(j) * k_pad, pm,
                                             _mm256_loadu_si256(reinterpret_cast<const __m256i*>(brow + j * patch)));
                continue;
            }
#endif
            for (int j = 0; j < bw; ++j)
                std::memcpy(o0 + static_cast<size_t>(j) * k_pad, brow + j * patch, sizeof(uint16_t) * patch);
        }
    }
};

// asynchronous banded prep: units claimed in order by nthreads workers (on a driver thread), per-band
// completion counters so the caller can upload band k while later bands are still being computed
struct InputJob {
    InputPrep P;
    int nb = 0, nthreads = 1, chunk = 8;
    std::vector<int> edges;
    std::unique_ptr<std::atomic<int>[]> done;
    std::atomic<int> next{0};
    std::thread driver;
};
}  // namespace
extern "C" {

// moge_input_band: patch rows [i0, i1) only (i.e. token rows [i0*bw, i1*bw)), written to `out` whose
// row 0 is token row i0*bw; the band with i1 == bh also gets the zero rows [bh*bw, rows_total). Work is
// split over (patch row, kh) pairs. Same arithmetic per output as the whole-image call.
static int input_band_impl(const float* img, const uint8_t* img8, int C, int Hi, int Wi, int64_t sc, int64_t shs,
                           int64_t sws, int Ho, int Wo, const int* ry_s, const float* ry_w, int ky,
                           const int* rx_s, const float* rx_w, int kx, const float* mean, const float* stdv,
                           int patch, uint16_t* out, int rows_total, int k_pad, int nthreads, int i0, int i1) {
    const int bh = Ho / patch;
    if (i0 < 0 || i1 > bh || i0 >= i1) return -1;
    InputPrep P{img, C, Hi, Wi, sc, shs, sws, Ho, Wo, ry_s, ry_w, ky, rx_s, rx_w, kx, mean, stdv, patch, rows_total, k_pad};
    P.img8 = img8;
    const int edges[2] = {i0, i1};
    P.init(i0, i1, edges, 1, &out);
    parallel_rows((i1 - i0) * patch, nthreads, [&](int ua, int ub) {
        InputPrep::Bufs B(P);
        for (int u = ua; u < ub; ++u) P.unit(i0 + u / patch, u % patch, B);
    });
    return 0;
}

int moge_input_band(const float* img, int C, int Hi, int Wi, int64_t sc, int64_t shs, int64_t sws,
                    int Ho, int Wo, const int* ry_s, const float* ry_w, int ky,
                    const int* rx_s, const float* rx_w, int kx, const float* mean, const float* stdv,
                    int patch, uint16_t* out, int rows_total, int k_pad, int nthreads, int i0, int i1) {
    return input_band_impl(img, nullptr, C, Hi, Wi, sc, shs, sws, Ho, Wo, ry_s, ry_w, ky, rx_s, rx_w, kx, mean, stdv,
                           patch, out, rows_total, k_pad, nthreads, i0, i1);
}

// round 12: the same from a uint8 image (strides in elements; HWC: sc = 1, sws = C): == moge_input_band on the
// float32 image u8 / 255.0f (the served canvas without materializing it as float32)
int moge_input_band_u8(const uint8_t* img, int C, int Hi, int Wi, int64_t sc, int64_t shs, int64_t sws,
                       int Ho, int Wo, const int* ry_s, const float* ry_w, int ky,
                       const int* rx_s, const float* rx_w, int kx, const float* mean, const float* stdv,
                       int patch, uint16_t* out, int rows_total, int k_pad, int nthreads, int i0, int i1) {
    return input_band_impl(nullptr, img, C, Hi, Wi, sc, shs, sws, Ho, Wo, ry_s, ry_w, ky, rx_s, rx_w, kx, mean, stdv,
                           patch, out, rows_total, k_pad, nthreads, i0, i1);
}

// moge_input_start: the whole image as nb bands (patch-row edges[0..nb], edges[0] = 0, edges[nb] = bh)
// into outs[k] (rows of band k; the last band also holds the trailing zero rows). Returns a handle;
// moge_input_wait_band(h, k) blocks until band k is complete, moge_input_join(h) joins and frees.
static void* input_start_impl(const float* img, const uint8_t* img8, int C, int Hi, int Wi, int64_t sc,
                              int64_t shs, int64_t sws, int Ho, int Wo, const int* ry_s, const float* ry_w, int ky,
                              const int* rx_s, const float* rx_w, int kx, const float* mean, const float* stdv,
                              int patch, int rows_total, int k_pad, int nthreads, int nb, const int* edges,
                              uint16_t* const* outs) {
    auto* J = new InputJob;
    J->P = InputPrep{img, C, Hi, Wi, sc, shs, sws, Ho, Wo, ry_s, ry_w, ky, rx_s, rx_w, kx, mean, stdv, patch, rows_total, k_pad};
    J->P.img8 = img8;
    J->nb = nb;
    J->nthreads = std::max(1, nthreads);
    if (const char* e = std::getenv("MOGE_IN_CHUNK")) J->chunk = std::max(1, std::atoi(e));
    J->edges.assign(edges, edges + nb + 1);
    J->P.init(0, Ho / patch, J->edges.data(), nb, outs);
    J->done.reset(new std::atomic<int>[nb]);
    for (int k = 0; k < nb; ++k) J->done[k].store(0);
    J->driver = std::thread([J] {
        const int units = J->P.bh * J->P.patch;
        run_parallel(std::min(J->nthreads, units), [J, units](int) {
            InputPrep::Bufs B(J->P);
            // units are claimed MOGE_IN_CHUNK (default 8) consecutive output rows at a time (round 11): the vertical
            // taps of neighbouring output rows overlap (~2.4 input rows apart, 6 taps), so one thread reuses its
            // input rows from cache (5 threads: 1.36-1.79 -> 1.08-1.33 ms; outputs identical)
            const int CH = J->chunk;
            for (;;) {
                const int u0 = J->next.fetch_add(CH, std::memory_order_relaxed);
                if (u0 >= units) break;
                for (int u = u0; u < std::min(u0 + CH, units); ++u) {
                    const int i = u / J->P.patch;
                    J->P.unit(i, u % J->P.patch, B);
                    J->done[J->P.band_of[i]].fetch_add(1, std::memory_order_release);
                }
            }
        });
    });
    return J;
}

void* moge_input_start(const float* img, int C, int Hi, int Wi, int64_t sc, int64_t shs, int64_t sws,
                       int Ho, int Wo, const int* ry_s, const float* ry_w, int ky,
                       const int* rx_s, const float* rx_w, int kx, const float* mean, const float* stdv,
                       int patch, int rows_total, int k_pad, int nthreads, int nb, const int* edges,
                       uint16_t* const* outs) {
    return input_start_impl(img, nullptr, C, Hi, Wi, sc, shs, sws, Ho, Wo, ry_s, ry_w, ky, rx_s, rx_w, kx, mean, stdv,
                            patch, rows_total, k_pad, nthreads, nb, edges, outs);
}

void* moge_input_start_u8(const uint8_t* img, int C, int Hi, int Wi, int64_t sc, int64_t shs, int64_t sws,
                          int Ho, int Wo, const int* ry_s, const float* ry_w, int ky,
                          const int* rx_s, const float* rx_w, int kx, const float* mean, const float* stdv,
                          int patch, int rows_total, int k_pad, int nthreads, int nb, const int* edges,
                          uint16_t* const* outs) {
    return input_start_impl(nullptr, img, C, Hi, Wi, sc, shs, sws, Ho, Wo, ry_s, ry_w, ky, rx_s, rx_w, kx, mean, stdv,
                            patch, rows_total, k_pad, nthreads, nb, edges, outs);
}

// round 12: 0 if the vectorized u8 / 255 of the input ring equals the IEEE division for all 256 values
int moge_check_u8_div() {
#ifdef MOGE_AVX512
    alignas(64) uint8_t u[256];
    alignas(64) float q[256];
    for (int i = 0; i < 256; ++i) u[i] = static_cast<uint8_t>(i);
    const __m512 c = _mm512_set1_ps(1.0f / 255.0f), k255 = _mm512_set1_ps(255.0f);
    for (int j = 0; j < 256; j += 16) {
        const __m512 v = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(_mm_loadu_si128(reinterpret_cast<const __m128i*>(u + j))));
        const __m512 q0 = _mm512_mul_ps(v, c);
        _mm512_storeu_ps(q + j, _mm512_fmadd_ps(_mm512_fnmadd_ps(q0, k255, v), c, q0));
    }
    int bad = 0;
    for (int i = 0; i < 256; ++i) {
        const float want = static_cast<float>(i) / 255.0f;
        if (std::memcmp(&q[i], &want, 4) != 0) ++bad;
    }
    return bad;
#else
    return 0;
#endif
}

int moge_input_wait_band(void* h, int k) {
    auto* J = static_cast<InputJob*>(h);
    const int need = (J->edges[k + 1] - J->edges[k]) * J->P.patch;
    int spins = 0;
    while (J->done[k].load(std::memory_order_acquire) < need) {
        _mm_pause();
        if (++spins > 256) { std::this_thread::yield(); spins = 0; }
    }
    return 0;
}

int moge_input_join(void* h) {
    auto* J = static_cast<InputJob*>(h);
    if (J->driver.joinable()) J->driver.join();
    delete J;
    return 0;
}

int moge_input(const float* img, int C, int Hi, int Wi, int64_t sc, int64_t shs, int64_t sws,
               int Ho, int Wo, const int* ry_s, const float* ry_w, int ky,
               const int* rx_s, const float* rx_w, int kx, const float* mean, const float* stdv,
               int patch, uint16_t* out, int rows_total, int k_pad, int nthreads) {
    return moge_input_band(img, C, Hi, Wi, sc, shs, sws, Ho, Wo, ry_s, ry_w, ky, rx_s, rx_w, kx, mean, stdv,
                           patch, out, rows_total, k_pad, nthreads, 0, Ho / patch);
}

// ---------------------------------------------------------------------------------------------
// Served postprocess (server/fastpath.py): the tail of MoGeModel.infer() + the server's
// crop / resize / re-projection / masking, fused into one pass that writes the response arrays.
//
// pts/nrm/mprob : canvas outputs of TtMoGe (points affine [Hc,Wc,3], unit normals [Hc,Wc,3], mask
//                 probability [Hc,Wc]), fp32 contiguous.
// region        : x0,y0,x1,y1 crop on the canvas (the letterbox is outside it); Ho x Wo output size.
// shift, ms     : recovered z-shift and metric scale (postprocess.canvas_to_camera semantics:
//                 z' = z + shift; canvas mask &= z' > 0; depth = z' * ms; points = (x*ms, y*ms, z'*ms)).
// nidx_y/nidx_x : nearest-neighbour source rows / cols inside the crop (torch's 'nearest' indices,
//                 computed by the caller with F.interpolate on an arange).
// u/v           : utils3d uv_map linspaces (Wo / Ho), kinv: row-major 3x4 of inv(4x4 intrinsics).
// Outputs (Ho x Wo): depth f32, points f32 x3, normal x3 (f16, round to nearest even, if nrm_half, else
// f32), mask u8 (0/1).
// Bilinear resize = torch upsample_bilinear2d (align_corners=False, no antialias), evaluated as
// torch's build does: fma(fma(x00,w0,x01*w1), h0, fma(x10,w0,x11*w1)*h1). Normalize = x / max(||x||,
// 1e-12) with ||x|| = sqrt((a*a + b*b) + c*c) (the resized map is NCHW: torch's strided reduction). Re-projection = [u*d, v*d, d, 1] . kinv rows, summed
// left to right. stats (per call): [count, min, max] of the valid depths (mask == 1), double.
namespace {
inline uint16_t f32_to_f16(float f) {
#if defined(__F16C__)
    return static_cast<uint16_t>(_cvtss_sh(f, 0));
#else   // portable round-to-nearest-even (== F16C / numpy for every non-NaN float, checked exhaustively)
    uint32_t x; std::memcpy(&x, &f, 4);
    const uint16_t sgn = static_cast<uint16_t>((x >> 16) & 0x8000u);
    const uint32_t ax = x & 0x7fffffffu;
    if (ax > 0x7f800000u) return static_cast<uint16_t>(sgn | 0x7e00u | ((ax >> 13) & 0x3ffu));   // NaN
    if (ax >= 0x477ff000u) return static_cast<uint16_t>(sgn | 0x7c00u);                           // >= 65520 -> inf
    if (ax < 0x38800000u) {                                                                     // half subnormal
        float a; std::memcpy(&a, &ax, 4);
        return static_cast<uint16_t>(sgn | static_cast<uint16_t>(std::nearbyint(a * 16777216.0f)));
    }
    const uint32_t r = ax - 0x38000000u;
    uint32_t res = r >> 13;
    const uint32_t rem = r & 0x1fffu;
    if (rem > 0x1000u || (rem == 0x1000u && (res & 1u))) ++res;
    return static_cast<uint16_t>(sgn | res);
#endif
}
}  // namespace

// output rows [ya_all, yb_all) only (moge_finish = all rows); stats cover those rows
int moge_finish_rows(const float* pts, const float* nrm, const float* mprob, int Hc, int Wc,
                int x0, int y0, int x1, int y1, float shift, float ms, int Ho, int Wo,
                const int* nidx_y, const int* nidx_x, const float* u, const float* v, const float* kinv,
                int force_projection, int apply_mask,
                float* out_depth, float* out_pts, void* out_nrm, int nrm_half, uint8_t* out_mask,
                double* stats, int nthreads, int ya_all, int yb_all) {
    const int rh = y1 - y0, rw = x1 - x0;
    if (rh <= 0 || rw <= 0 || x0 < 0 || y0 < 0 || x1 > Wc || y1 > Hc) return -1;
    if (ya_all < 0 || yb_all > Ho || ya_all >= yb_all) return -1;
    const float sh = static_cast<float>(rh) / static_cast<float>(Ho);
    const float sw = static_cast<float>(rw) / static_cast<float>(Wo);
    std::vector<int> cx0(Wo), cx1(Wo);
    std::vector<float> cw0(Wo), cw1(Wo);
    for (int x = 0; x < Wo; ++x) lin_index(x, rw, sw, cx0[x], cx1[x], cw0[x], cw1[x]);
    const float inf = std::numeric_limits<float>::infinity();
    const int nrows = yb_all - ya_all;
    nthreads = std::max(1, std::min(nthreads, nrows));
    std::vector<double> st(3 * nthreads, 0.0);
    int chunk = (nrows + nthreads - 1) / nthreads;
    auto work = [&](int t) {
        int ya = ya_all + t * chunk, yb = std::min(yb_all, ya + chunk);
        double cnt = 0, mn = inf, mx = -inf;
        for (int yo = ya; yo < yb; ++yo) {
            int r0, r1; float h0, h1;
            lin_index(yo, rh, sh, r0, r1, h0, h1);
            r0 += y0; r1 += y0;
            const int rn = y0 + nidx_y[yo];
            const float vv = v[yo];
            for (int xo = 0; xo < Wo; ++xo) {
                const int c0 = x0 + cx0[xo], c1 = x0 + cx1[xo];
                const float w0 = cw0[xo], w1 = cw1[xo];
                const int64_t i00 = static_cast<int64_t>(r0) * Wc + c0, i01 = static_cast<int64_t>(r0) * Wc + c1;
                const int64_t i10 = static_cast<int64_t>(r1) * Wc + c0, i11 = static_cast<int64_t>(r1) * Wc + c1;
                auto bil = [&](float a00, float a01, float a10, float a11) {
                    return std::fma(std::fma(a00, w0, a01 * w1), h0, std::fma(a10, w0, a11 * w1) * h1);
                };
                // depth taps: (z + shift) * ms
                const float d00 = (pts[3 * i00 + 2] + shift) * ms, d01 = (pts[3 * i01 + 2] + shift) * ms;
                const float d10 = (pts[3 * i10 + 2] + shift) * ms, d11 = (pts[3 * i11 + 2] + shift) * ms;
                const float d = bil(d00, d01, d10, d11);
                // normal
                float n[3];
                for (int c = 0; c < 3; ++c) n[c] = bil(nrm[3 * i00 + c], nrm[3 * i01 + c], nrm[3 * i10 + c], nrm[3 * i11 + c]);
                float nn = std::sqrt((n[0] * n[0] + n[1] * n[1]) + n[2] * n[2]);   // torch: strided (NCHW) norm, unfused
                nn = std::max(nn, 1e-12f);
                for (int c = 0; c < 3; ++c) n[c] = n[c] / nn;
                // mask: nearest of the canvas mask (prob > 0.5 and z' > 0), then finite / positive depth
                const int64_t in_ = static_cast<int64_t>(rn) * Wc + x0 + nidx_x[xo];
                bool m = (mprob[in_] > 0.5f) && ((pts[3 * in_ + 2] + shift) > 0.f);
                m = m && std::isfinite(d) && (d > 0.f);
                float p[3];
                if (force_projection) {
                    const float uu = u[xo];
                    const float ud = uu * d, vd = vv * d;
                    for (int c = 0; c < 3; ++c) {
                        const float* k = kinv + 4 * c;
                        p[c] = ((ud * k[0] + vd * k[1]) + d * k[2]) + 1.f * k[3];
                    }
                } else {
                    for (int c = 0; c < 2; ++c)
                        p[c] = bil(pts[3 * i00 + c] * ms, pts[3 * i01 + c] * ms, pts[3 * i10 + c] * ms, pts[3 * i11 + c] * ms);
                    p[2] = d;   // bilinear of (z + shift) * ms == the depth map
                }
                const int64_t o = static_cast<int64_t>(yo) * Wo + xo;
                float dd = d;
                if (apply_mask && !m) {
                    dd = inf; p[0] = p[1] = p[2] = inf; n[0] = n[1] = n[2] = 0.f;
                }
                out_depth[o] = dd;
                out_pts[3 * o] = p[0]; out_pts[3 * o + 1] = p[1]; out_pts[3 * o + 2] = p[2];
                if (nrm_half) {
                    uint16_t* on = static_cast<uint16_t*>(out_nrm) + 3 * o;
                    on[0] = f32_to_f16(n[0]); on[1] = f32_to_f16(n[1]); on[2] = f32_to_f16(n[2]);
                } else {
                    float* on = static_cast<float*>(out_nrm) + 3 * o;
                    on[0] = n[0]; on[1] = n[1]; on[2] = n[2];
                }
                out_mask[o] = m ? 1 : 0;
                if (m) { cnt += 1; mn = std::min(mn, static_cast<double>(dd)); mx = std::max(mx, static_cast<double>(dd)); }
            }
        }
        st[3 * t] = cnt; st[3 * t + 1] = mn; st[3 * t + 2] = mx;
    };
    run_parallel(nthreads, work);
    double cnt = 0, mn = inf, mx = -inf;
    for (int t = 0; t < nthreads; ++t) { cnt += st[3 * t]; mn = std::min(mn, st[3 * t + 1]); mx = std::max(mx, st[3 * t + 2]); }
    stats[0] = cnt; stats[1] = mn; stats[2] = mx;
    return 0;
}

int moge_finish(const float* pts, const float* nrm, const float* mprob, int Hc, int Wc,
                int x0, int y0, int x1, int y1, float shift, float ms, int Ho, int Wo,
                const int* nidx_y, const int* nidx_x, const float* u, const float* v, const float* kinv,
                int force_projection, int apply_mask,
                float* out_depth, float* out_pts, void* out_nrm, int nrm_half, uint8_t* out_mask,
                double* stats, int nthreads) {
    return moge_finish_rows(pts, nrm, mprob, Hc, Wc, x0, y0, x1, y1, shift, ms, Ho, Wo, nidx_y, nidx_x, u, v, kinv,
                            force_projection, apply_mask, out_depth, out_pts, out_nrm, nrm_half, out_mask, stats,
                            nthreads, 0, Ho);
}

// k-th smallest (0-based) of the depths where mask == 1 (positive finite floats: their bit patterns
// sort like the values). Radix select: a parallel 16-bit histogram of the high halves, then an
// nth_element over the one bucket. Returns NaN when k is out of range.
float moge_masked_kth(const float* depth, const uint8_t* mask, int64_t n, int64_t k, int nthreads) {
    nthreads = std::max(1, nthreads);
    std::vector<std::vector<int64_t>> hist(nthreads, std::vector<int64_t>(65536, 0));
    int64_t chunk = (n + nthreads - 1) / nthreads;
    auto pass1 = [&](int t) {
        int64_t a = t * chunk, b = std::min(n, a + chunk);
        auto& h = hist[t];
        for (int64_t i = a; i < b; ++i) if (mask[i]) { uint32_t u; std::memcpy(&u, depth + i, 4); ++h[u >> 16]; }
    };
    run_parallel(nthreads, pass1);
    int64_t acc = 0; int bucket = -1; int64_t before = 0;
    for (int bkt = 0; bkt < 65536; ++bkt) {
        int64_t c = 0; for (int t = 0; t < nthreads; ++t) c += hist[t][bkt];
        if (k < acc + c) { bucket = bkt; before = acc; break; }
        acc += c;
    }
    if (bucket < 0 || k < 0) return std::numeric_limits<float>::quiet_NaN();
    std::vector<float> vals;
    for (int64_t i = 0; i < n; ++i)
        if (mask[i]) { uint32_t u; std::memcpy(&u, depth + i, 4); if (static_cast<int>(u >> 16) == bucket) vals.push_back(depth[i]); }
    auto kk = static_cast<size_t>(k - before);
    std::nth_element(vals.begin(), vals.begin() + kk, vals.end());
    return vals[kk];
}

// standard base64 (RFC 4648 alphabet, '=' padding) of n bytes into out (4 * ceil(n / 3) chars)
int moge_b64(const uint8_t* src, int64_t n, char* out, int nthreads) {
    static const char tbl[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    const int64_t groups = n / 3;
    nthreads = std::max<int64_t>(1, std::min<int64_t>(nthreads, groups / 65536 + 1));
    int64_t chunk = (groups + nthreads - 1) / nthreads;
    auto work = [&](int t) {
        int64_t a = t * chunk, b = std::min(groups, a + chunk);
        const uint8_t* s = src + 3 * a;
        char* o = out + 4 * a;
        for (int64_t g = a; g < b; ++g, s += 3, o += 4) {
            uint32_t w = (uint32_t(s[0]) << 16) | (uint32_t(s[1]) << 8) | s[2];
            o[0] = tbl[w >> 18]; o[1] = tbl[(w >> 12) & 63]; o[2] = tbl[(w >> 6) & 63]; o[3] = tbl[w & 63];
        }
    };
    run_parallel(nthreads, work);
    const int64_t rem = n - 3 * groups;
    if (rem) {
        const uint8_t* s = src + 3 * groups;
        char* o = out + 4 * groups;
        uint32_t w = uint32_t(s[0]) << 16;
        if (rem == 2) w |= uint32_t(s[1]) << 8;
        o[0] = tbl[w >> 18]; o[1] = tbl[(w >> 12) & 63];
        o[2] = rem == 2 ? tbl[(w >> 6) & 63] : '=';
        o[3] = '=';
    }
    return 0;
}

// base64 of the concatenation of nseg byte segments (ptrs[i], lens[i]) without gathering them first:
// each thread encodes a range of 3-byte groups, walking the segments with a cursor (groups that straddle
// a segment boundary are assembled byte by byte). Same output as moge_b64 on the joined bytes.
int moge_b64_segs(const uint8_t* const* ptrs, const int64_t* lens, int nseg, char* out, int nthreads) {
    static const char tbl[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    std::vector<int64_t> start(nseg + 1, 0);
    for (int i = 0; i < nseg; ++i) start[i + 1] = start[i] + lens[i];
    const int64_t n = start[nseg];
    const int64_t groups = n / 3;
    nthreads = std::max<int64_t>(1, std::min<int64_t>(nthreads, groups / 65536 + 1));
    const int64_t chunk = (groups + nthreads - 1) / nthreads;
    auto work = [&](int t) {
        int64_t g = t * chunk;
        const int64_t gend = std::min(groups, g + chunk);
        if (g >= gend) return;
        int64_t pos = 3 * g;
        int s = int(std::upper_bound(start.begin(), start.end(), pos) - start.begin()) - 1;
        char* o = out + 4 * g;
        while (g < gend) {
            while (s < nseg && pos >= start[s + 1]) ++s;
            const int64_t off = pos - start[s];
            const int64_t avail = start[s + 1] - pos;
            int64_t bulk = std::min(avail / 3, gend - g);
            if (bulk > 0) {
                const uint8_t* p = ptrs[s] + off;
                for (int64_t k = 0; k < bulk; ++k, p += 3, o += 4) {
                    uint32_t w = (uint32_t(p[0]) << 16) | (uint32_t(p[1]) << 8) | p[2];
                    o[0] = tbl[w >> 18]; o[1] = tbl[(w >> 12) & 63]; o[2] = tbl[(w >> 6) & 63]; o[3] = tbl[w & 63];
                }
                g += bulk; pos += 3 * bulk;
                continue;
            }
            uint8_t b[3];
            int ss = s;
            for (int j = 0; j < 3; ++j) {
                const int64_t q = pos + j;
                while (q >= start[ss + 1]) ++ss;
                b[j] = ptrs[ss][q - start[ss]];
            }
            uint32_t w = (uint32_t(b[0]) << 16) | (uint32_t(b[1]) << 8) | b[2];
            o[0] = tbl[w >> 18]; o[1] = tbl[(w >> 12) & 63]; o[2] = tbl[(w >> 6) & 63]; o[3] = tbl[w & 63];
            o += 4; ++g; pos += 3;
        }
    };
    run_parallel(nthreads, work);
    const int64_t rem = n - 3 * groups;
    if (rem) {
        uint8_t b[2] = {0, 0};
        int ss = 0;
        for (int64_t j = 0; j < rem; ++j) {
            const int64_t q = 3 * groups + j;
            while (q >= start[ss + 1]) ++ss;
            b[j] = ptrs[ss][q - start[ss]];
        }
        char* o = out + 4 * groups;
        uint32_t w = (uint32_t(b[0]) << 16) | (uint32_t(b[1]) << 8);
        o[0] = tbl[w >> 18]; o[1] = tbl[(w >> 12) & 63];
        o[2] = rem == 2 ? tbl[(w >> 6) & 63] : '=';
        o[3] = '=';
    }
    return 0;
}

}  // extern "C"
namespace {
#if defined(__SSE4_1__)
template <int BPP>
inline void paeth_sse(const uint8_t* s, const uint8_t* p, uint8_t* c, int64_t rowbytes) {
    const __m128i zero = _mm_setzero_si128();
    __m128i a = zero, cc = zero;
    for (int64_t i = 0; i < rowbytes; i += BPP) {
        uint32_t ub, us;        // 4-byte loads / stores (3-byte memcpy stalls): buffers carry 4 bytes of padding
        std::memcpy(&ub, p + i, 4);
        std::memcpy(&us, s + i, 4);
        const __m128i b = _mm_unpacklo_epi8(_mm_cvtsi32_si128(static_cast<int>(ub)), zero);
        const __m128i x = _mm_unpacklo_epi8(_mm_cvtsi32_si128(static_cast<int>(us)), zero);
        __m128i pa = _mm_sub_epi16(b, cc), pb = _mm_sub_epi16(a, cc);
        __m128i pc = _mm_abs_epi16(_mm_add_epi16(pa, pb));
        pa = _mm_abs_epi16(pa);
        pb = _mm_abs_epi16(pb);
        const __m128i sm = _mm_min_epi16(pc, _mm_min_epi16(pa, pb));
        const __m128i pred = _mm_blendv_epi8(_mm_blendv_epi8(cc, b, _mm_cmpeq_epi16(sm, pb)), a,
                     _mm_cmpeq_epi16(sm, pa));
        a = _mm_and_si128(_mm_add_epi8(x, pred), _mm_set1_epi16(0xFF));
        const uint32_t r = static_cast<uint32_t>(_mm_cvtsi128_si32(_mm_packus_epi16(a, a)));
        std::memcpy(c + i, &r, 4);
        cc = b;
    }
}
#endif

}  // namespace
extern "C" {

// PNG scanline unfilter (filters 0..4, 8-bit samples, non-interlaced): src = the inflated IDAT stream,
// H rows of (1 + rowbytes) bytes; dst = H x rowbytes pixels, then (out_ch < bpp) the first out_ch
// channels of every pixel are kept in place (RGBA -> RGB, as PIL's convert('RGB')). Returns -1 on a bad
// filter type.
namespace {
struct Unfilter {
    int64_t rowbytes;
    int bpp, out_ch;
    std::vector<uint8_t> prev, cur, srow;
    Unfilter(int64_t rb, int b, int oc) : rowbytes(rb), bpp(b), out_ch(oc), prev(rb + 4, 0), cur(rb + 4, 0), srow(rb + 4, 0) {}
    // src: nrows filtered rows (filter byte + rowbytes each); dst: nrows output rows (W * out_ch)
    int rows(const uint8_t* src, int nrows, uint8_t* dst) {
        for (int y = 0; y < nrows; ++y) {
            const uint8_t* s = src + static_cast<int64_t>(y) * (rowbytes + 1);
            const int ft = s[0];
            ++s;
            uint8_t* c = cur.data();
            const uint8_t* p = prev.data();
            switch (ft) {
                case 0: std::memcpy(c, s, rowbytes); break;
                case 1:
                    for (int64_t i = 0; i < rowbytes; ++i) c[i] = static_cast<uint8_t>(s[i] + (i >= bpp ? c[i - bpp] : 0));
                    break;
                case 2:
                    for (int64_t i = 0; i < rowbytes; ++i) c[i] = static_cast<uint8_t>(s[i] + p[i]);
                    break;
                case 3:
                    for (int64_t i = 0; i < rowbytes; ++i)
                        c[i] = static_cast<uint8_t>(s[i] + ((static_cast<int>(i >= bpp ? c[i - bpp] : 0) + p[i]) >> 1));
                    break;
                case 4: {
#if defined(__SSE4_1__)
                    // libpng's SSE paeth: one pixel (bpp <= 4 bytes) per step as 16-bit lanes;
                    // |p-a| = |b-c|, |p-b| = |a-c|, |p-c| = |(b-c)+(a-c)|; ties a, then b, then c.
                    std::memcpy(srow.data(), s, rowbytes);   // padded copy of the filtered row
                    if (bpp == 3) paeth_sse<3>(srow.data(), p, c, rowbytes);
                    else if (bpp == 4) paeth_sse<4>(srow.data(), p, c, rowbytes);
                    else return -1;
#else
                    for (int64_t i = 0; i < rowbytes; ++i) {
                        const int a = i >= bpp ? c[i - bpp] : 0, b = p[i], cc = i >= bpp ? p[i - bpp] : 0;
                        const int pp = a + b - cc, pa = std::abs(pp - a), pb = std::abs(pp - b), pc = std::abs(pp - cc);
                        const int pred = (pa <= pb && pa <= pc) ? a : (pb <= pc ? b : cc);
                        c[i] = static_cast<uint8_t>(s[i] + pred);
                    }
#endif
                    break;
                }
                default: return -1;
            }
            const int64_t W = rowbytes / bpp;
            uint8_t* d = dst + static_cast<int64_t>(y) * W * out_ch;
            if (out_ch == bpp) std::memcpy(d, c, rowbytes);
            else for (int64_t x = 0; x < W; ++x) for (int k = 0; k < out_ch; ++k) d[x * out_ch + k] = c[x * bpp + k];
            std::swap(prev, cur);
        }
        return 0;
    }
};
}  // namespace
extern "C" {
int moge_png_unfilter(const uint8_t* src, int H, int64_t rowbytes, int bpp, int out_ch, uint8_t* dst) {
    Unfilter u(rowbytes, bpp, out_ch);
    return u.rows(src, H, dst);
}

// incremental form (rows unfiltered while the IDAT stream is still being inflated)
void* moge_unf_new(int64_t rowbytes, int bpp, int out_ch) { return new Unfilter(rowbytes, bpp, out_ch); }
int moge_unf_rows(void* st, const uint8_t* src, int nrows, uint8_t* dst) { return static_cast<Unfilter*>(st)->rows(src, nrows, dst); }
void moge_unf_free(void* st) { delete static_cast<Unfilter*>(st); }

}  // extern "C"
// PNG IDAT inflate with ISA-L (isal_inflate through function pointers the caller resolves from
// python-isal's shared object; the inflate_state layout of ISA-L 2.31 -- offsets checked by the caller's
// version gate and self test) pipelined with the serial Unfilter on a std::thread that follows the inflated
// row count. No GIL, no per-chunk Python objects. Returns 0, or -1 (bad stream / checksum / size: the
// caller falls back to PIL). z: H * (rowbytes + 1) scratch; dst: H x W x 3.
namespace {
constexpr size_t kIsalStateBytes = 87368;     // sizeof(struct inflate_state), ISA-L 2.31
constexpr size_t kIsalOffTotalOut = 12, kIsalOffNextIn = 16, kIsalOffAvailIn = 32;
constexpr size_t kIsalOffBlockState = 21160, kIsalOffCrcFlag = 21172;
constexpr int kIsalZlib = 3, kIsalBlockFinish = 5;
}  // namespace
extern "C" {
int64_t moge_isal_state_bytes() { return static_cast<int64_t>(kIsalStateBytes); }
int moge_png_inflate_unfilter(void* fn_init, void* fn_inflate, const uint8_t* const* chunks, const int64_t* lens,
                              int nchunks, int H, int64_t rowbytes, int bpp, uint8_t* z, uint8_t* dst,
                              double* times_us) {
    const auto t_start = std::chrono::steady_clock::now();
    auto us = [&] { return std::chrono::duration<double, std::micro>(std::chrono::steady_clock::now() - t_start).count(); };
    using InitFn = void (*)(void*);
    using InflFn = int (*)(void*);
    auto init = reinterpret_cast<InitFn>(fn_init);
    auto infl = reinterpret_cast<InflFn>(fn_inflate);
    thread_local std::vector<uint64_t> stbuf;
    if (stbuf.size() * 8 < kIsalStateBytes + 64) stbuf.assign((kIsalStateBytes + 64) / 8 + 1, 0);
    uint8_t* st = reinterpret_cast<uint8_t*>(stbuf.data());
    init(st);
    std::memcpy(st + kIsalOffCrcFlag, &kIsalZlib, 4);
    const int64_t stride = rowbytes + 1, total = static_cast<int64_t>(H) * stride;
    uint8_t* next_out = z;
    uint32_t avail_out = static_cast<uint32_t>(std::min<int64_t>(total, 0xFFFFFFFFll));
    std::memcpy(st + 0, &next_out, sizeof(next_out));
    std::memcpy(st + 8, &avail_out, 4);
    std::atomic<int> ready{0}, stop{0};
    Unfilter u(rowbytes, bpp, 3);
    int unf_rc = 0;
    // the worker sleeps on a condition variable between published row counts (a spinning worker that the
    // scheduler put on the inflating thread's core doubled the decode time on the loaded host)
    std::mutex mu;
    std::condition_variable cv;
    std::thread worker([&] {
        int done = 0;
        for (;;) {
            int r = ready.load(std::memory_order_acquire);
            if (r <= done) {
                if (done >= H || stop.load(std::memory_order_acquire)) return;
                std::unique_lock<std::mutex> l(mu);
                cv.wait(l, [&] { return ready.load(std::memory_order_acquire) > done || stop.load(std::memory_order_acquire); });
                continue;
            }
            const int n = std::min(r - done, 64);
            if (u.rows(z + static_cast<int64_t>(done) * stride, n, dst + static_cast<int64_t>(done) * (rowbytes / bpp) * 3)) { unf_rc = -1; return; }
            done += n;
        }
    });
    int rc = 0;
    auto publish = [&] {
        uint32_t tot;
        std::memcpy(&tot, st + kIsalOffTotalOut, 4);
        const int r = static_cast<int>(std::min<int64_t>(tot, total) / stride);
        if (r > ready.load(std::memory_order_relaxed)) {
            { std::lock_guard<std::mutex> l(mu); ready.store(r, std::memory_order_release); }
            cv.notify_one();
        }
    };
    thread_local std::vector<uint8_t> discard(1 << 16);
    auto overflow = [&] {   // output full: the image is complete, keep inflating into a discard buffer
        uint32_t ao0;
        std::memcpy(&ao0, st + 8, 4);
        if (ao0 != 0) return false;
        uint32_t tot;
        std::memcpy(&tot, st + kIsalOffTotalOut, 4);
        if (static_cast<int64_t>(tot) < total) return false;
        publish();
        uint8_t* no = discard.data();
        const uint32_t ao = static_cast<uint32_t>(discard.size());
        std::memcpy(st + 0, &no, sizeof(no));
        std::memcpy(st + 8, &ao, 4);
        return true;
    };
    for (int k = 0; k < nchunks && rc == 0; ++k) {
        const uint8_t* ni = chunks[k];
        uint32_t ai = static_cast<uint32_t>(lens[k]);
        std::memcpy(st + kIsalOffNextIn, &ni, sizeof(ni));
        std::memcpy(st + kIsalOffAvailIn, &ai, 4);
        uint32_t prev_left = ai + 1;
        while (rc == 0) {
            const int r = infl(st);
            if ((r == 0 || r == 2) && overflow()) continue;   // bytes beyond the image: inflated and ignored
            if (r != 0 && r != 1) { rc = -1; break; }      // ISAL_DECOMP_OK / ISAL_END_INPUT; errors -> fallback
            publish();
            uint32_t left;
            std::memcpy(&left, st + kIsalOffAvailIn, 4);
            int bs;
            std::memcpy(&bs, st + kIsalOffBlockState, 4);
            if (left == 0 || bs == kIsalBlockFinish) break;
            if (left >= prev_left) { rc = -1; break; }   // no progress: bail out to the fallback
            prev_left = left;
        }
    }
    if (rc == 0) {   // flush anything still buffered (no more input)
        for (int it = 0; it < 4; ++it) {
            int bs;
            std::memcpy(&bs, st + kIsalOffBlockState, 4);
            if (bs == kIsalBlockFinish) break;
            uint32_t ai = 0;
            std::memcpy(st + kIsalOffAvailIn, &ai, 4);
            const int r = infl(st);
            if ((r == 0 || r == 2) && overflow()) continue;
            if (r != 0 && r != 1) { rc = -1; break; }
            publish();
        }
        int bs;
        uint32_t tot;
        std::memcpy(&bs, st + kIsalOffBlockState, 4);
        std::memcpy(&tot, st + kIsalOffTotalOut, 4);
        if (bs != kIsalBlockFinish || static_cast<int64_t>(tot) < total) rc = -1;
    }
    if (times_us) times_us[0] = us();          // inflate done
    if (rc) {
        { std::lock_guard<std::mutex> l(mu); stop.store(1, std::memory_order_release); }
        cv.notify_one();
    }
    worker.join();
    if (times_us) times_us[1] = us();          // unfilter done
    return rc ? rc : unf_rc;
}
}  // extern "C"

extern "C" {

// base64 decode of a CLEAN standard-alphabet string (length % 4 == 0, '=' only as 1-2 trailing pad
// characters): returns the decoded length, or -1 for anything else (the caller then uses Python's
// base64.b64decode, whose non-validating mode skips foreign characters). On clean input the result
// equals base64.b64decode(validate=False) (leftover pad bits are ignored as in binascii).
int64_t moge_b64dec(const uint8_t* s, int64_t n, uint8_t* out, int nthreads) {
    if (n % 4) return -1;
    if (n == 0) return 0;
    int pad = 0;
    if (s[n - 1] == '=') { pad = 1; if (s[n - 2] == '=') pad = 2; }
    static int8_t tbl[256];
    static bool init = false;
    if (!init) {
        for (int i = 0; i < 256; ++i) tbl[i] = -1;
        const char* a = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
        for (int i = 0; i < 64; ++i) tbl[static_cast<uint8_t>(a[i])] = static_cast<int8_t>(i);
        init = true;
    }
    const int64_t groups = n / 4;
    const int64_t full = pad ? groups - 1 : groups;
    nthreads = std::max<int64_t>(1, std::min<int64_t>(nthreads, full / 65536 + 1));
    const int64_t chunk = (full + nthreads - 1) / nthreads;
    std::vector<int> bad(nthreads, 0);
    auto work = [&](int t) {
        const int64_t a = t * chunk, b = std::min(full, a + chunk);
        const uint8_t* p = s + 4 * a;
        uint8_t* o = out + 3 * a;
        int acc = 0;
        for (int64_t g = a; g < b; ++g, p += 4, o += 3) {
            const int v0 = tbl[p[0]], v1 = tbl[p[1]], v2 = tbl[p[2]], v3 = tbl[p[3]];
            acc |= v0 | v1 | v2 | v3;
            const uint32_t w = (uint32_t(v0 & 63) << 18) | (uint32_t(v1 & 63) << 12) | (uint32_t(v2 & 63) << 6) | uint32_t(v3 & 63);
            o[0] = uint8_t(w >> 16); o[1] = uint8_t(w >> 8); o[2] = uint8_t(w);
        }
        bad[t] = acc < 0;
    };
    run_parallel(nthreads, work);
    for (int t = 0; t < nthreads; ++t) if (bad[t]) return -1;
    int64_t len = 3 * full;
    if (pad) {
        const uint8_t* p = s + 4 * full;
        const int v0 = tbl[p[0]], v1 = tbl[p[1]];
        const int v2 = pad == 1 ? tbl[p[2]] : 0;
        if (v0 < 0 || v1 < 0 || v2 < 0) return -1;
        const uint32_t w = (uint32_t(v0) << 18) | (uint32_t(v1) << 12) | (uint32_t(v2) << 6);
        out[len++] = uint8_t(w >> 16);
        if (pad == 1) out[len++] = uint8_t(w >> 8);
    }
    return len;
}
}  // extern "C"

// ---------------------------------------------------------------------------------------------
// Served preprocess: postprocess.fit_to_canvas in one pass. Pillow's BILINEAR resize of an 8-bit RGB
// image (ImagingResample: separable, double-precision triangle-filter coefficients normalized per
// output pixel and converted to 22-bit fixed point, horizontal pass into uint8 rows then vertical
// pass, round-half-up + clip), pasted at (x0, y0) on a canvas filled with `pad`, converted to float32
// as u8 / 255.0f and written planar NCHW [3, Ch, Cw]. rw x rh = resized size (== Wi x Hi -> copy).
namespace {
constexpr int kPrec = 22;
inline uint8_t clip8(int v) { v >>= kPrec; return static_cast<uint8_t>(v < 0 ? 0 : (v > 255 ? 255 : v)); }
int pil_coeffs(int in_size, int out_size, std::vector<int>& bounds, std::vector<int32_t>& kk) {
    const double in0 = 0.0, in1 = static_cast<double>(in_size);
    const double scale = (in1 - in0) / out_size;
    double filterscale = scale < 1.0 ? 1.0 : scale;
    const double support = 1.0 * filterscale;
    const int ksize = static_cast<int>(std::ceil(support)) * 2 + 1;
    bounds.assign(2 * out_size, 0);
    kk.assign(static_cast<size_t>(out_size) * ksize, 0);
    std::vector<double> k(ksize);
    for (int xx = 0; xx < out_size; ++xx) {
        const double center = in0 + (xx + 0.5) * scale;
        double ww = 0.0;
        const double ss = 1.0 / filterscale;
        int xmin = static_cast<int>(center - support + 0.5);
        if (xmin < 0) xmin = 0;
        int xmax = static_cast<int>(center + support + 0.5);
        if (xmax > in_size) xmax = in_size;
        xmax -= xmin;
        for (int x = 0; x < xmax; ++x) {
            double t = (x + xmin - center + 0.5) * ss;
            if (t < 0.0) t = -t;
            const double w = t < 1.0 ? 1.0 - t : 0.0;
            k[x] = w;
            ww += w;
        }
        for (int x = 0; x < xmax; ++x) if (ww != 0.0) k[x] /= ww;
        for (int x = xmax; x < ksize; ++x) k[x] = 0;
        for (int x = 0; x < ksize; ++x) {
            const double v = k[x] * (1 << kPrec);
            kk[static_cast<size_t>(xx) * ksize + x] = static_cast<int32_t>(k[x] < 0 ? -0.5 + v : 0.5 + v);
        }
        bounds[2 * xx] = xmin;
        bounds[2 * xx + 1] = xmax;
    }
    return ksize;
}
}  // namespace

// U8 = false: float32 planar NCHW [3, Ch, Cw] = u8 / 255.0f (pad: the float pad value);
// U8 = true (round 12): the uint8 canvas itself, interleaved HWC [Ch, Cw, 3] (pad: the gray level), i.e. the
// float32 canvas before the / 255 (the input prep reads it as u8 / 255.0f: bit-identical model input)
extern "C++" {
template <bool U8>
static int fit_canvas_impl(const uint8_t* img, int Hi, int Wi, int64_t row_stride, int rw, int rh, int x0, int y0,
                           int Cw, int Ch, float pad, void* out_v, int nthreads) {
    if (x0 < 0 || y0 < 0 || x0 + rw > Cw || y0 + rh > Ch || rw < 1 || rh < 1) return -1;
    const bool need_h = rw != Wi, need_v = rh != Hi;
    std::vector<int> bh, bv;
    std::vector<int32_t> kh, kv;
    const int ksh = need_h ? pil_coeffs(Wi, rw, bh, kh) : 0;
    const int ksv = need_v ? pil_coeffs(Hi, rh, bv, kv) : 0;
    // horizontal pass over the source rows the vertical pass needs (all of them here: box = full image)
    const int ybox_first = need_v ? bv[0] : 0;
    const int ybox_last = need_v ? bv[2 * rh - 2] + bv[2 * rh - 1] : Hi;
    const int tmp_rows = ybox_last - ybox_first;
    std::vector<uint8_t> tmp;
    const uint8_t* src = img;
    int64_t src_stride = row_stride;
    if (need_h) {
        tmp.resize(static_cast<size_t>(tmp_rows) * rw * 3);
        parallel_rows(tmp_rows, nthreads, [&](int a, int b) {
            for (int yy = a; yy < b; ++yy) {
                const uint8_t* r = img + static_cast<int64_t>(yy + ybox_first) * row_stride;
                uint8_t* o = tmp.data() + static_cast<size_t>(yy) * rw * 3;
                for (int xx = 0; xx < rw; ++xx) {
                    const int xmin = bh[2 * xx], xmax = bh[2 * xx + 1];
                    const int32_t* k = kh.data() + static_cast<size_t>(xx) * ksh;
                    int s0 = 1 << (kPrec - 1), s1 = s0, s2 = s0;
                    const uint8_t* px = r + 3 * xmin;
                    if (ksh == 3 && xmin + 3 <= Wi) {   // upscale: fixed 3 taps (the zero-padded ones add 0)
                        s0 += px[0] * k[0] + px[3] * k[1] + px[6] * k[2];
                        s1 += px[1] * k[0] + px[4] * k[1] + px[7] * k[2];
                        s2 += px[2] * k[0] + px[5] * k[1] + px[8] * k[2];
                    } else {
                        for (int x = 0; x < xmax; ++x, px += 3) {
                            s0 += px[0] * k[x]; s1 += px[1] * k[x]; s2 += px[2] * k[x];
                        }
                    }
                    o[3 * xx] = clip8(s0); o[3 * xx + 1] = clip8(s1); o[3 * xx + 2] = clip8(s2);
                }
            }
        });
        src = tmp.data();
        src_stride = static_cast<int64_t>(rw) * 3;
    } else if (need_v) {
        src = img + static_cast<int64_t>(ybox_first) * row_stride;   // rows relative to ybox_first
    }
    const int64_t plane = static_cast<int64_t>(Ch) * Cw;
    float lut[256];
    for (int v = 0; v < 256; ++v) lut[v] = static_cast<float>(v) / 255.0f;   // == float32(u8) / 255
    if constexpr (U8) {
        uint8_t* out8 = static_cast<uint8_t*>(out_v);
        const uint8_t pad8 = static_cast<uint8_t>(pad);
        parallel_rows(Ch, nthreads, [&](int a, int b) {
            std::vector<int32_t> acc;
            for (int y = a; y < b; ++y) {
                uint8_t* o = out8 + static_cast<int64_t>(y) * Cw * 3;
                const int yy = y - y0;
                if (yy < 0 || yy >= rh) { std::memset(o, pad8, static_cast<size_t>(Cw) * 3); continue; }
                std::memset(o, pad8, static_cast<size_t>(x0) * 3);
                std::memset(o + static_cast<size_t>(x0 + rw) * 3, pad8, static_cast<size_t>(Cw - x0 - rw) * 3);
                uint8_t* ox = o + static_cast<size_t>(x0) * 3;
                if (need_v) {
                    const int ymin = bv[2 * yy] - ybox_first, ymax = bv[2 * yy + 1];
                    const int32_t* k = kv.data() + static_cast<size_t>(yy) * ksv;
                    const int n = 3 * rw;
                    acc.assign(n, 1 << (kPrec - 1));
                    for (int t = 0; t < ymax; ++t) {
                        const uint8_t* r = src + static_cast<int64_t>(ymin + t) * src_stride;
                        const int32_t kt = k[t];
                        int32_t* ap = acc.data();
                        for (int i = 0; i < n; ++i) ap[i] += r[i] * kt;
                    }
                    for (int i = 0; i < n; ++i) ox[i] = clip8(acc[i]);
                } else {
                    std::memcpy(ox, src + static_cast<int64_t>(yy) * src_stride, static_cast<size_t>(rw) * 3);
                }
            }
        });
        return 0;
    }
    float* out = static_cast<float*>(out_v);
    parallel_rows(Ch, nthreads, [&](int a, int b) {
        std::vector<int32_t> acc;
        for (int y = a; y < b; ++y) {
            float* o0 = out + static_cast<int64_t>(y) * Cw;
            float* o1 = o0 + plane;
            float* o2 = o1 + plane;
            const int yy = y - y0;
            if (yy < 0 || yy >= rh) {
                for (int x = 0; x < Cw; ++x) { o0[x] = pad; o1[x] = pad; o2[x] = pad; }
                continue;
            }
            for (int x = 0; x < x0; ++x) { o0[x] = pad; o1[x] = pad; o2[x] = pad; }
            for (int x = x0 + rw; x < Cw; ++x) { o0[x] = pad; o1[x] = pad; o2[x] = pad; }
            if (need_v) {
                const int ymin = bv[2 * yy] - ybox_first, ymax = bv[2 * yy + 1];
                const int32_t* k = kv.data() + static_cast<size_t>(yy) * ksv;
                const int n = 3 * rw;
                acc.assign(n, 1 << (kPrec - 1));
                for (int t = 0; t < ymax; ++t) {        // row-vector accumulation (vectorizes over x)
                    const uint8_t* r = src + static_cast<int64_t>(ymin + t) * src_stride;
                    const int32_t kt = k[t];
                    int32_t* ap = acc.data();
                    for (int i = 0; i < n; ++i) ap[i] += r[i] * kt;
                }
                for (int xx = 0; xx < rw; ++xx) {
                    o0[x0 + xx] = lut[clip8(acc[3 * xx])];
                    o1[x0 + xx] = lut[clip8(acc[3 * xx + 1])];
                    o2[x0 + xx] = lut[clip8(acc[3 * xx + 2])];
                }
            } else {
                const uint8_t* r = src + static_cast<int64_t>(yy) * src_stride;
                for (int xx = 0; xx < rw; ++xx) {
                    o0[x0 + xx] = lut[r[3 * xx]];
                    o1[x0 + xx] = lut[r[3 * xx + 1]];
                    o2[x0 + xx] = lut[r[3 * xx + 2]];
                }
            }
        }
    });
    return 0;
}

}  // extern "C++"

int moge_fit_canvas(const uint8_t* img, int Hi, int Wi, int64_t row_stride, int rw, int rh, int x0, int y0,
                    int Cw, int Ch, float pad, float* out, int nthreads) {
    return fit_canvas_impl<false>(img, Hi, Wi, row_stride, rw, rh, x0, y0, Cw, Ch, pad, out, nthreads);
}

int moge_fit_canvas_u8(const uint8_t* img, int Hi, int Wi, int64_t row_stride, int rw, int rh, int x0, int y0,
                       int Cw, int Ch, float pad, uint8_t* out, int nthreads) {
    return fit_canvas_impl<true>(img, Hi, Wi, row_stride, rw, rh, x0, y0, Cw, Ch, pad, out, nthreads);
}

}  // extern "C"