| // 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); |
| } |
| } |
| |
