{% if usesF16 %}
enable f16;
{% endif %}
{% macro offset_fn(fn_name, opShape, opRank, op_same, op_numel, outShape, outRank, out_numel) %}
fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif %}) -> u32 {
{% if out_numel == 0 %}
  return 0u;
{% elif op_numel == 1 %}
  return 0u;
{% elif op_same %}
  return out_index;
{% else %}
  var offset = 0u;
{% for axis in range(outRank) %}
{% set op_axis = axis - (outRank - opRank) %}
{% if op_axis >= 0 and opShape[op_axis] != 1 %}
{% set c_stride = namespace(value=1) %}
{% for j in range(axis + 1, outRank) %}
{% set c_stride.value = c_stride.value * outShape[j] %}
{% endfor %}
{% set op_stride = namespace(value=1) %}
{% for j in range(op_axis + 1, opRank) %}
{% set op_stride.value = op_stride.value * opShape[j] %}
{% endfor %}
{% if c_stride.value == 1 %}
  let coord{{ axis }} = out_index % {{ outShape[axis] }}u;
{% else %}
  let coord{{ axis }} = (out_index / {{ c_stride.value }}u) % {{ outShape[axis] }}u;
{% endif %}
{% if op_stride.value == 1 %}
  offset = offset + coord{{ axis }};
{% else %}
  offset = offset + coord{{ axis }} * {{ op_stride.value }}u;
{% endif %}
{% endif %}
{% endfor %}
  return offset;
{% endif %}
}
{%- endmacro %}{% macro broadcast_offset_call(fn_name, opShape, outShape, out_index) %}
{% set op_numel = namespace(value=1) %}
{% for d in opShape %}{% set op_numel.value = op_numel.value * d %}{% endfor %}
{% set out_numel = namespace(value=1) %}
{% for d in outShape %}{% set out_numel.value = out_numel.value * d %}{% endfor %}
{{ fn_name }}({% if out_numel.value != 0 and op_numel.value != 1 %}{{ out_index }}{% endif %})
{%- endmacro %}

{{ env.wgsl.resourceDeclarations }}

const HIDDEN: u32 = {{ hiddenSize }}u;
const EPSILON: f32 = {{ epsilon }};
const WG: u32 = {{ workgroupSize }}u;

var<workgroup> partial: array<f32, WG>;
var<workgroup> row_mean: f32;
var<workgroup> row_inv: f32;

{% set xNumel = namespace(value=1) %}
{% for dim in source.xShape %}
{% set xNumel.value = xNumel.value * dim %}
{% endfor %}
{% set scaleNumel = namespace(value=1) %}
{% for dim in source.scaleShape %}
{% set scaleNumel.value = scaleNumel.value * dim %}
{% endfor %}
{% if scaleNumel.value != 1 %}
{{ offset_fn("scale_offset", source.scaleShape, source.scaleShape | length, source.scaleShape == source.xShape, scaleNumel.value, source.xShape, source.xShape | length, xNumel.value) }}
{% endif %}

{% if hasBias %}
{% set biasNumel = namespace(value=1) %}
{% for dim in source.biasShape %}
{% set biasNumel.value = biasNumel.value * dim %}
{% endfor %}
{% if biasNumel.value != 1 %}
{{ offset_fn("bias_offset", source.biasShape, source.biasShape | length, source.biasShape == source.xShape, biasNumel.value, source.xShape, source.xShape | length, xNumel.value) }}
{% endif %}

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

// Reusing partial after this reduction requires a barrier between the read of
// partial[0] and the next write, or the next round can race the prior readers.
{% set trailingBarrier = trailingBarrier is defined and trailingBarrier %}
fn reduce_sum(value: f32, tid: u32) -> f32 {
  partial[tid] = value;
  workgroupBarrier();
{{ wgsl_tree_fold(["partial"], idx="tid", wg="WG", form="head") }}
{% if trailingBarrier %}
  let total = partial[0];
  workgroupBarrier();
  return total;
{% else %}
  return partial[0];
{% endif %}
}


@compute @workgroup_size(WG, 1, 1)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
  let row = wg.x + wg.y * params.rowStride;
  if (row >= params.rows) {
    return;
  }
  let tid = lid.x;
  let base = row * HIDDEN;

  var local_sum = 0.0;
  for (var d = tid; d < HIDDEN; d = d + WG) {
    let value = f32(x[base + d]);
    local_sum = local_sum + value;
  }

  let sum = reduce_sum(local_sum, tid);
  if (tid == 0u) {
    row_mean = sum / f32(HIDDEN);
  }
  workgroupBarrier();

  var local_var_sum = 0.0;
  for (var d = tid; d < HIDDEN; d = d + WG) {
    let diff = f32(x[base + d]) - row_mean;
    local_var_sum = local_var_sum + diff * diff;
  }

  let var_sum = reduce_sum(local_var_sum, tid);
  if (tid == 0u) {
    let variance = var_sum / f32(HIDDEN);
    row_inv = inverseSqrt(variance + EPSILON);
{% if writeMean %}
    mean_out[row] = row_mean;
{% endif %}
{% if writeInvStdDev %}
    inv_std_out[row] = row_inv;
{% endif %}
  }
  workgroupBarrier();

  for (var d = tid; d < HIDDEN; d = d + WG) {
    let index = base + d;
    let normalized = (f32(x[index]) - row_mean) * row_inv;
    var value = normalized * f32(scale[{% if scaleNumel.value == 1 %}0u{% else %}{{ broadcast_offset_call("scale_offset", source.scaleShape, source.xShape, "index") }}{% endif %}]);
{% if hasBias %}
    value = value + f32(bias[{% if biasNumel.value == 1 %}0u{% else %}{{ broadcast_offset_call("bias_offset", source.biasShape, source.xShape, "index") }}{% endif %}]);
{% endif %}
    y[index] = {{ scalar }}(value);
  }
}