| /* Normalize residual = input + skip, with an optional bias. Reductions use |
| * one workgroup per row; closed-form one-element rows use one invocation. */ |
| {% set degenerateRow = false %} |
| {{ env.wgsl.resourceDeclarations }} |
| |
| const HIDDEN: u32 = {{ hiddenSize }}u; |
| const WG: u32 = {{ workgroupSize }}u; |
| |
| var<workgroup> partial: array<f32, WG>; |
| {% macro wgsl_tree_reduce_f32(name, mode, buffer="partial", wg="WG", trailingBarrier=true, valueType="f32") %} |
| fn {{ name }}(value: {{ valueType }}, tid: u32) -> {{ valueType }} { |
| {{ buffer }}[tid] = value; |
| workgroupBarrier(); |
| // Ceil-halving keeps every lane when the workgroup size is not a power of |
| // two. For even n this matches the power-of-two tree order; for odd n, lanes |
| // [0, n-half) fold the upper tail while the middle lane carries forward. |
| var n: u32 = {{ wg }}; |
| loop { |
| let half = (n + 1u) / 2u; |
| if (tid < n - half) { |
| {{ buffer }}[tid] = {{ buffer }}[tid] + {{ buffer }}[tid + half]; |
| } |
| workgroupBarrier(); |
| n = half; |
| if (n == 1u) { |
| break; |
| } |
| } |
| // The default trailing barrier makes this helper safe for back-to-back calls: every lane reads |
| // slot 0 here, so the next call's first store must not run until all lanes have read it. |
| // `trailingBarrier=false` is safe only when the buffer is never written again before kernel exit. |
| let reduced = {{ buffer }}[0]; |
| workgroupBarrier(); |
| return reduced; |
| } |
| {% endmacro %} |
| {{ wgsl_tree_reduce_f32("reduce_sum", "add", "partial", "WG") }} |
| var<workgroup> row_inv: f32; |
| |
| fn residual_value(row: u32, d: u32) -> f32 { |
| let index = row * HIDDEN + d; |
| var value = f32(input[index]) + f32(skip[index]); |
| {% if hasBias %} |
| value = value + f32(bias[d]); |
| {% endif %} |
| return value; |
| } |
| |
| @compute @workgroup_size(WG, 1, 1) |
| fn main( |
| @builtin({{ "global_invocation_id" if degenerateRow else "workgroup_id" }}) {{ "gid" if degenerateRow else "wg" }}: vec3<u32>, |
| @builtin(local_invocation_id) lid: vec3<u32>) { |
| // Fold the row grid across workgroups; independent rows also include the |
| // invocation offset. The bounds guard drops the final dispatch tail. |
| let row = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u; |
| if (row >= params.rows) { |
| return; |
| } |
| let tid = lid.x; |
| |
| // RMS normalization uses one sum-of-squares sweep, without a mean or beta. |
| |
| var local_sq = 0.0; |
| for (var d: u32 = tid; d < HIDDEN; d = d + WG) { |
| let value = residual_value(row, d); |
| local_sq = local_sq + value * value; |
| } |
| let sq = reduce_sum(local_sq, tid); |
| if (tid == 0u) { |
| row_inv = inverseSqrt(sq / f32(HIDDEN) + params.epsilon); |
| {% if writeMean is defined and writeMean %} |
| mean[row] = 0.0; |
| {% endif %} |
| {% if writeInvStd is defined and writeInvStd and not (packedStatistics is defined and packedStatistics) %} |
| inv_std_var[row] = row_inv; |
| {% endif %} |
| } |
| workgroupBarrier(); |
| |
| for (var d: u32 = tid; d < HIDDEN; d = d + WG) { |
| let index = row * HIDDEN + d; |
| let residual = residual_value(row, d); |
| {% if writeResidualSum %} |
| input_skip_bias_sum[index] = {{ scalar }}(residual); |
| {% endif %} |
| output[index] = {{ scalar }}(residual * row_inv * f32(gamma[d])); |
| } |
| } |
| |
