File size: 5,013 Bytes
9008499 21ba9b4 9008499 21ba9b4 9008499 21ba9b4 9008499 21ba9b4 9008499 21ba9b4 9008499 | 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 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | // Pass 2 of the split-K Conv GEMM: sum the SPLIT_K raw partial slices that
// a Conv GEMM partial pass wrote, then apply the bias and epilogue
// that pass deliberately skipped, and publish the logical [batch, M, N] output.
//
// The partial scratch is [SPLIT_K, batch, PARTIAL_ROWS, PARTIAL_COLS] with both
// trailing axes padded to whole GEMM tiles; this pass indexes only the logical
// (batch, m, n), so the padded rows and columns are read by nobody and their
// contents never matter. One output element per thread, walking n fastest so
// consecutive threads touch consecutive addresses in every slice.
//
// Summing slices in index order changes the f32 association relative to one
// uninterrupted K loop, so the two orders need not be bit-identical.
{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
// 2D-folded flat index: gid.y carries the high bits past the dispatch's
// per-axis workgroup fold width.
let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
if ({{ name }} >= {{ bound }}) {
return;
}{% endmacro %}
{{ env.wgsl.resourceDeclarations }}
{% set applyActivation = hasActivation is defined and hasActivation %}
{% if applyActivation %}
{% if hasActivation %}
// Apply the fused activation in the f32 accumulator before the single output
// cast, avoiding an intermediate convolution tensor.
{% macro fused_act_fn(mode, alpha, beta) %}
fn fused_act(v: f32) -> f32 {
{% if mode == "Relu" %}
return max(v, 0.0);
{% elif mode == "Clip" %}
return clamp(v, f32({{ alpha }}), f32({{ beta }}));
{% elif mode == "LeakyRelu" %}
return select(v * f32({{ alpha }}), v, v >= 0.0);
{% elif mode == "Sigmoid" %}
return 1.0 / (1.0 + exp(-v));
{% elif mode == "Tanh" %}
// tanh is already +/-1 to f32 precision at the clamp bounds. Clamping before
// the builtin preserves that saturated result for larger accumulators.
return tanh(clamp(v, -10.0, 10.0));
{% elif mode == "HardSigmoid" %}
return clamp(f32({{ alpha }}) * v + f32({{ beta }}), 0.0, 1.0);
{% elif mode == "QuickGelu" %}
{% if alpha == 1 %}
// alpha == 1 is SiLU: the multiply folds away (the spelling yolov8-class detectors emit).
return v * (1.0 / (1.0 + exp(-v)));
{% else %}
return v * (1.0 / (1.0 + exp(-(f32({{ alpha }}) * v))));
{% endif %}
{% elif mode == "Elu" %}
return select(f32({{ alpha }}) * (exp(v) - 1.0), v, v >= 0.0);
{% elif mode == "ThresholdedRelu" %}
return select(0.0, v, v > f32({{ alpha }}));
{% elif mode == "Softplus" %}
// log(1 + exp(v)) spelled overflow-safe: max(v, 0) + log(1 + exp(-|v|)).
return max(v, 0.0) + log(1.0 + exp(-abs(v)));
{% elif mode == "Erf" %}
// Abramowitz & Stegun 7.1.26 (|error| < 1.5e-7), the same polynomial the standalone
// Erf/Gelu kernels use, so a fused and an unfused graph agree to f32 rounding.
let a = abs(v);
let t = 1.0 / (1.0 + 0.3275911 * a);
let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
return sign(v) * (1.0 - poly * exp(-a * a));
{% elif mode == "Gelu" and alpha == 0 %}
// ONNX Gelu, erf form: 0.5 * v * (1 + erf(v / sqrt(2))); activation_params[0] == 0.
let u = v * 0.70710678118654752;
let a = abs(u);
let t = 1.0 / (1.0 + 0.3275911 * a);
let poly = ((((1.061405429 * t - 1.453152027) * t + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t;
return 0.5 * v * (1.0 + sign(u) * (1.0 - poly * exp(-a * a)));
{% elif mode == "Gelu" or mode == "FastGelu" %}
// tanh approximation (ONNX Gelu approximate="tanh" arrives as activation_params[0] != 0;
// com.microsoft.FastGelu always). tanh saturates to +/-1 well inside the clamp.
let inner = v * (0.035677408136300125 * v * v + 0.7978845608028654);
return v * (0.5 + 0.5 * tanh(clamp(inner, -10.0, 10.0)));
{% else %}
return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0);
{% endif %}
}{% endmacro %}
{{ fused_act_fn(activation, actAlpha, actBeta) }}
{% endif %}
{% endif %}
const M: u32 = {{ M }}u;
const N: u32 = {{ N }}u;
const SPLIT_K: u32 = {{ splitK }}u;
const PARTIAL_ROWS: u32 = {{ mPadded }}u;
const PARTIAL_COLS: u32 = {{ nPadded }}u;
const PARTIAL_SLICE_STRIDE: u32 = {{ batchCount }}u * PARTIAL_ROWS * PARTIAL_COLS;
const COUNT: u32 = {{ batchCount }}u * M * N;
const WORKGROUP_SIZE: u32 = {{ reduceWorkgroupSize }}u;
@compute @workgroup_size({{ reduceWorkgroupSize }}, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
{{ flat_index_2d("WORKGROUP_SIZE", "idx", "COUNT") }}
let col = idx % N;
let row = (idx / N) % M;
let batch = idx / (M * N);
var src = (batch * PARTIAL_ROWS + row) * PARTIAL_COLS + col;
var acc = 0.0;
for (var z = 0u; z < SPLIT_K; z = z + 1u) {
acc = acc + partial[src];
src = src + PARTIAL_SLICE_STRIDE;
}
{% if hasBias %}
acc = acc + f32(bias[row]);
{% endif %}
{% if applyActivation %}
acc = fused_act(acc);
{% endif %}
y[idx] = {{ T }}(acc);
}
|