| {% 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 }} |
| |
| // Rank-2 [N, C] inference vec4 specialization. Vectors run across adjacent channels, so |
| // scale/bias/mean/var are also bound as vec4<f32> and C must be divisible by 4. |
| @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 channel4 = i % params.channels4; |
| let alpha = scale[channel4] * inverseSqrt(input_var[channel4] + vec4<f32>(params.epsilon)); |
| // Subtract the mean before scaling. The expanded form |
| // x * alpha + (bias - mean * alpha) can lose the small residual (x - mean) |
| // when |mean| is much larger than |x - mean|. |
| y[i] = (x[i] - input_mean[channel4]) * alpha + bias[channel4]; |
| } |
| |
