{% macro wgsl_tree_fold_stmt(a, op, idx, svar) %}
{% if op == "max" %}
{{ a }}[{{ idx }}] = max({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]);
{%- else %}
{{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] + {{ a }}[{{ idx }} + {{ svar }}];
{%- endif %}
{% endmacro %}
{% macro wgsl_tree_fold(arrays, op="add", idx="lid", wg="WORKGROUP_SIZE", svar="stride", typed=false, form="tail", breakInline=false, bodyInline=false, barrierFirst=false) %}
  var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u;
  loop {
{% if form == "head" %}
{% if breakInline %}
    if ({{ svar }} == 0u) { break; }
{% else %}
    if ({{ svar }} == 0u) {
      break;
    }
{% endif %}
{% endif %}
{% if bodyInline %}
    if ({{ idx }} < {{ svar }}) { {{ wgsl_tree_fold_stmt(arrays[0], op, idx, svar) }} }
{% else %}
    if ({{ idx }} < {{ svar }}) {
{% for a in arrays %}
      {{ wgsl_tree_fold_stmt(a, op, idx, svar) }}
{% endfor %}
    }
{% endif %}
{% if form == "head" %}
{% if barrierFirst %}
    workgroupBarrier();
    {{ svar }} = {{ svar }} / 2u;
{% else %}
    {{ svar }} = {{ svar }} / 2u;
    workgroupBarrier();
{% endif %}
{% else %}
    workgroupBarrier();
    if ({{ svar }} == 1u) {
      break;
    }
    {{ svar }} = {{ svar }} / 2u;
{% endif %}
  }
{%- endmacro %}

/* Split-K partial sums for tensors with few planes and a large spatial extent.
   A workgroup-per-plane kernel exposes too little parallelism, so this pass
   splits each plane across SPLIT workgroups. Each accumulates a raw sum and
   sum-of-squares over its slice. The combine pass produces mean and inverse
   standard deviation, and the apply pass normalizes. */
{% set vectorized = vectorized if vectorized is defined else false %}
{% set useSubgroups = useSubgroups if useSubgroups is defined else false %}
{% if usesF16 %}
enable f16;
{% endif %}
{% if useSubgroups %}
enable subgroups;
{% endif %}
{% set LOAD_OPEN = "vec4<f32>(" if usesF16 else "" %}
{% set LOAD_CLOSE = ")" if usesF16 else "" %}
{{ env.wgsl.resourceDeclarations }}

const WG: u32 = {{ workgroupSize }}u;
const SPLIT: u32 = {{ split }}u;

{% if useSubgroups %}
// One slot per possible subgroup avoids assuming any mapping from local
// invocation IDs to subgroup membership.
var<workgroup> subgroup_partials: array<vec2<f32>, WG>;
{% else %}
var<workgroup> red_sum: array<f32, WG>;
var<workgroup> red_sq: array<f32, WG>;
{% endif %}

@compute @workgroup_size(WG, 1, 1)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>,
        @builtin(num_workgroups) nwg: vec3<u32>{% if useSubgroups %},
        @builtin(subgroup_invocation_id) subgroup_lane: u32,
        @builtin(subgroup_id) subgroup_id: u32,
        @builtin(num_subgroups) num_subgroups: u32{% endif %}) {
  let plane = wg.x + wg.y * nwg.x;
  if (plane >= params.planes) {
    return;
  }
  let k = wg.z;
  let tid = lid.x;
{% if vectorized %}
  let spatial = params.spatial / 4u;
{% else %}
  let spatial = params.spatial;
{% endif %}
  let chunk = (spatial + SPLIT - 1u) / SPLIT;
  let start = k * chunk;
  var end = start + chunk;
  if (end > spatial) { end = spatial; }
  let base = plane * spatial;

  var s = 0.0;
  var sq = 0.0;
  var i = start + tid;
  loop {
    if (i >= end) { break; }
{% if vectorized %}
    let v = {{ LOAD_OPEN }}input[base + i]{{ LOAD_CLOSE }};
    s = s + v.x + v.y + v.z + v.w;
    sq = sq + dot(v, v);
{% else %}
    let v = f32(input[base + i]);
    s = s + v;
    sq = sq + v * v;
{% endif %}
    i = i + WG;
  }
{% if useSubgroups %}
  let subgroup_total = vec2<f32>(subgroupAdd(s), subgroupAdd(sq));
  if (subgroup_lane == 0u) {
    subgroup_partials[subgroup_id] = subgroup_total;
  }
  workgroupBarrier();
  if (tid == 0u) {
    var total = vec2<f32>(0.0);
    for (var subgroup = 0u; subgroup < num_subgroups; subgroup = subgroup + 1u) {
      total = total + subgroup_partials[subgroup];
    }
    let idx = (plane * SPLIT + k) * 2u;
    partials[idx] = total.x;
    partials[idx + 1u] = total.y;
  }
{% else %}
  red_sum[tid] = s;
  red_sq[tid] = sq;
  workgroupBarrier();
{{ wgsl_tree_fold(["red_sum", "red_sq"], idx="tid", wg="WG", typed=true, form="head", breakInline=true) }}
  if (tid == 0u) {
    let idx = (plane * SPLIT + k) * 2u;
    partials[idx] = red_sum[0];
    partials[idx + 1u] = red_sq[0];
  }
{% endif %}
}