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