File size: 2,953 Bytes
69e74ac f1a8138 69e74ac f1a8138 69e74ac | 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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
/* One workgroup normalizes each row of residual = input + skip, with an
* optional bias. */
{% if usesF16 %}
enable f16;
{% endif %}
{{ 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) %}
fn {{ name }}(value: f32, tid: u32) -> f32 {
{{ 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) {
{% if mode == "max" %}
{{ buffer }}[tid] = max({{ buffer }}[tid], {{ buffer }}[tid + half]);
{% else %}
{{ buffer }}[tid] = {{ buffer }}[tid] + {{ buffer }}[tid + half];
{% endif %}
}
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];
{% if trailingBarrier %}
workgroupBarrier();
{% endif %}
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(workgroup_id) wg: vec3<u32>,
@builtin(local_invocation_id) lid: vec3<u32>) {
// 2D-folded row index: wg.y carries the high bits past the per-axis dispatch fold width.
// Reduces to wg.x when the dispatch does not fold;
// the row >= params.rows guard drops the over-dispatched 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);
}
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]));
}
}
|