// 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.
{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
  // per-axis workgroup fold width.
  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
  if ({{ name }} >= {{ bound }}) {
    return;
  }{% endmacro %}
{{ env.wgsl.resourceDeclarations }}
{% set applyActivation = hasActivation is defined and hasActivation %}
{% if applyActivation %}
{% if hasActivation %}
// Apply the fused activation in the f32 accumulator before the single output
// cast, avoiding an intermediate convolution tensor.
{% macro fused_act_fn(mode, alpha, beta) %}
fn fused_act(v: f32) -> f32 {
{% 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);
{% elif mode == "QuickGelu" %}
{% if alpha == 1 %}
  // alpha == 1 is SiLU: the multiply folds away (the spelling yolov8-class detectors emit).
  return v * (1.0 / (1.0 + exp(-v)));
{% else %}
  return v * (1.0 / (1.0 + exp(-(f32({{ alpha }}) * v))));
{% endif %}
{% elif mode == "Elu" %}
  return select(f32({{ alpha }}) * (exp(v) - 1.0), v, v >= 0.0);
{% elif mode == "ThresholdedRelu" %}
  return select(0.0, v, v > f32({{ alpha }}));
{% elif mode == "Softplus" %}
  // log(1 + exp(v)) spelled overflow-safe: max(v, 0) + log(1 + exp(-|v|)).
  return max(v, 0.0) + log(1.0 + exp(-abs(v)));
{% elif mode == "Erf" %}
  // Abramowitz & Stegun 7.1.26 (|error| < 1.5e-7), the same polynomial the standalone
  // Erf/Gelu kernels use, so a fused and an unfused graph agree to f32 rounding.
  let a = abs(v);
  let t = 1.0 / (1.0 + 0.3275911 * a);
  let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
  return sign(v) * (1.0 - poly * exp(-a * a));
{% elif mode == "Gelu" and alpha == 0 %}
  // ONNX Gelu, erf form: 0.5 * v * (1 + erf(v / sqrt(2))); activation_params[0] == 0.
  let u = v * 0.70710678118654752;
  let a = abs(u);
  let t = 1.0 / (1.0 + 0.3275911 * a);
  let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
  return 0.5 * v * (1.0 + sign(u) * (1.0 - poly * exp(-a * a)));
{% elif mode == "Gelu" or mode == "FastGelu" %}
  // tanh approximation (ONNX Gelu approximate="tanh" arrives as activation_params[0] != 0;
  // com.microsoft.FastGelu always). tanh saturates to +/-1 well inside the clamp.
  let inner = v * (0.035677408136300125 * v * v + 0.7978845608028654);
  return v * (0.5 + 0.5 * tanh(clamp(inner, -10.0, 10.0)));
{% else %}
  return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
{% endif %}
}{% endmacro %}
{{ fused_act_fn(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>) {
{{ flat_index_2d("WORKGROUP_SIZE", "idx", "COUNT") }}
  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);
}