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