| {% 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 }} |
| @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }}) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| {{ flat_index_2d(tunables.WORKGROUP_SIZE, "index") }} |
| let spatial = params.height * params.width; |
| let channel = (index / spatial) % params.channels; |
| {% if usesF16 %} |
| // f16 tensors use an f32 epsilon, so widen the normalization arithmetic to |
| // f32 and narrow only on store. WGSL does not allow the mixed f16 + f32 add. |
| let normalized = (f32(x[index]) - f32(input_mean[channel])) * inverseSqrt(f32(input_var[channel]) + params.epsilon); |
| y[index] = {{ scalar }}(normalized * f32(scale[channel]) + f32(bias[channel])); |
| {% else %} |
| let normalized = (x[index] - input_mean[channel]) * inverseSqrt(input_var[channel] + params.epsilon); |
| y[index] = normalized * scale[channel] + bias[channel]; |
| {% endif %} |
| } |
| |
