File size: 1,296 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 23 24 25 | {% 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 %}
}
|