| // 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); |
| } |
| |
