Kernels:

---

webgpu-kernels

/

ai.onnx.BatchNormalization

Like 0

Follow

WebGPU Kernels Community 24

{% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
{% if note == "dispatch-limit" %}
  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
  // per-axis workgroup fold width (outputs > 16.7M elements).
{% elif note == "limit" %}
  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
  // per-axis workgroup fold width.
{% elif note == "device-axis" %}
  // The flat dispatch is folded across x/y at a fixed per-axis workgroup
  // width; gid.y carries the high portion of the output index.
{% elif note == "vec4-limit" %}
  // 2D-folded flat vec4 index: gid.y carries the high bits past the dispatch's
  // per-axis workgroup fold width (the dispatch caps x and spills into y).
{% elif note == "element-limit" %}
  // 2D-folded flat element index: gid.y carries the high bits past the
  // dispatch's per-axis workgroup fold width.
{% elif note == "dispatch" %}
  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
  // per-axis workgroup fold width.
{% endif %}
{% if bound == "" %}
  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
{%- elif guardInline %}
  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
  if ({{ name }} >= {{ bound }}) { return; }
{%- else %}
  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
  if ({{ name }} >= {{ bound }}) {
    return;
  }
{%- endif %}
{% 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("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];
}