| {% 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]); |
| } |
| |
