File size: 1,176 Bytes
1f79789 b635b3b 1f79789 | 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 26 27 28 29 30 31 32 33 | {{ 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]; }
}
|