File size: 1,842 Bytes
0af9165 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | // 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 %}
}
|