{% macro wgsl_tree_fold_stmt(a, op, idx, svar) %}
{% if op == "max" or op == "min" %}
{{ a }}[{{ idx }}] = {{ op }}({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]);{% else %}
{{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] {{ "*" if op == "prod" else "+" }} {{ 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, reuse=false) %}
  var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u;
  loop {
    if ({{ svar }} == 0u) { break; }
    if ({{ idx }} < {{ svar }}) {
{% for a in arrays %}
      {{ wgsl_tree_fold_stmt(a, op, idx, svar) }}
{% endfor %}
    }
    {{ svar }} = {{ svar }} / 2u;
    workgroupBarrier();
  }{% endmacro %}
{% if useSubgroups %}
enable subgroups;
{% endif %}
{{ env.wgsl.resourceDeclarations }}

const HIDDEN: u32 = {{ hidden }}u;
const HIDDEN_V: u32 = {{ hiddenVec }}u;
const WG: u32 = {{ wg }}u;

var<workgroup> sg_partials: array<f32, WG>;

fn reduce_scalar(value: f32{% if useSubgroups %}, sg_lane: u32, sg_id: u32, num_sg: u32{% else %}, tid: u32{% endif %}) -> f32 {
{% if useSubgroups %}
  let s = subgroupAdd(value);
  if (num_sg == 1u) {
    return s;
  }
  if (sg_lane == 0u) {
    sg_partials[sg_id] = s;
  }
  workgroupBarrier();
  var total = 0.0;
  for (var i = 0u; i < num_sg; i = i + 1u) {
    total = total + sg_partials[i];
  }
  return total;
{% else %}
  // No-subgroup tier: workgroup barrier tree-reduction (WG is a power of two).
  sg_partials[tid] = value;
  workgroupBarrier();
{{ wgsl_tree_fold(["sg_partials"], idx="tid", wg="WG", form="head", breakInline=true) }}
  return sg_partials[0];
{% endif %}
}

// 4 contiguous residual elements (input[idx] + skip[skip_idx] [+ bias]) at vec4
// index `vi`. skip_idx == idx for the normal (non-broadcast) path; for a skip
// that broadcasts across the leading/batch dim uses a folded index.
fn residual_value(idx: u32, skip_idx: u32{% if hasBias %}, vi: u32{% endif %}) -> vec4<f32> {
  var value = vec4<f32>(input[idx]) + vec4<f32>(skip[skip_idx]);
{% if hasBias %}
  value = value + vec4<f32>(bias[vi]);
{% endif %}
  return value;
}

@compute @workgroup_size(WG, 1, 1)
fn main(
  @builtin(workgroup_id) wg_id: vec3<u32>,
  @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
  @builtin(subgroup_invocation_id) sg_lane: u32,
  @builtin(subgroup_id) sg_id: u32,
  @builtin(num_subgroups) num_sg: u32{% endif %}
) {
  let row = wg_id.x + wg_id.y * params.rowStride;
  if (row >= params.rows) {
    return;
  }
  let tid = lid.x;
  let base = row * HIDDEN_V;
  let skip_base = base;

  var acc = 0.0;
  for (var i = tid; i < HIDDEN_V; i = i + WG) {
    let v = residual_value(base + i, skip_base + i{% if hasBias %}, i{% endif %});
    acc = acc + dot(v, v);
  }

  let total = reduce_scalar(acc{% if useSubgroups %}, sg_lane, sg_id, num_sg{% else %}, tid{% endif %});
  let row_inv = inverseSqrt(total / f32(HIDDEN) + params.epsilon);

  for (var i = tid; i < HIDDEN_V; i = i + WG) {
    let idx = base + i;
    let residual = residual_value(idx, skip_base + i{% if hasBias %}, i{% endif %});
{% if writeResidualSum %}
    input_skip_bias_sum[idx] = {{ vecType }}(residual);
{% endif %}
    output[idx] = {{ vecType }}(residual * row_inv * vec4<f32>(gamma[i]));
  }
}