File size: 1,202 Bytes
27ad418 ab820a8 27ad418 a0e12a2 27ad418 a0e12a2 ab820a8 27ad418 a0e12a2 ab820a8 a0e12a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | {% 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];
}
|