// Pass 2 of the split-K Conv GEMM: sum the SPLIT_K raw partial slices that
// a Conv GEMM partial pass wrote, then apply the bias and epilogue
// that pass deliberately skipped, and publish the logical [batch, M, N] output.
//
// The partial scratch is [SPLIT_K, batch, PARTIAL_ROWS, PARTIAL_COLS] with both
// trailing axes padded to whole GEMM tiles; this pass indexes only the logical
// (batch, m, n), so the padded rows and columns are read by nobody and their
// contents never matter. One output element per thread, walking n fastest so
// consecutive threads touch consecutive addresses in every slice.
//
// Summing slices in index order changes the f32 association relative to one
// uninterrupted K loop, so the two orders need not be bit-identical.
{{ env.wgsl.resourceDeclarations }}
{% set applyActivation = hasActivation is defined and hasActivation %}
{% if applyActivation %}
{% set hasActivation = hasActivation is defined and hasActivation %}
{% set activation = activation | default("") %}
{% set actAlpha = actAlpha | default(0.0) %}
{% set actBeta = actBeta | default(0.0) %}
{% if hasActivation %}
// Apply the fused activation in the f32 accumulator before the single output
// cast, avoiding an intermediate convolution tensor.
{% macro fused_act_return(mode, alpha, beta) -%}
{% if mode == "Relu" %}
  return max(v, 0.0);
{% elif mode == "Clip" %}
  return clamp(v, f32({{ alpha }}), f32({{ beta }}));
{% elif mode == "LeakyRelu" %}
  return select(v * f32({{ alpha }}), v, v >= 0.0);
{% elif mode == "Sigmoid" %}
  return 1.0 / (1.0 + exp(-v));
{% elif mode == "Tanh" %}
  // tanh is already +/-1 to f32 precision at the clamp bounds. Clamping before
  // the builtin preserves that saturated result for larger accumulators.
  return tanh(clamp(v, -10.0, 10.0));
{% elif mode == "HardSigmoid" %}
  return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
{% else %}
  return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
{% endif %}
{%- endmacro -%}
fn fused_act(v: f32) -> f32 {
{{ fused_act_return(activation, actAlpha, actBeta) -}}
}
{% endif %}

{% endif %}

const M: u32 = {{ M }}u;
const N: u32 = {{ N }}u;
const SPLIT_K: u32 = {{ splitK }}u;
const PARTIAL_ROWS: u32 = {{ mPadded }}u;
const PARTIAL_COLS: u32 = {{ nPadded }}u;
const PARTIAL_SLICE_STRIDE: u32 = {{ batchCount }}u * PARTIAL_ROWS * PARTIAL_COLS;
const COUNT: u32 = {{ batchCount }}u * M * N;
const WORKGROUP_SIZE: u32 = {{ reduceWorkgroupSize }}u;

@compute @workgroup_size({{ reduceWorkgroupSize }}, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  // 2D-folded flat index: gid.y carries the high bits past
  // the per-axis dispatch fold width and reduces to the 1D form at y=0.
  let idx = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WORKGROUP_SIZE;
  if (idx >= COUNT) {
    return;
  }
  let col = idx % N;
  let row = (idx / N) % M;
  let batch = idx / (M * N);

  var src = (batch * PARTIAL_ROWS + row) * PARTIAL_COLS + col;
  var acc = 0.0;
  for (var z = 0u; z < SPLIT_K; z = z + 1u) {
    acc = acc + partial[src];
    src = src + PARTIAL_SLICE_STRIDE;
  }
{% if hasBias %}
  acc = acc + f32(bias[row]);
{% endif %}
{% if applyActivation %}
  acc = fused_act(acc);
{% endif %}
  y[idx] = {{ T }}(acc);
}