File size: 1,129 Bytes
27ad418 ab820a8 27ad418 a0e12a2 27ad418 a0e12a2 ab820a8 27ad418 a0e12a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
// 2D-folded flat index: gid.y carries the high bits past the dispatch's
// per-axis workgroup fold width.
let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
if ({{ name }} >= {{ bound }}) {
return;
}{% endmacro %}
{{ env.wgsl.resourceDeclarations }}
// Inference-only vec4 specialization: 128-bit loads/stores over x/y. Gated to
// spatial % 4 == 0 so each vec4 stays inside one channel; the per-component
// calculation is:
// (x - mean) * inverseSqrt(var + epsilon) * scale + bias
@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
{{ flat_index_2d(tunables.WORKGROUP_SIZE, "i", "params.count4") }}
let channel = (i / params.spatial4) % params.channels;
let normalized = (x[i] - vec4<f32>(input_mean[channel])) * inverseSqrt(input_var[channel] + params.epsilon);
y[i] = normalized * vec4<f32>(scale[channel]) + vec4<f32>(bias[channel]);
}
|