#version 450
// Q4_K int8-dp4a decode matvec with VECTORIZED uvec2 weight loads. Same math + thread layout as
// mul_mat_q4k_dp4a (column dp4a, one output row/invocation, the shipped 1.8x default) but the weight
// SSBO is typed uvec2 and the 8 qs u32 of each 64-weight chunk are read as FOUR uvec2 (8-byte)
// transactions instead of eight scalar uint -- wider memory transactions for better bus efficiency on
// the bandwidth-access-bound RADV path (occupancy/ILP were both flat; transaction size is the untried
// lever). Alignment-safe: a Q4_K super-block is 36 u32 (even), super-blocks start even, the qs region
// starts at base+4 u32 (even) -> every uvec2 load is 8-byte aligned. Scalar d/dmin/scale reads extract
// the component via wu(). Decode == mul_mat_q4k.
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
#extension GL_EXT_spirv_intrinsics : require
spirv_instruction(extensions = ["SPV_KHR_integer_dot_product"], capabilities = [6018, 6019], id = 4453)
int sdot_accsat(int a, int b, int acc, spirv_literal int fmt);
layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
layout(set = 0, binding = 0) readonly buffer W { uvec2 w2[]; }; // Q4_K blocks as uvec2 (18 uvec2/block)
layout(set = 0, binding = 1) readonly buffer XQ { uint xq[]; };
layout(set = 0, binding = 2) readonly buffer XS { float xs[]; };
layout(set = 0, binding = 3) readonly buffer XSUM { float xsum[];};
layout(set = 0, binding = 4) writeonly buffer Y { float y[]; };
layout(push_constant) uniform Pc { uint m0; uint mcount; uint nout; uint k; uint woff; };
const uint QK_K = 256u;
const uint BLK_U32 = 36u;
const int FMT4x8 = 0;
// One u32 at u32-index i, extracted from the uvec2 view.
uint wu(uint i) { return w2[i >> 1u][i & 1u]; }
uint scale_byte(uint base, uint b) {
uint word = wu(base + 1u + (b >> 2u));
return (word >> ((b & 3u) * 8u)) & 0xFFu;
}
uint get_scale_min_k4(uint base, uint j) {
uint sc, m;
if (j < 4u) {
sc = scale_byte(base, j) & 63u;
m = scale_byte(base, j + 4u) & 63u;
} else {
uint s_j4 = scale_byte(base, j + 4u);
uint s_jm4 = scale_byte(base, j - 4u);
uint s_j = scale_byte(base, j);
sc = (s_j4 & 0x0Fu) | ((s_jm4 >> 6u) << 4u);
m = (s_j4 >> 4u) | ((s_j >> 6u) << 4u);
}
return (sc << 8u) | m;
}
void main() {
uint n = gl_GlobalInvocationID.x;
if (n >= nout) { return; }
uint nblocks = k / QK_K;
uint rowbase = woff + n * nblocks * BLK_U32;
float acc = 0.0;
for (uint blk = 0u; blk < nblocks; blk++) {
uint base = rowbase + blk * BLK_U32;
uint dword = wu(base);
float d = float(unpackHalf2x16(dword).x);
float dmin = float(unpackHalf2x16(dword).y);
uint qbase = base + 4u; // u32 index of the 128 qs bytes (even -> uvec2-aligned)
uint ablk = blk * 8u;
for (uint c = 0u; c < 4u; c++) {
uint is = c * 2u;
uint sm1 = get_scale_min_k4(base, is);
uint sm2 = get_scale_min_k4(base, is + 1u);
float d1 = d * float(sm1 >> 8u);
float m1 = dmin * float(sm1 & 0xFFu);
float d2 = d * float(sm2 >> 8u);
float m2 = dmin * float(sm2 & 0xFFu);
uint qcw2 = (qbase + c * 8u) >> 1u; // uvec2 index of this chunk's 8 qs u32 = 4 uvec2
uint blo = ablk + is;
uint xqlo = blo * 8u;
uint xqhi = (blo + 1u) * 8u;
int dlo = 0;
int dhi = 0;
for (uint p = 0u; p < 4u; p++) { // 4 uvec2 loads cover the 8 qs u32
uvec2 qv = w2[qcw2 + p];
uint j = p * 2u;
// low byte-lane of qv.x / qv.y -> two consecutive qs u32
uint w0 = qv.x;
uint w1 = qv.y;
dlo = sdot_accsat(int(w0 & 0x0F0F0F0Fu), int(xq[xqlo + j]), dlo, FMT4x8);
dhi = sdot_accsat(int((w0 >> 4u) & 0x0F0F0F0Fu), int(xq[xqhi + j]), dhi, FMT4x8);
dlo = sdot_accsat(int(w1 & 0x0F0F0F0Fu), int(xq[xqlo + j + 1u]), dlo, FMT4x8);
dhi = sdot_accsat(int((w1 >> 4u) & 0x0F0F0F0Fu), int(xq[xqhi + j + 1u]), dhi, FMT4x8);
}
acc += d1 * xs[blo] * float(dlo) - m1 * xsum[blo];
acc += d2 * xs[blo + 1u] * float(dhi) - m2 * xsum[blo + 1u];
}
}
y[n] = acc;
}