// Apply per-plane statistics: y = (x - mean) * inverseStddev * scale + bias.
// The vectorized route packs four adjacent spatial values per invocation; each
// packed load and store remains within one plane.
{% set vectorized = vectorized if vectorized is defined else false %}
{% if usesF16 %}
enable f16;
{% endif %}
{% set LOAD_OPEN = "vec4<f32>(" if usesF16 else "" %}
{% set LOAD_CLOSE = ")" if usesF16 else "" %}
{% set STORE_OPEN = "vec4<f16>(" if usesF16 else "" %}
{% set STORE_CLOSE = ")" if usesF16 else "" %}
{% set CHAN_OPEN = "f32(" if usesF16 else "" %}
{% set CHAN_CLOSE = ")" if usesF16 else "" %}
{{ env.wgsl.resourceDeclarations }}

const WG: u32 = {{ applyWorkgroupSize }}u;

@compute @workgroup_size(WG, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
{% if not vectorized %}
  // 2D-folded flat index: gid.y carries the high bits after dispatch folding.
{% endif %}
  let index = gid.x + gid.y * nwg.x * WG;
  if (index >= params.count) {
    return;
  }
{% if vectorized %}
  // The vectorized path requires each plane to contain a multiple of four
  // values, so a packed load/store never crosses an instance boundary.
  let plane = (index * 4u) / params.spatial;
{% else %}
  let plane = index / params.spatial;
{% endif %}
  let channel = plane % params.channels;
  let mean = stats[plane * 2u];
  let inv_std = stats[plane * 2u + 1u];
{% if vectorized %}
  let value = {{ LOAD_OPEN }}input[index]{{ LOAD_CLOSE }};
  output[index] = {{ STORE_OPEN }}(value - vec4<f32>(mean)) * vec4<f32>(inv_std * {{ CHAN_OPEN }}scale[channel]{{ CHAN_CLOSE }}) + vec4<f32>({{ CHAN_OPEN }}bias[channel]{{ CHAN_CLOSE }}){{ STORE_CLOSE }};
{% else %}
  output[index] = {{ scalar }}((f32(input[index]) - mean) * inv_std * f32(scale[channel]) + f32(bias[channel]));
{% endif %}
}