// Pass 2 of the split-K column-wise reduction. One invocation per output column
// folds the segment partials and applies the selected reduction's final step.
// Segments are folded in ascending order for deterministic results. This order
// differs from the single-pass reduction but remains within the f32 tolerance.
{% set yv = "f16(" if outputF16 else "" %}
{% set vy = ")" if outputF16 else "" %}
{{ env.wgsl.resourceDeclarations }}
{% macro wgsl_minmax_identity(name, op, scalar="f32") %}
/* Exact {{ op }} reduction identity. WGSL rejects infinity during constant
 * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */
fn {{ name }}() -> {{ scalar }} {
  var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }};
  return bitcast<f32>(bits);
}{% endmacro %}

const WG: u32 = {{ workgroupSize }}u;
const SPLIT: u32 = {{ split }}u;
{% if op == "logsumexp" %}
const F32_MIN: f32 = -3.4028234663852886e38;
const F32_MAX: f32 = 3.4028234663852886e38;

fn is_nan_f32(value: f32) -> bool {
  let bits = bitcast<u32>(value);
  return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u;
}
{% elif op == "max" or op == "min" %}
{{ wgsl_minmax_identity("reduction_identity", op) }}
{% endif %}

@compute @workgroup_size(WG, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>,
        @builtin(num_workgroups) nwg: vec3<u32>) {
  // The start already folds gid.y in, so the stride must span every y row too;
  // an x-only stride would send y = 0 lanes over columns the y >= 1 rows own.
  let stride = nwg.x * nwg.y * WG;
  let start = (gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG) + gid.x;
  for (var col = start; col < params.cols; col = col + stride) {
{% if op == "logsumexp" %}
    // Merge SPLIT (segMax, segSumExp) pairs stably; carry NaN / +Inf markers.
    var nan_value = 0.0;
    var has_nan = false;
    var global_max = F32_MIN;
    for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
      let nv = partials[(2u * SPLIT + seg) * params.cols + col];
      if (nv != 0.0 || is_nan_f32(nv)) {
        has_nan = true;
        nan_value = nv;
      }
      global_max = max(global_max, partials[seg * params.cols + col]);
    }
    var sum = 0.0;
    for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
      let seg_max = partials[seg * params.cols + col];
      let seg_sum = partials[(SPLIT + seg) * params.cols + col];
      sum = sum + seg_sum * exp(seg_max - global_max);
    }
    let has_positive_inf = global_max > F32_MAX;
    let finite_or_inf = select(global_max + log(sum), global_max, has_positive_inf);
    y[col] = {{ yv }}select(finite_or_inf, nan_value, has_nan){{ vy }};
{% else %}
{% if op == "max" %}
    var total = reduction_identity();
{% elif op == "min" %}
    var total = reduction_identity();
{% elif op == "prod" %}
    var total = 1.0;
{% else %}
    var total = 0.0;
{% endif %}
    for (var seg = 0u; seg < SPLIT; seg = seg + 1u) {
      let p = partials[seg * params.cols + col];
{% if op == "max" or op == "min" %}
      total = {{ op }}(total, p);
{% elif op == "prod" %}
      total = total * p;
{% else %}
      total = total + p;
{% endif %}
    }
{% if op == "l2" %}
    y[col] = {{ yv }}sqrt(total){{ vy }};
{% elif op == "logsum" %}
    y[col] = {{ yv }}log(total){{ vy }};
{% elif op == "mean" %}
    y[col] = {{ yv }}total / f32(params.rows){{ vy }};
{% else %}
{% if outputF16 %}
    y[col] = f16(total);
{% else %}
    y[col] = total;
{% endif %}
{% endif %}
{% endif %}
  }
}