| {% 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 }} |
| // Separate optional outputs keep the reduction within eight storage bindings. |
| @compute @workgroup_size({{ workgroupSize }}) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| {{ flat_index_2d(workgroupSize, "row", "params.rows", guardInline=true) }} |
| let stats = row_stats[row]; |
| mean[row] = stats.x; |
| inv_std_var[row] = stats.y; |
| } |
| |
