File size: 2,277 Bytes
48220c7 9088f09 48220c7 9088f09 48220c7 | 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 | {% macro wgsl_tree_fold_stmt(a, op, idx, svar) %}
{% if op == "max" or op == "min" %}
{{ a }}[{{ idx }}] = {{ op }}({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]);{% else %}
{{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] {{ "*" if op == "prod" else "+" }} {{ a }}[{{ idx }} + {{ svar }}];{% endif %}{% endmacro %}
{% macro wgsl_tree_fold(arrays, op="add", idx="lid", wg="WORKGROUP_SIZE", svar="stride", typed=false, form="tail", breakInline=false, bodyInline=false, barrierFirst=false, reuse=false) %}
var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u;
loop {
if ({{ svar }} == 0u) { break; }
if ({{ idx }} < {{ svar }}) { {{ wgsl_tree_fold_stmt(arrays[0], op, idx, svar) }} }
{{ svar }} = {{ svar }} / 2u;
workgroupBarrier();
}{% endmacro %}
/* Split-K partial sum-of-squares for tensors with few rows and a large hidden
dimension. A workgroup-per-row kernel exposes too little parallelism in this
regime, so this pass splits each row across SPLIT workgroups
(row = wg.x, split index = wg.z). Each workgroup accumulates a partial
sum-of-squares over its HIDDEN/SPLIT slice and writes one partial to scratch.
The normalize pass folds the SPLIT partials per row. Split-K reassociates the
f32 sum, so this route is not bit-identical to the unsplit reduction. */
{{ env.wgsl.resourceDeclarations }}
const HIDDEN: u32 = {{ hiddenSize }}u;
const WG: u32 = {{ workgroupSize }}u;
const SPLIT: u32 = {{ split }}u;
var<workgroup> red: array<f32, WG>;
@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 + wg.y * params.rowStride;
if (row >= params.rows) {
return;
}
let k = wg.z;
let tid = lid.x;
let chunk = (HIDDEN + SPLIT - 1u) / SPLIT;
let start = k * chunk;
var end = start + chunk;
if (end > HIDDEN) { end = HIDDEN; }
let base = row * HIDDEN;
var acc = 0.0;
var d = start + tid;
loop {
if (d >= end) { break; }
let v = f32(x[base + d]);
acc = acc + v * v;
d = d + WG;
}
red[tid] = acc;
workgroupBarrier();
{{ wgsl_tree_fold(["red"], idx="tid", wg="WG", typed=true, form="head", breakInline=true, bodyInline=true) }}
if (tid == 0u) {
partials[row * SPLIT + k] = red[0];
}
}
|