File size: 2,032 Bytes
5a12de3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | // GEMV specialization for y[N] = a[K] @ B[K, N]. Each workgroup owns 32
// consecutive vec4 column groups, or 128 output columns.
{{ env.wgsl.resourceDeclarations }}
{% set OUT = outputBuffer if outputBuffer is defined else "c" %}
const LANES: u32 = {{ gemvLanes }}u;
// SLICES partitions the K reduction across the workgroup's second dimension.
// Thread zero of each column group combines the slice partials in index order.
const SLICES: u32 = {{ gemvSlices }}u;
var<workgroup> partials: array<vec4<f32>, LANES * SLICES>;
@compute @workgroup_size({{ gemvLanes }}, {{ gemvSlices }}, 1)
fn main(
@builtin(workgroup_id) workgroup_id: vec3<u32>,
@builtin(local_invocation_id) lid: vec3<u32>
) {
let lane = lid.x;
let slice = lid.y;
let cg = workgroup_id.x * LANES + lane;
var acc = vec4<f32>(0.0);
if (cg < params.N4) {
{% if unrollK2 is defined and unrollK2 %}
// Two K positions per iteration amortize loop/address arithmetic on long
// decode projections while preserving each slice's exact strided order.
var k = slice;
for (; k + SLICES < params.K; k = k + 2u * SLICES) {
acc = acc + f32(a[k]) * vec4<f32>(b[k * params.N4 + cg]);
let k1 = k + SLICES;
acc = acc + f32(a[k1]) * vec4<f32>(b[k1 * params.N4 + cg]);
}
if (k < params.K) {
acc = acc + f32(a[k]) * vec4<f32>(b[k * params.N4 + cg]);
}
{% else %}
for (var k = slice; k < params.K; k = k + SLICES) {
acc = acc + f32(a[k]) * vec4<f32>(b[k * params.N4 + cg]);
}
{% endif %}
}
partials[slice * LANES + lane] = acc;
workgroupBarrier();
if (slice == 0u && cg < params.N4) {
var total = partials[lane];
for (var s = 1u; s < SLICES; s = s + 1u) {
total = total + partials[s * LANES + lane];
}
{% if alphaScale is defined and alphaScale != 1 %}
// This specialization bakes the output multiplier into the shader and
// applies it after combining the K slices.
total = total * f32({{ alphaScale }});
{% endif %}
{{ OUT }}[cg] = vec4<{{ T }}>(total);
}
}
|