| // 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. |
| {{ env.wgsl.resourceDeclarations }} |
| {% set applyActivation = hasActivation is defined and hasActivation %} |
| {% if applyActivation %} |
| {% set hasActivation = hasActivation is defined and hasActivation %} |
| {% set activation = activation | default("") %} |
| {% set actAlpha = actAlpha | default(0.0) %} |
| {% set actBeta = actBeta | default(0.0) %} |
| {% if hasActivation %} |
| // Apply the fused activation in the f32 accumulator before the single output |
| // cast, avoiding an intermediate convolution tensor. |
| {% macro fused_act_return(mode, alpha, beta) -%} |
| {% 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); |
| {% else %} |
| return v * clamp(v * 0.16666666666666666 + 0.5, 0.0, 1.0); |
| {% endif %} |
| {%- endmacro -%} |
| fn fused_act(v: f32) -> f32 { |
| {{ fused_act_return(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>) { |
| // 2D-folded flat index: gid.y carries the high bits past |
| // the per-axis dispatch fold width and reduces to the 1D form at y=0. |
| let idx = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WORKGROUP_SIZE; |
| if (idx >= COUNT) { |
| return; |
| } |
| 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); |
| } |
| |
