File size: 3,472 Bytes
5a12de3 | 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 95 96 | // 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 %}
}
}
|