{% 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 %}

/* One workgroup normalizes each row of residual = input + skip, with an
 * optional bias. */
{% set degenerateRow = (not simplified) and hiddenSize == 1 %}
{% if usesF16 %}
enable f16;
{% endif %}
{% if useSubgroups and not degenerateRow %}
enable subgroups;
{% endif %}
{{ env.wgsl.resourceDeclarations }}

{% if not degenerateRow or writeResidualSum %}
const HIDDEN: u32 = {{ hiddenSize }}u;
{% endif %}
const WG: u32 = {{ workgroupSize }}u;
{% if simplified %}

var<workgroup> partial: array<f32, WG>;
{% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true) %}
fn {{ name }}(value: f32, tid: u32) -> f32 {
  {{ buffer }}[tid] = value;
  workgroupBarrier();
  // Ceil-halving keeps every lane when the workgroup size is not a power of
  // two. For even n this matches the power-of-two tree order; for odd n, lanes
  // [0, n-half) fold the upper tail while the middle lane carries forward.
  var n: u32 = {{ wg }};
  loop {
    let half = (n + 1u) / 2u;
    if (tid < n - half) {
{% if mode == "max" %}
      {{ buffer }}[tid] = max({{ buffer }}[tid], {{ buffer }}[tid + half]);
{% else %}
      {{ buffer }}[tid] = {{ buffer }}[tid] + {{ buffer }}[tid + half];
{% endif %}
    }
    workgroupBarrier();
    n = half;
    if (n == 1u) {
      break;
    }
  }
  // The default trailing barrier makes this helper safe for back-to-back calls: every lane reads
  // slot 0 here, so the next call's first store must not run until all lanes have read it.
  // `trailingBarrier=false` is safe only when the buffer is never written again before kernel exit.
  let reduced = {{ buffer }}[0];
{% if trailingBarrier %}
  workgroupBarrier();
{% endif %}
  return reduced;
}
{% endmacro %}

{{ wgsl_tree_reduce_f32("reduce_sum", "add", "partial", "WG") }}
var<workgroup> row_inv: f32;
{% else %}
{% if not degenerateRow %}

var<workgroup> pair_partial: array<vec2<f32>, WG>;

{% if useSubgroups %}
fn reduce_pair(value: vec2<f32>, sg_lane: u32, sg_id: u32, num_sg: u32) -> vec2<f32> {
  let s = vec2<f32>(subgroupAdd(value.x), subgroupAdd(value.y));
  if (num_sg == 1u) {
    return s;
  }
  if (sg_lane == 0u) {
    pair_partial[sg_id] = s;
  }
  workgroupBarrier();
  var total = vec2<f32>(0.0, 0.0);
  for (var i = 0u; i < num_sg; i = i + 1u) {
    total = total + pair_partial[i];
  }
  return total;
}
{% else %}
fn reduce_pair(value: vec2<f32>, tid: u32) -> vec2<f32> {
  pair_partial[tid] = value;
  workgroupBarrier();
{{ wgsl_tree_fold(["pair_partial"], idx="tid", wg="WG", form="head") }}
  return pair_partial[0];
}
{% endif %}
{% endif %}
{% endif %}

{% if not degenerateRow or writeResidualSum %}
fn residual_value(row: u32, d: u32) -> f32 {
  let index = row * HIDDEN + d;
  var value = f32(input[index]) + f32(skip[index]);
{% if hasBias %}
  value = value + f32(bias[d]);
{% endif %}
  return value;
}
{% endif %}

@compute @workgroup_size(WG, 1, 1)
fn main(
  @builtin(workgroup_id) wg: vec3<u32>,
  @builtin(num_workgroups) nwg: vec3<u32>{% if not degenerateRow %},
  @builtin(local_invocation_id) lid: vec3<u32>{% endif %}{% if useSubgroups and not degenerateRow %},
  @builtin(subgroup_invocation_id) sg_lane: u32,
  @builtin(subgroup_id) sg_id: u32,
  @builtin(num_subgroups) num_sg: u32{% endif %}
) {
  // 2D-folded row index: wg.y carries the high bits past the maxComputeWorkgroupsPerDimension
  // workgroup-per-dimension dispatch limit. Reduces to wg.x when nwg.y == 1;
  // the row >= params.rows guard drops the over-dispatched tail.
  let row = wg.x + wg.y * nwg.x;
  if (row >= params.rows) {
    return;
  }
{% if not degenerateRow %}
  let tid = lid.x;
{% endif %}
{% if simplified %}

  // RMS normalization uses one sum-of-squares sweep, without a mean or beta.

  var local_sq = 0.0;
  for (var d: u32 = tid; d < HIDDEN; d = d + WG) {
    let value = residual_value(row, d);
    local_sq = local_sq + value * value;
  }
  let sq = reduce_sum(local_sq, tid);
  if (tid == 0u) {
    row_inv = inverseSqrt(sq / f32(HIDDEN) + params.epsilon);
  }
  workgroupBarrier();

  for (var d: u32 = tid; d < HIDDEN; d = d + WG) {
    let index = row * HIDDEN + d;
    let residual = residual_value(row, d);
{% if writeResidualSum %}
    input_skip_bias_sum[index] = {{ scalar }}(residual);
{% endif %}
    output[index] = {{ scalar }}(residual * row_inv * f32(gamma[d]));
  }
{% elif degenerateRow %}

  // HIDDEN == 1: the row's mean is its only element, so the centered value and
  // the variance are exactly zero and the output reduces to beta. The closed
  // form avoids computing that zero by subtracting two equal rounded values.
  let row_inv = inverseSqrt(params.epsilon);
{% if writeResidualSum %}
  let residual = residual_value(row, 0u);
  input_skip_bias_sum[row] = {{ scalar }}(residual);
{% endif %}
  // 0.0 * row_inv keeps the IEEE result when epsilon == 0 makes row_inv +Inf.
  output[row] = {{ scalar }}(0.0 * row_inv * f32(gamma[0]){% if hasBeta %} + f32(beta[0]){% endif %});
{% else %}

  // Shifted moments: accumulating (x - x[0], (x - x[0])^2) keeps the sums
  // small for rows with a large common offset; every thread reconstructs the
  // row mean and variance from the merged pair.
  let shift = residual_value(row, 0u);
  var acc = vec2<f32>(0.0, 0.0);
  for (var d = tid; d < HIDDEN; d = d + WG) {
    let centered = residual_value(row, d) - shift;
    acc.x = acc.x + centered;
    acc.y = acc.y + centered * centered;
  }

{% if useSubgroups %}
  let totals = reduce_pair(acc, sg_lane, sg_id, num_sg);
{% else %}
  let totals = reduce_pair(acc, tid);
{% endif %}
  let mean_delta = totals.x / f32(HIDDEN);
  let row_mean = shift + mean_delta;
  let variance = max(totals.y / f32(HIDDEN) - mean_delta * mean_delta, 0.0);
  let row_inv = inverseSqrt(variance + params.epsilon);
  for (var d = tid; d < HIDDEN; d = d + WG) {
    let index = row * HIDDEN + d;
    let residual = residual_value(row, d);
{% if writeResidualSum %}
    input_skip_bias_sum[index] = {{ scalar }}(residual);
{% endif %}
    output[index] = {{ scalar }}((residual - row_mean) * row_inv * f32(gamma[d]){% if hasBeta %} + f32(beta[d]){% endif %});
  }
{% endif %}
}