HLSL wave intrinsics on Vulkan, Metal and WebGPU

Use HLSL wave (subgroup) intrinsics such as WavePrefixCountBits, WaveActiveSum and WaveReadLaneFirst in one shader compiled to Vulkan, Metal and WebGPU, and find out which subgroup features the device has to support.

OneShader 1.2.0 compiles the wave and quad intrinsics to all three targets: SPIR-V GroupNonUniform* instructions as DXC chooses them, Metal simd_* and quad_* functions, and WGSL’s subgroups feature. The supported intrinsics, on 32-bit int, uint and float scalars and vectors:

• WaveGetLaneCount, WaveGetLaneIndex, WaveIsFirstLane
• WaveActiveAnyTrue, WaveActiveAllTrue, WaveActiveAllEqual, WaveActiveBallot, WaveActiveCountBits
• WaveActiveSum, WaveActiveProduct, WaveActiveMin, WaveActiveMax, WaveActiveBitAnd, WaveActiveBitOr, WaveActiveBitXor
• WavePrefixSum, WavePrefixProduct, WavePrefixCountBits
• WaveReadLaneAt, WaveReadLaneFirst
• QuadReadAcrossX, QuadReadAcrossY, QuadReadAcrossDiagonal, QuadReadLaneAt

Shaders/Compact.hlsl keeps the values above 0.5, with one atomic per wave instead of one per thread:

StructuredBuffer<float> Values : register(t0);
RWStructuredBuffer<float> Kept : register(u1);
RWByteAddressBuffer Count : register(u2);
 
[shader("compute")]
[numthreads(64, 1, 1)]
void CSMain(uint3 id : SV_DispatchThreadID)
{
    float v = Values[id.x];
    bool keep = v > 0.5;
 
    // One atomic per wave instead of one per thread.
    uint laneOffset = WavePrefixCountBits(keep);
    uint waveCount = WaveActiveCountBits(keep);
    uint waveBase = 0;
    if (WaveIsFirstLane())
        Count.InterlockedAdd(0, waveCount, waveBase);
    waveBase = WaveReadLaneFirst(waveBase);
 
    if (keep)
        Kept[waveBase + laneOffset] = v;
}

Program.cs, with RecursiveAI.OneShader 1.2.0:

using OneShader;
using OneShader.Layouts;
using OneShader.Layouts.WebGpu;
 
CompileResult result = ShaderCompiler.CompileFile("Shaders/Compact.hlsl", new CompileOptions
{
    Targets = [CompileTarget.SpirV, CompileTarget.Msl, CompileTarget.Wgsl],
});
foreach (var d in result.Diagnostics) Console.WriteLine(d);
if (!result.Success) return;
 
foreach (ShaderBlob blob in result.Blobs)
    Console.WriteLine($"{blob.Target}: {blob.Data.Length} bytes");
File.WriteAllBytes("compact.spv", result.Find("CSMain", CompileTarget.SpirV)!.ToArray());
File.WriteAllText("compact.metal", result.Find("CSMain", CompileTarget.Msl)!.AsText());
File.WriteAllText("compact.wgsl", result.Find("CSMain", CompileTarget.Wgsl)!.AsText());
 
// What the device must support: Vulkan subgroup operations, and the WebGPU feature to request.
var entry = result.Reflection!.EntryPoints.Single();
Console.WriteLine($"SubgroupOperations: {entry.SubgroupOperations} (0x{(uint)entry.SubgroupOperations:x})");
 
EntryLayout.TryCreate(result.Reflection!, "CSMain", out EntryLayout? cs, out _);
PipelineLayout pipeline = PipelineLayout.Create(cs!);
if (pipeline.TryToWebGpu(out WebGpuPipelineLayout? webgpu, out var error,
        new WebGpuLayoutOptions { ReadOnlyStorageBuffers = ["Values"] }))
    Console.WriteLine($"WebGPU requiredFeatures: [{string.Join(", ", webgpu!.RequiredFeatures ?? [])}]");
else
    Console.WriteLine(error);

SpirV: 1576 bytes
Msl: 1067 bytes
Wgsl: 939 bytes
SubgroupOperations: Basic, Ballot (0x9)
WebGPU requiredFeatures: [subgroups]

What each target got, from the files the program wrote:

The SPIR-V passes spirv-val for Vulkan 1.2, and declares only the capabilities it uses (GroupNonUniform and GroupNonUniformBallot here).

Checking the device

• Vulkan: SubgroupOperations uses the values of VkSubgroupFeatureFlagBits (0x1 Basic, 0x8 Ballot), so compare it with subgroupSupportedOperations before creating the pipeline. The generated entry class and Layout carry the same value in build-package projects.
• WebGPU: create the device with the features in RequiredFeatures. The WGSL starts with enable subgroups; and diagnostic(off, subgroup_uniformity);, because HLSL doesn’t require wave operations to be in uniform control flow.
• Metal: SIMD-group functions need Metal 2.1 or later, which the MSL 2.4 output already requires.

Limits in 1.2.0

• WaveMatch and WaveMultiPrefix* are errors, with a suggestion:

  match.hlsl(3,61): error OSH0350: The intrinsic 'WaveMatch' is not supported: it needs partitioned subgroup operations, which no target offers portably (SPIR-V has only a vendor extension; MSL and WGSL have none). Build the groups with WaveActiveBallot and WaveReadLaneFirst in a loop instead.

• Vertex shaders on WebGPU. WebGPU’s subgroup functions exist only in compute and fragment shaders, so a vertex shader that uses one compiles for SPIR-V but not for WGSL:

  vsw.hlsl(2,63): error OSH5202: The intrinsic 'WaveActiveSum' is not available in a vertex shader in WGSL: WebGPU's subgroup functions exist only in compute and fragment shaders.

  On Metal, the lane index and count aren’t available in vertex shaders either.

• Checking WGSL with naga. naga 30 doesn’t parse enable subgroups; yet, so it can’t validate this output. The repository checks wave shaders in a browser on WebGPU instead (WebGPU execution parity).

The instruction-level mapping per target is in SPIR-V, MSL and WGSL.

Every intrinsic mapped to SPIR-V, Metal and WGSL: HLSL wave intrinsics mapping table.