// Native part of engine/rr/enc0_m.comp (generated by engine/rr/gen_enc0_m.py): declarations, after the translated
// header. Binding 1 holds NGX's weight image followed by engine/rr/enc0_pack.py's blob (offsets in enc0_m_w.h already
// include its base); binding 2 is the arena.
#extension GL_EXT_shader_8bit_storage : require
#extension GL_KHR_shader_subgroup_vote : require
#extension GL_EXT_control_flow_attributes : require
#extension GL_GOOGLE_include_directive : require
#include "enc0_m_w.h"
#include "rr_exact.glsl"
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u;
layout(set = 0, binding = 1, std430) readonly buffer W8U { uint8_t w[]; } w8u;
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC 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
#define CM gl_CooperativeMatrixLayoutColumnMajor
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#define Q8(dst, src) Q8PAIRS(dst, src)
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
shared uint KVS[2 * 768];        // the current head's K, then V^T; before the attention the key pixels' features
#define XF KVS                   // the key pixels' 32 FP8 features (plane 0 | plane 1), slot-major
shared uint XT[96 * 8];          // the normalised key tokens
#define RR_KS KVS
#define RR_KS0 0u
#define RR_VS KVS
#define RR_VS0 768u
#define MG XT                    // the merge rows (16 x 128 FP8) reuse XT
shared float16_t O112[64 * 4];   // surface 112 per token, for surface 128's 2x2 sums
ivec2 keyXY(uint t, int x0, int y0)
{
    if (t < 64u) return ivec2(x0 + int(t & 7u), y0 + int(t >> 3u));
    const int i = int(t & 7u), side = int((t - 64u) >> 3u);
    return ivec2(side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8), side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i));
}
uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d)
{
    return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u);
}
uvec2 mergeXY(uint r) { return r < 8u ? uvec2(r >> 1u, r & 1u) : uvec2((r - 8u) >> 1u, 2u + (r & 1u)); }
// the translated prologue's shared buffer sh: input_noise (10 x 10 pixels, 4 f16) at byte 0, input_tensor (two planes
// of 10 x 10 pixels x 16 FP8 features) at byte 800
float16_t noiseAt(uint gx, uint gy, uint c)
{
    const uint b = 80u * gy + 8u * gx + 2u * c;
    return hb((sh[b >> 2] >> (8u * (b & 2u))) & 0xffffu);
}