Specialization constants on Vulkan, Metal and WebGPU

// Pipeline-time constants: one compiled shader, many pipelines.
[[vk::constant_id(0)]] const int Radius = 4;
[[vk::constant_id(1)]] const bool Horizontal = true;
[[vk::constant_id(2)]] const float Sigma = 2.0;
 
Texture2D<float4> Source : register(t0);
[[vk::image_format("rgba8")]] RWTexture2D<float4> Target : register(u1);
 
[shader("compute")]
[numthreads(8, 8, 1)]
void CSMain(uint3 id : SV_DispatchThreadID)
{
    int2 step = Horizontal ? int2(1, 0) : int2(0, 1);
    float4 sum = 0;
    float total = 0;
    for (int i = -Radius; i <= Radius; i++)
    {
        float w = exp(-(i * i) / (2 * Sigma * Sigma));
        sum += Source[int2(id.xy) + step * i] * w;
        total += w;
    }
    Target[id.xy] = sum / total;
}

using OneShader.Layouts;
using OneShader.Layouts.Metal;
using OneShader.Layouts.Vulkan;
using Spec.Shaders;
 
// The constants, their ids and defaults, generated from the shader.
Console.WriteLine($"Radius: id {Blur.SpecializationConstants.Radius.Id}, default {Blur.SpecializationConstants.Radius.Default}");
Console.WriteLine($"CSMain reads ids [{string.Join(", ", Blur.CSMain.SpecConstantIds.ToArray())}]");
 
// A vertical blur with radius 8: set fields on the generated struct.
Blur.Specialization s = Blur.Specialization.Defaults;
s.Radius = 8;
s.Horizontal = 0;   // a bool is a 4-byte 0 or 1
 
// Vulkan: the contents of VkSpecializationInfo.
SpecializationInfo vk = s.ToVulkan(Blur.CSMain.SpecConstantIds);
foreach (SpecializationMapEntry e in vk.MapEntries)
    Console.WriteLine($"Vulkan: constantID {e.ConstantId}, offset {e.Offset}, size {e.Size}");
Console.WriteLine($"Vulkan: data {Convert.ToHexString(vk.Data)}");
 
// Metal: MTLFunctionConstantValues entries, by index and MTLDataType.
foreach (FunctionConstantValue v in s.ToMetal(Blur.CSMain.SpecConstantIds))
    Console.WriteLine($"Metal: index {v.Index}, MTLDataType {v.Type.DataType}, bits 0x{v.Bits:x8}");
 
// WebGPU: the compute stage's constants record, keyed by decimal id.
Console.WriteLine($"WebGPU: {Record(s.ToWebGpu(Blur.CSMain.SpecConstantIds))}");
 
// The same by name, from the entry layout (this also works for shaders compiled at run time).
SpecializationValues byName = SpecializationValues.For(Blur.CSMain.Layout).Set("Radius", 8).Set("Horizontal", false);
Console.WriteLine($"By name: {Record(byName.ToWebGpu())}");
 
// What each target compiled the constant to.
Console.WriteLine(Line(Blur.CSMain.Wgsl, "override Radius"));
Console.WriteLine(Line(Blur.CSMain.Msl, "function_constant(0)"));
 
static string Record(IReadOnlyDictionary<string, double> c) =>
    "{ " + string.Join(", ", c.Select(kv => $"\"{kv.Key}\": {kv.Value}")) + " }";
 
static string Line(ReadOnlySpan<byte> utf8, string text) =>
    System.Text.Encoding.UTF8.GetString(utf8).Split('\n').First(l => l.Contains(text)).Trim();

Radius: id 0, default 4
CSMain reads ids [0, 1, 2]
Vulkan: constantID 0, offset 0, size 4
Vulkan: constantID 1, offset 4, size 4
Vulkan: constantID 2, offset 8, size 4
Vulkan: data 080000000000000000000040
Metal: index 0, MTLDataType 29, bits 0x00000008
Metal: index 1, MTLDataType 53, bits 0x00000000
Metal: index 2, MTLDataType 3, bits 0x40000000
WebGPU: { "0": 8, "1": 0, "2": 2 }
By name: { "0": 8, "1": 0, "2": 2 }
@id(0) override Radius: i32 = 4i;
constant int Radius [[function_constant(0)]];

using OneShader;
 
// Fix Radius at compile time: it becomes a plain constant that the optimizer folds.
// The others stay pipeline-time constants.
CompileResult result = ShaderCompiler.CompileFile("Shaders/Blur.hlsl", new CompileOptions
{
    SpecializationOverrides = [new("Radius", 2), new("Blurr", 1)],
});
foreach (var d in result.Diagnostics)
    Console.WriteLine(d);
Console.WriteLine($"left as constants: {string.Join(", ", result.Reflection!.SpecConstants.Select(c => $"{c.Name} (id {c.Id})"))}");
 
CompileResult bad = ShaderCompiler.CompileFile("Shaders/Blur.hlsl", new CompileOptions
{
    SpecializationOverrides = [new("Radius", 2.5)],
});
foreach (var d in bad.Diagnostics)
    Console.WriteLine(d);

Shaders/Blur.hlsl(1,1): warning OSH0335: The shader declares no specialization constant named 'Blurr', so the override is ignored. Check the spelling against the [[vk::constant_id]] declarations.
left as constants: Horizontal (id 1), Sigma (id 2)
Shaders/Blur.hlsl(1,1): error OSH0335: The value 2.5 does not fit the int specialization constant 'Radius': use a whole number in its range.