| // Pass 2 of the split-K column-wise reduction. One invocation per output column |
| // folds the segment partials and applies the selected reduction's final step. |
| // Segments are folded in ascending order for deterministic results. This order |
| // differs from the single-pass reduction but remains within the f32 tolerance. |
| {% set yv = "f16(" if outputF16 else "" %} |
| {% set vy = ")" if outputF16 else "" %} |
| {{ env.wgsl.resourceDeclarations }} |
| {% macro wgsl_minmax_identity(name, op, scalar="f32") %} |
| /* Exact {{ op }} reduction identity. WGSL rejects infinity during constant |
| * evaluation, so an f32 identity is constructed from its IEEE-754 bits. */ |
| fn {{ name }}() -> {{ scalar }} { |
| var bits = {{ "0xff800000u" if op == "max" else "0x7f800000u" }}; |
| return bitcast<f32>(bits); |
| }{% endmacro %} |
| |
| const WG: u32 = {{ workgroupSize }}u; |
| const SPLIT: u32 = {{ split }}u; |
| {% if op == "logsumexp" %} |
| const F32_MIN: f32 = -3.4028234663852886e38; |
| const F32_MAX: f32 = 3.4028234663852886e38; |
| |
| fn is_nan_f32(value: f32) -> bool { |
| let bits = bitcast<u32>(value); |
| return (bits & 0x7f800000u) == 0x7f800000u && (bits & 0x007fffffu) != 0u; |
| } |
| {% elif op == "max" or op == "min" %} |
| {{ wgsl_minmax_identity("reduction_identity", op) }} |
| {% endif %} |
| |
| @compute @workgroup_size(WG, 1, 1) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>, |
| @builtin(num_workgroups) nwg: vec3<u32>) { |
| // The start already folds gid.y in, so the stride must span every y row too; |
| // an x-only stride would send y = 0 lanes over columns the y >= 1 rows own. |
| let stride = nwg.x * nwg.y * WG; |
| let start = (gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG) + gid.x; |
| for (var col = start; col < params.cols; col = col + stride) { |
| {% if op == "logsumexp" %} |
| // Merge SPLIT (segMax, segSumExp) pairs stably; carry NaN / +Inf markers. |
| var nan_value = 0.0; |
| var has_nan = false; |
| var global_max = F32_MIN; |
| for (var seg = 0u; seg < SPLIT; seg = seg + 1u) { |
| let nv = partials[(2u * SPLIT + seg) * params.cols + col]; |
| if (nv != 0.0 || is_nan_f32(nv)) { |
| has_nan = true; |
| nan_value = nv; |
| } |
| global_max = max(global_max, partials[seg * params.cols + col]); |
| } |
| var sum = 0.0; |
| for (var seg = 0u; seg < SPLIT; seg = seg + 1u) { |
| let seg_max = partials[seg * params.cols + col]; |
| let seg_sum = partials[(SPLIT + seg) * params.cols + col]; |
| sum = sum + seg_sum * exp(seg_max - global_max); |
| } |
| let has_positive_inf = global_max > F32_MAX; |
| let finite_or_inf = select(global_max + log(sum), global_max, has_positive_inf); |
| y[col] = {{ yv }}select(finite_or_inf, nan_value, has_nan){{ vy }}; |
| {% else %} |
| {% if op == "max" %} |
| var total = reduction_identity(); |
| {% elif op == "min" %} |
| var total = reduction_identity(); |
| {% elif op == "prod" %} |
| var total = 1.0; |
| {% else %} |
| var total = 0.0; |
| {% endif %} |
| for (var seg = 0u; seg < SPLIT; seg = seg + 1u) { |
| let p = partials[seg * params.cols + col]; |
| {% if op == "max" or op == "min" %} |
| total = {{ op }}(total, p); |
| {% elif op == "prod" %} |
| total = total * p; |
| {% else %} |
| total = total + p; |
| {% endif %} |
| } |
| {% if op == "l2" %} |
| y[col] = {{ yv }}sqrt(total){{ vy }}; |
| {% elif op == "logsum" %} |
| y[col] = {{ yv }}log(total){{ vy }}; |
| {% elif op == "mean" %} |
| y[col] = {{ yv }}total / f32(params.rows){{ vy }}; |
| {% else %} |
| {% if outputF16 %} |
| y[col] = f16(total); |
| {% else %} |
| y[col] = total; |
| {% endif %} |
| {% endif %} |
| {% endif %} |
| } |
| } |
| |
