#version 460
// d4r engine, preset M: the input stage, written from the algorithm in M_NOTES.md. Per render pixel it gathers
// 20 features (the colour, four colours of the previous result along three motion candidates, the recurrent
// feature, a motion-change measure), rounds them to FP8 with NGX's dither, and embeds each 2x2-pixel token
// (128 feature slots) into 64 channels with one matrix product per 4x4-token block. It also writes the three
// motion candidates the reconstruction uses.
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_float_e4m3 : require
#extension GL_GOOGLE_include_directive : require
#include "round_half.glsl"
layout(local_size_x = 32) in;       // one wave: 16 tokens = 8x8 render pixels, two pixels per lane
layout(set = 0, binding = 0, std430) readonly buffer Params
{
    ivec2 size;             // render size
    ivec2 shift;            // pixels the block grid is shifted by
    ivec2 tok;              // token grid
    vec2 jitter, prevJitter;
    vec2 mvScale, mvOffs;   // motion texel -> pixels of the 3x image
    vec2 rnd2hi, hi2rnd;
    vec2 hiSize;
    float changeScale, changeGain;
    uint depthFlag, reset, frame;
    float expo, expoPrev;
    float halfInvX, halfInvY, renderInvX, renderInvY; // host-rounded CUDA texture coordinates
} p;
layout(set = 0, binding = 1) uniform sampler2D colorTex;
layout(set = 0, binding = 2) uniform sampler2D motionTex;
layout(set = 0, binding = 3) uniform sampler2D depthTex;
layout(set = 0, binding = 4) uniform sampler2D prevMvTex;      // candidate A of the previous frame, bilinear
layout(set = 0, binding = 5) uniform sampler2D halfTex;        // half-resolution previous result, bilinear
layout(set = 0, binding = 6) uniform sampler2D featTex;        // recurrent feature, bilinear
layout(set = 0, binding = 7, rg16f) uniform writeonly image2D mvAOut;
layout(set = 0, binding = 8, rg16f) uniform writeonly image2D mvBOut;
layout(set = 0, binding = 9, rg16f) uniform writeonly image2D mvCOut;
layout(set = 0, binding = 10, std430) readonly buffer W8 { floate4m3_t w[]; } w8;      // [128][64]
layout(set = 0, binding = 11, std430) readonly buffer W16 { f16vec4 w[]; } w16;        // bias, 16 equal rows of 64
layout(set = 0, binding = 12, std430) readonly buffer Noise { uint16_t n[]; } noise;   // 8 x 16 x 16
layout(set = 0, binding = 13, std430) writeonly buffer Out { floate4m3_t a[]; } outb;  // [tok.y][tok.x][64]
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HACC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define RM gl_CooperativeMatrixLayoutRowMajor
// NVIDIA's embedding GEMM is chained m16n8k32 mma with an f16 accumulator: the bias is the initial C and every
// 32 products are summed with one f16 rounding (swin_model.py mma_chain; the same k32 step the native kernels'
// SWIN_EXACT build implements). K32_ROUND reproduces that: after each pair of k16 cooperative products the
// accumulator is rounded to f16 through a per-invocation f16 scratch slot (ROUND_HALF_SHARED).
#ifndef K32_ROUND
#define K32_ROUND 1
#endif
#if K32_ROUND
#define R32(m) ROUND_HALF_SHARED(m, RS)
#else
#define R32(m)
#endif
shared floate4m3_t FE[16 * 136];
shared float16_t RS[gl_WorkGroupSize.x * 8u];
const float L10 = 3.3219281;
// LDR (NGX's *_ldr_* kernels, the game did not set the HDR flag): display-referred colour, no exposure and no
// range compression anywhere.
#ifndef LDR
#define LDR 0
#endif
// PLANES: the tokens in NGX's layout, two planes of 32 channels ([2][tok.y][tok.x][32]), for the HIP network
#ifndef PLANES
#define PLANES 0
#endif
#if LDR
float compress1(float v) { return v; }
#else
float compress1(float v) {
    precise float mapped = fma(v, 100.0, 1.0);
    precise float lg = log2(max(mapped, 1e-4));
    precise float base10 = lg * 0.30103;
    precise float result = base10 * 0.36;
    return result;
}
#endif
vec3 compress(vec3 v) { return vec3(compress1(v.x), compress1(v.y), compress1(v.z)); }
#if LDR
float colour1(float v) { return isinf(v) || isnan(v) ? 0.0 : min(v, 32752.0); }
#else
float colour1(float v) { return compress1(isinf(v) || isnan(v) ? 0.0 : min(v, 32752.0)); }
#endif
ivec2 texel(ivec2 t) { return ivec2(t.x > 0 ? min(t.x, p.size.x - 1) : 0, t.y > 0 ? min(t.y, p.size.y - 1) : 0); }
vec2 motion(ivec2 t) { return p.mvScale * (p.mvOffs + texelFetch(motionTex, texel(t), 0).xy); }
// NGX's dithered truncation of an f16 value to the FP8 grid
floate4m3_t dither(float v, uint n)
{
    const float16_t h = float16_t(isinf(v) || isnan(v) ? 0.0 : min(v, 1024.0));
    const uint e = uint(float16BitsToUint16(h)) & 0xFC00u;
    precise float16_t added = h + uint16BitsToFloat16(uint16_t(e | n));
    precise float16_t t = added - uint16BitsToFloat16(uint16_t(e));
    const float16_t r = uint16BitsToFloat16(uint16_t(uint(float16BitsToUint16(t)) & 0xFF80u));
    return floate4m3_t(clamp(r, float16_t(-448.0), float16_t(448.0)));
}
void gatherPixel(uint pi, ivec2 block, float e1, float e2, bool fusedFourth)
{
    const uint px = pi & 7u, py = pi >> 3;
    const ivec2 raw = block + ivec2(px, py), a = abs(raw);
    const ivec2 m = ivec2(a.x < p.size.x ? a.x : 2 * p.size.x - a.x - 2, a.y < p.size.y ? a.y : 2 * p.size.y - a.y - 2);
    const bool outside = raw.x < 0 || raw.y < 0 || a.x >= p.size.x || a.y >= p.size.y;
    // the neighbours with the smallest and the largest depth
    ivec2 lo = m, hi = m;
    float dlo = texelFetch(depthTex, texel(m), 0).r, dhi = dlo;
    const ivec2 offs[8] = ivec2[8](ivec2(-1, -1), ivec2(0, -1), ivec2(1, -1), ivec2(-1, 0), ivec2(1, 0), ivec2(-1, 1), ivec2(0, 1), ivec2(1, 1));
    for (int k = 0; k < 8; ++k)
    {
        const float d = texelFetch(depthTex, texel(m + offs[k]), 0).r;
        if (d < dlo) { dlo = d; lo = m + offs[k]; }
        if (d > dhi) { dhi = d; hi = m + offs[k]; }
    }
    const vec2 mvA = motion(p.depthFlag != 0u ? lo : hi), mvB = motion(p.depthFlag != 0u ? hi : lo), mvC = motion(m);
    if (!outside)
    {
        imageStore(mvAOut, m, vec4(mvA, 0.0, 0.0)); imageStore(mvBOut, m, vec4(mvB, 0.0, 0.0)); imageStore(mvCOut, m, vec4(mvC, 0.0, 0.0));
    }
    const vec3 own = texelFetch(colorTex, texel(m), 0).rgb;
    const vec2 centre = vec2(m) + 0.5, f = centre - p.jitter, inv = 1.0 / vec2(p.size);
    // how much the motion changed since the previous frame
    const vec2 prev = textureLod(prevMvTex, clamp(((f + p.prevJitter) + mvA * p.hi2rnd) * vec2(p.renderInvX, p.renderInvY), vec2(0.0), vec2(p.size) * vec2(p.renderInvX, p.renderInvY)), 0.0).xy;
    const vec2 dm = (p.changeScale * 540.0 / float(p.size.y)) * (mvA * p.hi2rnd - p.hi2rnd * prev);
    const float d2 = fma(dm.x, dm.x, dm.y * dm.y);
    const float change = p.changeGain * (log2(d2 > 0.0 ? sqrt(d2) + 1.0 : 1.0) * 0.6931472);
    // the previous result along the candidates
    precise vec2 q[4];
    q[0] = mvA + f * p.rnd2hi; q[1] = mvC + f * p.rnd2hi; q[2] = mvB + f * p.rnd2hi;
    // The captured PTX fuses the first pixel's fourth coordinate, but reuses a rounded
    // centre*scale product for the second pixel. A single FMA for both changes FP8 dither.
    q[3] = fusedFourth ? fma(centre, p.rnd2hi, mvC) : centre * p.rnd2hi + mvC;
    vec3 h[4];
    for (int k = 0; k < 4; ++k)
    {
        precise vec3 s = own;
        if (!(q[k].x < 0.0 || q[k].y < 0.0 || q[k].x > p.hiSize.x || q[k].y > p.hiSize.y || p.reset != 0u))
#if LDR
            s = textureLod(halfTex, q[k] * vec2(p.halfInvX, p.halfInvY), 0.0).rgb;
#else
        {
            precise vec3 packedHistory = textureLod(halfTex, q[k] * vec2(p.halfInvX, p.halfInvY), 0.0).rgb * 2.45;
            precise vec3 scaled = clamp(packedHistory, 1e-4, 3.0) * (1.0 / 0.36);
            precise vec3 exponent = scaled * L10;
            precise vec3 linear = (exp2(exponent) - 1.0) * (1.0 / 100.0);
            s = linear * (1.0 / e2);
        }
#endif
        h[k] = compress(e1 * s);
    }
    vec4 feat = vec4(0.0);
    float mf = 0.0;
    if (p.reset == 0u)
    {
        mf = change;
        if (!(q[0].x < 0.0 || q[0].y < 0.0 || q[0].x > p.hiSize.x || q[0].y > p.hiSize.y))
        {
            precise vec2 normalized = fma(q[0], 1.0 / p.hiSize, p.prevJitter * inv);
            precise vec2 texels = normalized * vec2(p.size);
            precise vec2 uv = texels * vec2(p.renderInvX, p.renderInvY);
            feat = textureLod(featTex, uv, 0.0);
        }
    }
    const vec3 c = e1 * own;
    const vec3 col = vec3(colour1(c.r), colour1(c.g), colour1(c.b));
    const uint n = uint(noise.n[(p.frame << 8) | uint((a.y & 15) << 4) | uint(a.x & 15)]) >> 9;
    // 32 slots: 16 per plane
    float v[32];
    for (int k = 0; k < 32; ++k) v[k] = 0.0;
    v[0] = col.r; v[1] = col.g; v[2] = feat.x; v[3] = feat.y; v[4] = col.b; v[5] = h[0].r; v[6] = feat.z; v[7] = feat.w;
    v[8] = h[0].g; v[9] = h[0].b; v[12] = h[1].r; v[13] = h[1].g;
    v[16] = h[1].b; v[17] = h[2].r; v[20] = h[2].g; v[21] = h[2].b; v[24] = h[3].r; v[25] = h[3].g; v[28] = h[3].b; v[29] = mf;
    const uint at = ((py >> 1) * 4u + (px >> 1)) * 136u + ((py & 1u) * 2u + (px & 1u)) * 32u;
    for (uint k = 0u; k < 32u; ++k) FE[at + k] = dither(v[k], n);
}
void main()
{
    const uint lane = gl_LocalInvocationID.x;
    const ivec2 block = 8 * ivec2(gl_WorkGroupID.xy) - p.shift;
    const float e1 = p.expo == 0.0 ? 1.0 : p.expo, e2 = p.expoPrev == 0.0 ? 1.0 : p.expoPrev;
    // Static variants preserve the first pixel's FMA; RADV merges a runtime FMA/mul-add selection.
    gatherPixel(lane, block, e1, e2, true);
    gatherPixel(lane + 32u, block, e1, e2, false);
    memoryBarrierShared(); barrier();
    // Four output accumulators live at once: every FE k16 A operand is loaded once and applied to the four
    // weight B operands. Each output column group keeps its own K32 chain: two k16 products, then the previous
    // accumulator added and rounded to f16 (ROUND_HALF_SHARED), exactly as the sequential version did.
    HACC b0; coopMatLoad(b0, w16.w, 0u, 16u, RM);
    HACC b1; coopMatLoad(b1, w16.w, 4u, 16u, RM);
    HACC b2; coopMatLoad(b2, w16.w, 8u, 16u, RM);
    HACC b3; coopMatLoad(b3, w16.w, 12u, 16u, RM);
    ACC a0 = ACC(b0), a1 = ACC(b1), a2 = ACC(b2), a3 = ACC(b3);
    for (uint kt = 0u; kt < 8u; kt += 2u)
    {
        A8 x0, x1;
        coopMatLoad(x0, FE, kt * 16u, 136u, RM);
        coopMatLoad(x1, FE, (kt + 1u) * 16u, 136u, RM);
        ACC p0 = ACC(0.0), p1 = ACC(0.0), p2 = ACC(0.0), p3 = ACC(0.0);
        B8 w0; coopMatLoad(w0, w8.w, kt * 16u * 64u + 0u, 64u, RM); p0 = coopMatMulAdd(x0, w0, p0);
        B8 w1; coopMatLoad(w1, w8.w, kt * 16u * 64u + 16u, 64u, RM); p1 = coopMatMulAdd(x0, w1, p1);
        B8 w2; coopMatLoad(w2, w8.w, kt * 16u * 64u + 32u, 64u, RM); p2 = coopMatMulAdd(x0, w2, p2);
        B8 w3; coopMatLoad(w3, w8.w, kt * 16u * 64u + 48u, 64u, RM); p3 = coopMatMulAdd(x0, w3, p3);
        B8 y0; coopMatLoad(y0, w8.w, (kt + 1u) * 16u * 64u + 0u, 64u, RM); p0 = coopMatMulAdd(x1, y0, p0);
        B8 y1; coopMatLoad(y1, w8.w, (kt + 1u) * 16u * 64u + 16u, 64u, RM); p1 = coopMatMulAdd(x1, y1, p1);
        B8 y2; coopMatLoad(y2, w8.w, (kt + 1u) * 16u * 64u + 32u, 64u, RM); p2 = coopMatMulAdd(x1, y2, p2);
        B8 y3; coopMatLoad(y3, w8.w, (kt + 1u) * 16u * 64u + 48u, 64u, RM); p3 = coopMatMulAdd(x1, y3, p3);
        // The k32 product is independent of C: add the previous half accumulator only after both k16s.
        a0 = p0 + a0; R32(a0);
        a1 = p1 + a1; R32(a1);
        a2 = p2 + a2; R32(a2);
        a3 = p3 + a3; R32(a3);
    }
    // Each lane owns one accumulator column (lane & 15) and eight rows (lane >= 16 ? 8..15 : 0..7). The clamped
    // accumulator is converted to FP8 straight into the packed token, so no staging region or second barrier.
    //
    // A lane's eight rows are the tokens (block >> 1) + (r & 3, r >> 2) in the low half and (r & 3, 2 + (r >> 2))
    // in the high half, so their byte offsets from the lane's first row are affine in r, (r >> 2) * tok.x +
    // (r & 3) tokens. One base word and eight constants therefore address all eight rows, and a block whose
    // whole 4x4 token square is on screen (every block but the image border) needs no per-row bounds test.
#if PLANES
#define TOK_STRIDE 32u
#else
#define TOK_STRIDE 64u
#endif
    const uint c16 = lane & 15u, rbase = (lane >> 4u) << 3u;
    const ivec2 bt = block >> 1;
    const uint b8 = uint((int(bt.y) + int(rbase >> 2)) * p.tok.x + int(bt.x)) * TOK_STRIDE + c16;
    const uint X8 = uint(p.tok.x) * TOK_STRIDE;
#if PLANES
#define DO_STORE(r, off) { \
        const uint tb_ = b8 + (off), pb_ = uint(p.tok.x * p.tok.y) * 32u; \
        outb.a[tb_] = floate4m3_t(clamp(a0[r], -6.2831855, 6.2831855)); \
        outb.a[tb_ + 16u] = floate4m3_t(clamp(a1[r], -6.2831855, 6.2831855)); \
        outb.a[tb_ + pb_] = floate4m3_t(clamp(a2[r], -6.2831855, 6.2831855)); \
        outb.a[tb_ + pb_ + 16u] = floate4m3_t(clamp(a3[r], -6.2831855, 6.2831855)); }
