| {% 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 |
| // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements). |
| {% elif note == "limit" %} |
| // 2D-folded flat index: gid.y carries the high bits past the |
| // maxComputeWorkgroupsPerDimension limit. |
| {% elif note == "device-axis" %} |
| // The flat dispatch is folded across x/y at the device's per-axis workgroup |
| // limit; 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 |
| // maxComputeWorkgroupsPerDimension limit (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 |
| // maxComputeWorkgroupsPerDimension limit. |
| {% elif note == "dispatch" %} |
| // 2D-folded flat index: gid.y carries the high bits past the |
| // maxComputeWorkgroupsPerDimension dispatch limit. |
| {% endif %} |
| {% if bound == "" %} |
| let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u; |
| {%- elif guardInline %} |
| let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u; |
| if ({{ name }} >= {{ bound }}) { return; } |
| {%- else %} |
| let {{ name }} = gid.x + gid.y * nwg.x * {{ 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>, @builtin(num_workgroups) nwg: 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)); |
| y[i] = x[i] * alpha + (bias[channel4] - input_mean[channel4] * alpha); |
| } |
| |
