WGSL: explicit loop unrolling to replace host-generated subgroup-matrix kernels (@unroll, @no_unroll) #11073

Request

Add loop-unrolling controls last mentioned and discussed in #4110. Subgroup-matrix kernels serve as a concrete motivation and implementation test case.

Implementation remained deferred beyond WGSL V1 despite agreement on @unroll and @no_unroll. Could this workload (amongst many others) justify revisiting that priority? The goal is to expose constant fragment-array indices without generating repeated WGSL statements on the host or losing large chunks of perf. This would allow for much more readable code, less bloat in shader libraries and together with other improvements also allow for compilation optimizations through Vulkan Spec Constants in WebGPU implementations.

Previous discussion

The May 23 minutes raised several objections: implementations could improve (WebGPU, i.e. Tint/Naga) compiler decisions instead of exposing workarounds; there was the mention that the FXC compile-time of the issues author problem that might disappear with newer compilers; and a hint could help one platform while hurting another.

The May 30 minutes record agreement on both attributes, deferring numeric factors. Reservations remained about ignored hints and portable value. A September 13 clarification in #4110 confirmed that implementation was still deferred.

My specific case concerns runtime local-memory traffic and maintaining host-generated matrix kernels, rather than an FXC compile-time workaround, and adovactes for a directive that allows to hint toward unrolling, not prevent it, however the two surely should be added together. Dynamic-array indexing was also raised later in #4110; this adds a concrete subgroup-matrix workload to that discussion, see next section.

Problem and example

Internal Dawn/WGSL GEMM benchmarking at NVIDIA generates accumulator initialization, fragment loads, MMA updates, and stores on the host to obtain constant indices. Constant loop bounds alone have not reliably produced equivalent code: the induction variable remains dynamic unless the compiler expands or specializes its uses. The compilers involved are the GPU Driver, as well as the WebGPU WGSL compilers (tint, naga) acting as a passtrough but adding substantial intermediate variables for the backends to optimize away. Additionally in Dawn DXC is on the chain for the D3D backend, because Dawn does not target DXIL directly. It is not realistic to expect that these compilers to reliably make the correct decision on spilling the opaque matrix types into gloablly addressable memory wihtout any user control.

This reduced excerpt follows the accumulator-update pattern of a optimized, near SoL WGSL GEMM. Fragment declarations, initialization, and loads are omitted; the spelling follows the experimental subgroup-matrix API.

const ROW_TILES: u32 = 2u;
const COL_TILES: u32 = 2u;

// a_frags and b_frags each contain two matrix fragments;
// acc_frags contains four result fragments.
for (var m = 0u; m < ROW_TILES; m++) {
    for (var n = 0u; n < COL_TILES; n++) {
        acc_frags[m * COL_TILES + n] = subgroupMatrixMultiplyAccumulate(
            a_frags[m], b_frags[n], acc_frags[m * COL_TILES + n]);
    }
}

My generator instead emits (and similarly do libraries like ORT WebGPU EP, HF WebGPU kernels):

acc_frags[0u] = subgroupMatrixMultiplyAccumulate(
    a_frags[0u], b_frags[0u], acc_frags[0u]);
acc_frags[1u] = subgroupMatrixMultiplyAccumulate(
    a_frags[0u], b_frags[1u], acc_frags[1u]);
acc_frags[2u] = subgroupMatrixMultiplyAccumulate(
    a_frags[1u], b_frags[0u], acc_frags[2u]);
acc_frags[3u] = subgroupMatrixMultiplyAccumulate(
    a_frags[1u], b_frags[1u], acc_frags[3u]);

Both axes need expansion to expose constant indices. Partial unrolling can leave dynamic accesses and cause global memory spills in the GPU ISA. A remaining dynamic access in initialization or the store epilogue can also prevent register promotion.

This does not mean every matrix operation needs unrolling, but exploration through host side code gen as well as organization and readability of final product are poor. One needs to dump the assembled shader for inspection on a per config basis.

Evidence

My benchmarking shows speedups of 500%+ simply by forcing unrolling and thus const indexing of loops.

The public Vulkan benchmark's shmem.comp uses [[unroll]] on loops over coopmat arrays, with specialization-dependent bounds. Its large outer K loop is unannotated. In the local checkout's supplied fp16-to-fp32 SPIR-V, the annotated loops retain dynamic accesses and OpLoopMerge ... Unroll: expansion is delegated downstream, but the drive can pick it up reliably from the unroll hint and const loop bounds. This underlines the directives are passed through as hints, not guarantees.

Proposal sketch and backend mappings

Revisit mutually exclusive per-loop @unroll and @no_unroll hints, retaining automatic optimization by default. From what I could tell the original objection that WebGPU compilers could own something like this is pretty much incorrect. Tint seems to mainly aim for a losseless passtrough, which totally fine in my opinion. Please correct me if I am wrong. For the example, annotate both tile loops with @unroll to request full expansion when the trip counts are known. Numeric factors and other controls can remain out of scope.

References: HLSL attributes, Clang loop controls, Apple's Metal example using full unrolling, and SPIR-V loop controls.

HLSL attributes are inputs to DXC, not necessarily unchanged driver inputs. Metal behavior needs verification across supported compiler versions. SPIR-V explicitly provides performance hints. Again: these mappings convey intent; they do not ensure identical code generation or performance and that totally suffices for a btter dev experience. The native languagues do not provide anymore. Of course the consideration that using the directives in WGSL could different backends in diferent ways is true, but that goes for pretty much each WGSL source change. For the example of a GEMM that could be: Software pipelining schedule, tile sizes, staging path, etc.

Questions and next step

• Bounds and specialization: Which loop forms are supported, and can requests remain effective after pipeline overrides resolve? This need not change existing restrictions on override-dependent types or arrays.
• Compiler responsibility: Can Tint/Naga or native compiler improvements reliably address this without new author controls? From what I have seen this is not something Tint wants to do and if it did want to, it would mean that devs loose control in what would likely be a very unstable solution. Tint should not own driver like scheduling decisions.
• Testing: CTS can cover valid placement, mutual exclusion, and unchanged behavior. Optimization effectiveness belongs in backend tests and benchmarks.

Vulkan specialization-constant reuse and module caching remain separate Dawn/Tint implementation topics, but could build on added unroll controls in the case of subgroupMatrix GEMMs.

Related: #4110, #4195, and the subgroup-matrix proposal.