File size: 1,235 Bytes
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 34 | {{ env.wgsl.resourceDeclarations }}
const HIDDEN: u32 = {{ hiddenSize }}u;
const SPATIAL: u32 = {{ spatial }}u;
const CPG: u32 = {{ channelsPerGroup }}u;
const GROUPS: u32 = {{ numGroups }}u;
const WG: u32 = {{ workgroupSize }}u;
const SPLIT: u32 = {{ split }}u;
const EPSILON: f32 = {{ epsilon }};
@compute @workgroup_size(WG, 1, 1)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
let row = wg.x;
let part = wg.z;
if (row >= params.rows) { return; }
var pair = vec2<f32>(0.0);
for (var p = 0u; p < SPLIT; p = p + 1u) { pair = pair + partials[row * SPLIT + p]; }
let base = row * HIDDEN;
let shift = f32(x[base]);
let mean_d = pair.x / f32(HIDDEN);
let variance = max(pair.y / f32(HIDDEN) - mean_d * mean_d, 0.0);
let mean = shift + mean_d;
let inv_std = inverseSqrt(variance + EPSILON);
let group = row % GROUPS;
let chunk = (HIDDEN + SPLIT - 1u) / SPLIT;
let start = part * chunk;
let end = min(start + chunk, HIDDEN);
for (var d = start + lid.x; d < end; d = d + WG) {
let channel = group * CPG + d / SPATIAL;
let index = base + d;
y[index] = {{ scalar }}((f32(x[index]) - mean) * inv_std * f32(scale[channel]) + f32(bias[channel]));
}
}
|