#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];
// 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
#include "enc0_features.glsl"
void gatherPixel(uint pi, ivec2 block, float e1, float e2, bool fusedFourth)
{
    floate4m3_t fe[32];
    gatherFeatures(pi, block, e1, e2, fusedFourth, fe);
    const uint px = pi & 7u, py = pi >> 3;
    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] = fe[k];
}
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
}