| {{ 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>) { |
| // 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; } |
| 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])); |
| } |
| } |
| |
