103 lines
3.2 KiB
Plaintext
103 lines
3.2 KiB
Plaintext
|
#version 450
|
||
|
|
|
||
|
|
/*
|
||
|
|
* Iris Vulkan GEMM (Q8_0 weights)
|
||
|
|
*
|
||
|
|
* Computes C[M,N] = alpha * op(A)[M,K] @ op(B)[K,N] + beta * C[M,N]
|
||
|
|
* where A and C are f32 and B is a GGML Q8_0-quantized weight matrix. Used for
|
||
|
|
* weight-bound linear layers: A is the f32 activation, B is the Q8_0 weight.
|
||
|
|
*
|
||
|
|
* GGML Q8_0 packs the weight into blocks of 32 elements. Each block is 34
|
||
|
|
* bytes laid out as: an fp16 scale (2 bytes) followed by 32 signed int8
|
||
|
|
* quants. A weight element at flat index e dequantizes to
|
||
|
|
* value = float(qs[e % 32]) * fp16_to_f32(scale[e / 32])
|
||
|
|
* Blocks never cross a matrix row because every weight's contraction
|
||
|
|
* dimension (K) is a multiple of 32, so the flat index e == row*K + col maps
|
||
|
|
* cleanly onto the block stream.
|
||
|
|
*
|
||
|
|
* The block buffer is read as a uint[] (no 16-bit storage extension needed):
|
||
|
|
* the fp16 scale at the block start is always 16-bit aligned within a word
|
||
|
|
* (block byte offset 34*b is congruent to 0 or 2 mod 4), and the int8 quants
|
||
|
|
* are extracted with byte-granular shifts.
|
||
|
|
*/
|
||
|
|
|
||
|
|
layout(local_size_x = 16, local_size_y = 16) in;
|
||
|
|
|
||
|
|
layout(std430, binding = 0) readonly buffer ABuf { float A[]; };
|
||
|
|
layout(std430, binding = 1) readonly buffer BBuf { uint Bw[]; };
|
||
|
|
layout(std430, binding = 2) buffer CBuf { float C[]; };
|
||
|
|
|
||
|
|
layout(push_constant) uniform PC {
|
||
|
|
uint M, N, K;
|
||
|
|
uint lda, ldb, ldc;
|
||
|
|
uint ta, tb;
|
||
|
|
float alpha, beta;
|
||
|
|
} pc;
|
||
|
|
|
||
|
|
shared float As[16][16];
|
||
|
|
shared float Bs[16][16];
|
||
|
|
|
||
|
|
/* Dequantize the Q8_0 weight element at flat index e. */
|
||
|
|
float load_q8(uint e) {
|
||
|
|
uint block = e >> 5u; /* e / 32 */
|
||
|
|
uint within = e & 31u; /* e % 32 */
|
||
|
|
uint base = block * 34u; /* byte offset of this block */
|
||
|
|
|
||
|
|
/* fp16 scale at byte `base` (16-bit aligned within its word). */
|
||
|
|
uint sword = Bw[base >> 2u];
|
||
|
|
vec2 shalf = unpackHalf2x16(sword);
|
||
|
|
float scale = ((base & 2u) == 0u) ? shalf.x : shalf.y;
|
||
|
|
|
||
|
|
/* int8 quant at byte base + 2 + within. */
|
||
|
|
uint qoff = base + 2u + within;
|
||
|
|
uint qword = Bw[qoff >> 2u];
|
||
|
|
uint b = (qword >> ((qoff & 3u) * 8u)) & 0xffu;
|
||
|
|
int q = (b < 128u) ? int(b) : int(b) - 256; /* sign extend */
|
||
|
|
|
||
|
|
return float(q) * scale;
|
||
|
|
}
|
||
|
|
|
||
|
|
void main() {
|
||
|
|
uint row = gl_GlobalInvocationID.y;
|
||
|
|
uint col = gl_GlobalInvocationID.x;
|
||
|
|
uint tx = gl_LocalInvocationID.x;
|
||
|
|
uint ty = gl_LocalInvocationID.y;
|
||
|
|
|
||
|
|
float acc = 0.0;
|
||
|
|
uint numTiles = (pc.K + 15u) / 16u;
|
||
|
|
|
||
|
|
for (uint t = 0u; t < numTiles; t++) {
|
||
|
|
uint kA = t * 16u + tx;
|
||
|
|
uint kB = t * 16u + ty;
|
||
|
|
|
||
|
|
float av = 0.0;
|
||
|
|
if (row < pc.M && kA < pc.K) {
|
||
|
|
av = (pc.ta == 0u) ? A[row * pc.lda + kA]
|
||
|
|
: A[kA * pc.lda + row];
|
||
|
|
}
|
||
|
|
As[ty][tx] = av;
|
||
|
|
|
||
|
|
float bv = 0.0;
|
||
|
|
if (col < pc.N && kB < pc.K) {
|
||
|
|
uint e = (pc.tb == 0u) ? (kB * pc.ldb + col)
|
||
|
|
: (col * pc.ldb + kB);
|
||
|
|
bv = load_q8(e);
|
||
|
|
}
|
||
|
|
Bs[ty][tx] = bv;
|
||
|
|
|
||
|
|
barrier();
|
||
|
|
|
||
|
|
for (uint k = 0u; k < 16u; k++) {
|
||
|
|
acc += As[ty][k] * Bs[k][tx];
|
||
|
|
}
|
||
|
|
|
||
|
|
barrier();
|
||
|
|
}
|
||
|
|
|
||
|
|
if (row < pc.M && col < pc.N) {
|
||
|
|
uint ci = row * pc.ldc + col;
|
||
|
|
float prev = (pc.beta != 0.0) ? C[ci] : 0.0;
|
||
|
|
C[ci] = pc.alpha * acc + pc.beta * prev;
|
||
|
|
}
|
||
|
|
}
|