src/vector/knn.wgsl
// Batched int8 kNN over a q8 corpus. The score is the EXACT same integer dot the
// CPU's `dot_i8_scalar` computes: sum over dim of (q_i8 * r_i8), accumulated in
// i32. Integer addition is associative and cannot overflow here (dim <= 65536,
// |product| <= 127*127 = 16129 -> |acc| <= 1.06e9 < 2^31), so ANY accumulation
// order on ANY device yields the SAME i32. Dividing by Q*Q on the host therefore
// reproduces the CPU cosine BIT-EXACTLY. That is the oracle, by construction.
struct Params {
n: u32, // rows in the corpus
dim_words: u32, // dim / 4 (i8 lanes packed 4-per-u32)
q_tile: u32, // queries resident in this dispatch
nblocks: u32, // row blocks per query for the top-k stage
k: u32, // neighbours kept per (query, block)
q_base: u32, // word offset of this tile inside the query buffer
cand_base: u32, // element offset of this tile inside the candidate buffers
_pad2: u32,
};
@group(0) @binding(0) var<uniform> params: Params;
@group(0) @binding(1) var<storage, read> rows: array<u32>; // TRANSPOSED: dim_words * n
@group(0) @binding(2) var<storage, read> queries: array<u32>; // q_tile * dim_words
@group(0) @binding(3) var<storage, read_write> scores: array<i32>; // q_tile * n
@group(0) @binding(4) var<storage, read_write> cand_score: array<i32>; // q_tile*nblocks*k
@group(0) @binding(5) var<storage, read_write> cand_row: array<u32>;
// 16 queries x 192 words (dim 768) = 3072 u32 = 12 KiB. Sized for the worst case
// this rig runs; the host refuses a dispatch that would exceed it.
const MAX_QT: u32 = 16u;
const MAX_DW: u32 = 256u;
var<workgroup> qcache: array<u32, 4096>; // MAX_QT * MAX_DW / 16 -> see host check
// ---- stage 1: score -------------------------------------------------------
// One workgroup covers a strided set of rows; the whole query tile lives in
// workgroup memory, so each row's bytes are read ONCE and reused q_tile times.
// That reuse is the entire reason this is not memory-bound at 1/q_tile of peak.
@compute @workgroup_size(256)
fn score(@builtin(global_invocation_id) gid: vec3<u32>,
@builtin(local_invocation_id) lid: vec3<u32>,
@builtin(num_workgroups) nwg: vec3<u32>) {
let dw = params.dim_words;
let qt = params.q_tile;
let total = qt * dw;
// Cooperative load of the query tile.
var i = lid.x;
loop {
if (i >= total) { break; }
qcache[i] = queries[params.q_base + i];
i = i + 256u;
}
workgroupBarrier();
let stride = nwg.x * 256u;
var row = gid.x;
loop {
if (row >= params.n) { break; }
// Word-transposed corpus: adjacent threads own adjacent ROWS, so at a
// fixed word w their loads are consecutive u32 -- one coalesced burst per
// warp instead of 32 separate 768-byte-strided cache lines. MEASURED at
// n=79k/dim=768/1024q on a 4090: row-major 0.256 s, transposed 0.071 s.
//
// The sixteen accumulators are NAMED SCALARS, not `array<i32, 16>`, and
// the query loop is unrolled by hand. A WGSL array indexed by a loop
// variable lands in scratch memory on every backend that cannot prove the
// index range, which turns each of the 192 words into 16 local-memory
// round trips. Named scalars stay in registers. This is the same trick
// `dot_i8_vnni` plays on the CPU with its eight zmm accumulators, for the
// same reason: keep the partial sums where the ALU can reach them.
//
// q_tile is ALWAYS 16 (host-enforced by QUERY_TILE); tiles shorter than
// that are zero-padded by the host, and the padding lanes' scores are
// dropped when the candidates are merged.
var a0: i32 = 0;
var a1: i32 = 0;
var a2: i32 = 0;
var a3: i32 = 0;
var a4: i32 = 0;
var a5: i32 = 0;
var a6: i32 = 0;
var a7: i32 = 0;
var a8: i32 = 0;
var a9: i32 = 0;
var a10: i32 = 0;
var a11: i32 = 0;
var a12: i32 = 0;
var a13: i32 = 0;
var a14: i32 = 0;
var a15: i32 = 0;
for (var w: u32 = 0u; w < dw; w = w + 1u) {
let rv = rows[w * params.n + row];
a0 = dot4I8Packed(qcache[0u * dw + w], rv) + a0;
a1 = dot4I8Packed(qcache[1u * dw + w], rv) + a1;
a2 = dot4I8Packed(qcache[2u * dw + w], rv) + a2;
a3 = dot4I8Packed(qcache[3u * dw + w], rv) + a3;
a4 = dot4I8Packed(qcache[4u * dw + w], rv) + a4;
a5 = dot4I8Packed(qcache[5u * dw + w], rv) + a5;
a6 = dot4I8Packed(qcache[6u * dw + w], rv) + a6;
a7 = dot4I8Packed(qcache[7u * dw + w], rv) + a7;
a8 = dot4I8Packed(qcache[8u * dw + w], rv) + a8;
a9 = dot4I8Packed(qcache[9u * dw + w], rv) + a9;
a10 = dot4I8Packed(qcache[10u * dw + w], rv) + a10;
a11 = dot4I8Packed(qcache[11u * dw + w], rv) + a11;
a12 = dot4I8Packed(qcache[12u * dw + w], rv) + a12;
a13 = dot4I8Packed(qcache[13u * dw + w], rv) + a13;
a14 = dot4I8Packed(qcache[14u * dw + w], rv) + a14;
a15 = dot4I8Packed(qcache[15u * dw + w], rv) + a15;
}
scores[0u * params.n + row] = a0;
scores[1u * params.n + row] = a1;
scores[2u * params.n + row] = a2;
scores[3u * params.n + row] = a3;
scores[4u * params.n + row] = a4;
scores[5u * params.n + row] = a5;
scores[6u * params.n + row] = a6;
scores[7u * params.n + row] = a7;
scores[8u * params.n + row] = a8;
scores[9u * params.n + row] = a9;
scores[10u * params.n + row] = a10;
scores[11u * params.n + row] = a11;
scores[12u * params.n + row] = a12;
scores[13u * params.n + row] = a13;
scores[14u * params.n + row] = a14;
scores[15u * params.n + row] = a15;
row = row + stride;
}
}
// ---- stage 2: block-local top-k ------------------------------------------
// Each workgroup owns one (query, row-block). Every thread keeps its own top-k
// over a strided slice, then thread 0 merges the 256 thread-local lists. The
// comparison is (score DESC, row-id ASC) -- byte-identical to the CPU `top_k`
// tie-break, which is what makes the merged result the same list, not a similar
// one.
const K_MAX: u32 = 8u;
var<workgroup> ts: array<i32, 2048>; // 256 threads * K_MAX
var<workgroup> tr: array<u32, 2048>;
fn better(sa: i32, ra: u32, sb: i32, rb: u32) -> bool {
if (sa != sb) { return sa > sb; }
return ra < rb;
}
@compute @workgroup_size(256)
fn topk(@builtin(workgroup_id) wid: vec3<u32>,
@builtin(local_invocation_id) lid: vec3<u32>) {
let k = params.k;
let q = wid.x / params.nblocks;
let b = wid.x % params.nblocks;
let per = (params.n + params.nblocks - 1u) / params.nblocks;
let start = b * per;
var end = start + per;
if (end > params.n) { end = params.n; }
var ls: array<i32, 8>;
var lr: array<u32, 8>;
for (var j: u32 = 0u; j < K_MAX; j = j + 1u) {
ls[j] = -2147483648;
lr[j] = 4294967295u;
}
var row = start + lid.x;
loop {
if (row >= end) { break; }
let s = scores[q * params.n + row];
// insertion into the thread-local sorted top-k
if (better(s, row, ls[k - 1u], lr[k - 1u])) {
var j: i32 = i32(k) - 1;
loop {
if (j <= 0) { break; }
if (better(s, row, ls[j - 1], lr[j - 1])) {
ls[j] = ls[j - 1]; lr[j] = lr[j - 1];
j = j - 1;
} else { break; }
}
ls[j] = s; lr[j] = row;
}
row = row + 256u;
}
for (var j: u32 = 0u; j < K_MAX; j = j + 1u) {
ts[lid.x * K_MAX + j] = ls[j];
tr[lid.x * K_MAX + j] = lr[j];
}
workgroupBarrier();
if (lid.x != 0u) { return; }
var gs: array<i32, 8>;
var gr: array<u32, 8>;
for (var j: u32 = 0u; j < K_MAX; j = j + 1u) {
gs[j] = -2147483648;
gr[j] = 4294967295u;
}
for (var c: u32 = 0u; c < 256u * K_MAX; c = c + 1u) {
let s = ts[c];
let r = tr[c];
if (r == 4294967295u) { continue; }
if (better(s, r, gs[k - 1u], gr[k - 1u])) {
var j: i32 = i32(k) - 1;
loop {
if (j <= 0) { break; }
if (better(s, r, gs[j - 1], gr[j - 1])) {
gs[j] = gs[j - 1]; gr[j] = gr[j - 1];
j = j - 1;
} else { break; }
}
gs[j] = s; gr[j] = r;
}
}
let out = params.cand_base + (q * params.nblocks + b) * k;
for (var j: u32 = 0u; j < k; j = j + 1u) {
cand_score[out + j] = gs[j];
cand_row[out + j] = gr[j];
}
}