| {% macro wgsl_tree_fold_stmt(a, op, idx, svar) %} |
| {% if op == "max" %} |
| {{ a }}[{{ idx }}] = max({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]); |
| {%- else %} |
| {{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] + {{ 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) %} |
| var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u; |
| loop { |
| {% if form == "head" %} |
| {% if breakInline %} |
| if ({{ svar }} == 0u) { break; } |
| {% else %} |
| if ({{ svar }} == 0u) { |
| break; |
| } |
| {% endif %} |
| {% endif %} |
| {% if bodyInline %} |
| if ({{ idx }} < {{ svar }}) { {{ wgsl_tree_fold_stmt(arrays[0], op, idx, svar) }} } |
| {% else %} |
| if ({{ idx }} < {{ svar }}) { |
| {% for a in arrays %} |
| {{ wgsl_tree_fold_stmt(a, op, idx, svar) }} |
| {% endfor %} |
| } |
| {% endif %} |
| {% if form == "head" %} |
| {% if barrierFirst %} |
| workgroupBarrier(); |
| {{ svar }} = {{ svar }} / 2u; |
| {% else %} |
| {{ svar }} = {{ svar }} / 2u; |
| workgroupBarrier(); |
| {% endif %} |
| {% else %} |
| workgroupBarrier(); |
| if ({{ svar }} == 1u) { |
| break; |
| } |
| {{ svar }} = {{ svar }} / 2u; |
| {% endif %} |
| } |
| {%- 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]; |
| } |
| } |
| |
