| {{ env.wgsl.resourceDeclarations }} |
| |
| const HIDDEN: u32 = {{ hiddenSize }}u; |
| const WG: u32 = {{ workgroupSize }}u; |
| const SPLIT: u32 = {{ split }}u; |
| var<workgroup> reduction: array<vec2<f32>, WG>; |
| |
| @compute @workgroup_size(WG, 1, 1) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) { |
| // wg.y carries the row index past the dispatch fold width; wg.z is the part. |
| let row = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u; |
| let part = wg.z; |
| if (row >= params.rows) { return; } |
| let tid = lid.x; |
| let chunk = (HIDDEN + SPLIT - 1u) / SPLIT; |
| let start = part * chunk; |
| let end = min(start + chunk, HIDDEN); |
| let base = row * HIDDEN; |
| let shift = f32(x[base]); |
| var pair = vec2<f32>(0.0); |
| for (var d = start + tid; d < end; d = d + WG) { |
| let value = f32(x[base + d]) - shift; |
| pair = pair + vec2<f32>(value, value * value); |
| } |
| reduction[tid] = pair; |
| workgroupBarrier(); |
| for (var stride = WG >> 1u; stride > 0u; stride = stride >> 1u) { |
| if (tid < stride) { reduction[tid] = reduction[tid] + reduction[tid + stride]; } |
| workgroupBarrier(); |
| } |
| if (tid == 0u) { partials[row * SPLIT + part] = reduction[0]; } |
| } |
| |