#else
#define DO_STORE(r, off) { \
        const uint tb_ = b8 + (off); \
        outb.a[tb_] = floate4m3_t(clamp(a0[r], -6.2831855, 6.2831855)); \
        outb.a[tb_ + 16u] = floate4m3_t(clamp(a1[r], -6.2831855, 6.2831855)); \
        outb.a[tb_ + 32u] = floate4m3_t(clamp(a2[r], -6.2831855, 6.2831855)); \
        outb.a[tb_ + 48u] = floate4m3_t(clamp(a3[r], -6.2831855, 6.2831855)); }
#endif
#define ROW_OFF(r) ((r) < 4 ? (r) * TOK_STRIDE : X8 + ((r) - 4) * TOK_STRIDE)
    if (bt.x >= 0 && bt.y >= 0 && bt.x + 3 < p.tok.x && bt.y + 3 < p.tok.y)
    {
        DO_STORE(0, ROW_OFF(0)); DO_STORE(1, ROW_OFF(1)); DO_STORE(2, ROW_OFF(2)); DO_STORE(3, ROW_OFF(3));
        DO_STORE(4, ROW_OFF(4)); DO_STORE(5, ROW_OFF(5)); DO_STORE(6, ROW_OFF(6)); DO_STORE(7, ROW_OFF(7));
    }
    else
    {
#define STORE_ROW(r) { \
        const ivec2 t_ = bt + ivec2(int(rbase + uint(r)) & 3, int(rbase + uint(r)) >> 2); \
        if (t_.x >= 0 && t_.y >= 0 && t_.x < p.tok.x && t_.y < p.tok.y) DO_STORE(r, ROW_OFF(r)) }
        STORE_ROW(0); STORE_ROW(1); STORE_ROW(2); STORE_ROW(3);
        STORE_ROW(4); STORE_ROW(5); STORE_ROW(6); STORE_ROW(7);
#undef STORE_ROW
    }
#undef DO_STORE
#undef ROW_OFF
#undef TOK_STRIDE
}