SPIRVShaderResources: add specialization constants support

hzqst · TheMostDiligent · commit 2385c1e87e69 · 2026-02-25T23:36:45.000-08:00
diff --git a/Graphics/ShaderTools/include/SPIRVShaderResources.hpp b/Graphics/ShaderTools/include/SPIRVShaderResources.hpp
@@ -44,6 +44,7 @@
 //    | Accel Structs     |
 //    | Push Constants    | _ _ _ m_TotalResources
 //    | Stage Inputs      |
+//    | Spec Constants    |
 //    | Resource Names    |
 
 #include <memory>
@@ -192,6 +193,28 @@ struct SPIRVShaderStageInputAttribs
 };
 static_assert(sizeof(SPIRVShaderStageInputAttribs) % sizeof(void*) == 0, "Size of SPIRVShaderStageInputAttribs struct must be multiple of sizeof(void*)");
 
+// sizeof(SPIRVSpecializationConstantAttribs) == 24, msvc x64
+struct SPIRVSpecializationConstantAttribs
+{
+    // clang-format off
+    SPIRVSpecializationConstantAttribs(const char*            _Name,
+                                       uint32_t               _SpecId,
+                                       uint32_t               _Size,
+                                       SHADER_CODE_BASIC_TYPE _BasicType) :
+        Name     {_Name},
+        SpecId   {_SpecId},
+        Size     {_Size},
+        BasicType{_BasicType}
+    {}
+    // clang-format on
+
+    const char* const            Name;
+    const uint32_t               SpecId;
+    const uint32_t               Size; // Byte size of the scalar type
+    const SHADER_CODE_BASIC_TYPE BasicType;
+};
+static_assert(sizeof(SPIRVSpecializationConstantAttribs) % sizeof(void*) == 0, "Size of SPIRVSpecializationConstantAttribs struct must be multiple of sizeof(void*)");
+
 /// Diligent::SPIRVShaderResources class
 class SPIRVShaderResources
 {
@@ -271,8 +294,9 @@ class SPIRVShaderResources
     Uint32 GetNumAccelStructs ()const noexcept{ return GetNumResources(ResourceClass::AccelStruct);     }
     Uint32 GetNumPushConstants()const noexcept{ return GetNumResources(ResourceClass::PushConstant);    }
 
-    Uint32 GetTotalResources()      const noexcept { return m_Offsets[static_cast<size_t>(ResourceClass::NumClasses)]; }
-    Uint32 GetNumShaderStageInputs()const noexcept { return m_NumShaderStageInputs; }
+    Uint32 GetTotalResources()        const noexcept { return m_Offsets[static_cast<size_t>(ResourceClass::NumClasses)]; }
+    Uint32 GetNumShaderStageInputs()  const noexcept { return m_NumShaderStageInputs; }
+    Uint32 GetNumSpecConstants()      const noexcept { return m_NumSpecConstants; }
 
     const SPIRVShaderResourceAttribs& GetUB          (Uint32 n)const noexcept{ return GetResAttribs(ResourceClass::UniformBuffer,   n); }
     const SPIRVShaderResourceAttribs& GetSB          (Uint32 n)const noexcept{ return GetResAttribs(ResourceClass::StorageBuffer,   n); }
@@ -294,6 +318,14 @@ class SPIRVShaderResources
         return reinterpret_cast<const SPIRVShaderStageInputAttribs*>(ResourceMemoryEnd)[n];
     }
 
+    const SPIRVSpecializationConstantAttribs& GetSpecConstant(Uint32 n) const noexcept
+    {
+        VERIFY(n < m_NumSpecConstants, "Specialization constant index (", n, ") is out of range. Total spec constant count: ", m_NumSpecConstants);
+        const SPIRVShaderResourceAttribs*   ResourceMemoryEnd = reinterpret_cast<const SPIRVShaderResourceAttribs*>(m_MemoryBuffer.get()) + GetTotalResources();
+        const SPIRVShaderStageInputAttribs* StageInputsEnd    = reinterpret_cast<const SPIRVShaderStageInputAttribs*>(ResourceMemoryEnd) + m_NumShaderStageInputs;
+        return reinterpret_cast<const SPIRVSpecializationConstantAttribs*>(StageInputsEnd)[n];
+    }
+
     const ShaderCodeBufferDesc* GetUniformBufferDesc(Uint32 Index) const
     {
         if (Index >= GetNumUBs())
@@ -452,6 +484,7 @@ class SPIRVShaderResources
     void Initialize(IMemoryAllocator&       Allocator,
                     const ResourceCounters& Counters,
                     Uint32                  NumShaderStageInputs,
+                    Uint32                  NumSpecConstants,
                     size_t                  ResourceNamesPoolSize,
                     StringPool&             ResourceNamesPool);
 
@@ -486,8 +519,13 @@ class SPIRVShaderResources
         return const_cast<SPIRVShaderStageInputAttribs&>(const_cast<const SPIRVShaderResources*>(this)->GetShaderStageInputAttribs(n));
     }
 
+    SPIRVSpecializationConstantAttribs& GetSpecConstant(Uint32 n) noexcept
+    {
+        return const_cast<SPIRVSpecializationConstantAttribs&>(const_cast<const SPIRVShaderResources*>(this)->GetSpecConstant(n));
+    }
+
     // Memory buffer that holds all resources as continuous chunk of memory:
-    // |  UBs  |  SBs  |  StrgImgs  |  SmplImgs  |  ACs  |  SepSamplers  |  SepImgs  | Stage Inputs | Resource Names |
+    // |  UBs  |  SBs  |  StrgImgs  |  SmplImgs  |  ACs  |  SepSamplers  |  SepImgs  | Stage Inputs | Spec Constants | Resource Names |
     std::unique_ptr<void, STDDeleterRawMem<void>> m_MemoryBuffer;
     std::unique_ptr<void, STDDeleterRawMem<void>> m_UBReflectionBuffer;
 
@@ -498,6 +536,7 @@ class SPIRVShaderResources
     std::array<OffsetType, static_cast<size_t>(ResourceClass::NumClasses) + 1> m_Offsets;
 
     OffsetType m_NumShaderStageInputs = 0;
+    OffsetType m_NumSpecConstants     = 0;
 
     SHADER_TYPE m_ShaderType = SHADER_TYPE_UNKNOWN;
 
diff --git a/Graphics/ShaderTools/src/SPIRVShaderResources.cpp b/Graphics/ShaderTools/src/SPIRVShaderResources.cpp
@@ -613,9 +613,57 @@ SPIRVShaderResources::SPIRVShaderResources(IMemoryAllocator&     Allocator,
     ResCounters.NumPushConstants = static_cast<Uint32>(resources.push_constant_buffers.size());
     static_assert(Uint32{SPIRVShaderResourceAttribs::ResourceType::NumResourceTypes} == 13, "Please set the new resource type counter here");
 
+    // Specialization constants reflection
+    struct SpecConstInfo
+    {
+        std::string            Name;
+        uint32_t               SpecId    = 0;
+        uint32_t               Size      = 0;
+        SHADER_CODE_BASIC_TYPE BasicType = SHADER_CODE_BASIC_TYPE_UNKNOWN;
+    };
+    std::vector<SpecConstInfo> SpecConstants;
+    Uint32                     NumSpecConstants = 0;
+
+    {
+        diligent_spirv_cross::SmallVector<diligent_spirv_cross::SpecializationConstant> spec_consts =
+            Compiler.get_specialization_constants();
+        for (const diligent_spirv_cross::SpecializationConstant& sc : spec_consts)
+        {
+            const diligent_spirv_cross::SPIRConstant& Constant = Compiler.get_constant(sc.id);
+            const diligent_spirv_cross::SPIRType&     Type     = Compiler.get_type(Constant.constant_type);
+
+            // Only support scalar specialization constants
+            if (Type.vecsize != 1 || Type.columns != 1)
+            {
+                LOG_WARNING_MESSAGE("Specialization constant '", Compiler.get_name(sc.id),
+                                    "' (SpecId=", sc.constant_id, ") in shader '", CI.Name,
+                                    "' is not a scalar type and will be skipped.");
+                continue;
+            }
+
+            SpecConstInfo Info;
+            Info.Name   = Compiler.get_name(sc.id);
+            Info.SpecId = sc.constant_id;
+            // OpTypeBool has width==1 in SPIRV-Cross; use 4 bytes (VkBool32) for bool specialization constants
+            Info.Size      = Type.basetype == diligent_spirv_cross::SPIRType::Boolean ? 4 : Type.width / 8;
+            Info.BasicType = SpirvBaseTypeToShaderCodeBasicType(Type.basetype);
+
+            if (Info.Name.empty())
+            {
+                LOG_WARNING_MESSAGE("Specialization constant with SpecId=", sc.constant_id,
+                                    " in shader '", CI.Name, "' has no name (OpName) and will be skipped.");
+                continue;
+            }
+
+            ResourceNamesPoolSize += Info.Name.length() + 1;
+            SpecConstants.emplace_back(std::move(Info));
+        }
+        NumSpecConstants = static_cast<Uint32>(SpecConstants.size());
+    }
+
     // Resource names pool is only needed to facilitate string allocation.
     StringPool ResourceNamesPool;
-    Initialize(Allocator, ResCounters, NumShaderStageInputs, ResourceNamesPoolSize, ResourceNamesPool);
+    Initialize(Allocator, ResCounters, NumShaderStageInputs, NumSpecConstants, ResourceNamesPoolSize, ResourceNamesPool);
 
     // Uniform buffer reflections
     std::vector<ShaderCodeBufferDescX> UBReflections;
@@ -842,6 +890,22 @@ SPIRVShaderResources::SPIRVShaderResources(IMemoryAllocator&     Allocator,
         VERIFY_EXPR(CurrStageInput == GetNumShaderStageInputs());
     }
 
+    if (!SpecConstants.empty())
+    {
+        Uint32 CurrSpecConst = 0;
+        for (const SpecConstInfo& SC : SpecConstants)
+        {
+            new (&GetSpecConstant(CurrSpecConst++)) SPIRVSpecializationConstantAttribs //
+                {
+                    ResourceNamesPool.CopyString(SC.Name),
+                    SC.SpecId,
+                    SC.Size,
+                    SC.BasicType //
+                };
+        }
+        VERIFY_EXPR(CurrSpecConst == GetNumSpecConstants());
+    }
+
     VERIFY(ResourceNamesPool.GetRemainingSize() == 0, "Names pool must be empty");
 
     if (m_ShaderType == SHADER_TYPE_COMPUTE)
@@ -862,6 +926,7 @@ SPIRVShaderResources::SPIRVShaderResources(IMemoryAllocator&     Allocator,
 void SPIRVShaderResources::Initialize(IMemoryAllocator&       Allocator,
                                       const ResourceCounters& Counters,
                                       Uint32                  NumShaderStageInputs,
+                                      Uint32                  NumSpecConstants,
                                       size_t                  ResourceNamesPoolSize,
                                       StringPool&             ResourceNamesPool)
 {
@@ -890,12 +955,16 @@ void SPIRVShaderResources::Initialize(IMemoryAllocator&       Allocator,
     VERIFY(NumShaderStageInputs <= MaxOffset, "Max offset exceeded");
     m_NumShaderStageInputs = static_cast<OffsetType>(NumShaderStageInputs);
 
+    VERIFY(NumSpecConstants <= MaxOffset, "Max offset exceeded");
+    m_NumSpecConstants = static_cast<OffsetType>(NumSpecConstants);
+
     size_t AlignedResourceNamesPoolSize = AlignUp(ResourceNamesPoolSize, sizeof(void*));
 
     static_assert(sizeof(SPIRVShaderResourceAttribs) % sizeof(void*) == 0, "Size of SPIRVShaderResourceAttribs struct must be multiple of sizeof(void*)");
     // clang-format off
     size_t MemorySize = GetTotalResources()           * sizeof(SPIRVShaderResourceAttribs) +
                         m_NumShaderStageInputs        * sizeof(SPIRVShaderStageInputAttribs) +
+                        m_NumSpecConstants            * sizeof(SPIRVSpecializationConstantAttribs) +
                         AlignedResourceNamesPoolSize  * sizeof(char);
 
     VERIFY_EXPR(GetNumUBs()           == Counters.NumUBs);
@@ -917,7 +986,8 @@ void SPIRVShaderResources::Initialize(IMemoryAllocator&       Allocator,
         m_MemoryBuffer  = std::unique_ptr<void, STDDeleterRawMem<void>>(pRawMem, Allocator);
         char* NamesPool = reinterpret_cast<char*>(m_MemoryBuffer.get()) +
             GetTotalResources() * sizeof(SPIRVShaderResourceAttribs) +
-            m_NumShaderStageInputs * sizeof(SPIRVShaderStageInputAttribs);
+            m_NumShaderStageInputs * sizeof(SPIRVShaderStageInputAttribs) +
+            m_NumSpecConstants * sizeof(SPIRVSpecializationConstantAttribs);
         ResourceNamesPool.AssignMemory(NamesPool, ResourceNamesPoolSize);
     }
 }
@@ -951,6 +1021,9 @@ SPIRVShaderResources::~SPIRVShaderResources()
     for (Uint32 n = 0; n < GetNumShaderStageInputs(); ++n)
         GetShaderStageInputAttribs(n).~SPIRVShaderStageInputAttribs();
 
+    for (Uint32 n = 0; n < GetNumSpecConstants(); ++n)
+        GetSpecConstant(n).~SPIRVSpecializationConstantAttribs();
+
     for (Uint32 n = 0; n < GetNumAccelStructs(); ++n)
         GetAccelStruct(n).~SPIRVShaderResourceAttribs();
 
@@ -1191,6 +1264,18 @@ std::string SPIRVShaderResources::DumpResources() const
     );
     VERIFY_EXPR(ResNum == GetTotalResources());
 
+    if (GetNumSpecConstants() > 0)
+    {
+        ss << std::endl
+           << "Specialization constants (" << GetNumSpecConstants() << "):";
+        for (Uint32 n = 0; n < GetNumSpecConstants(); ++n)
+        {
+            const auto& SC = GetSpecConstant(n);
+            ss << std::endl
+               << "  '" << SC.Name << "' SpecId=" << SC.SpecId << " Size=" << SC.Size;
+        }
+    }
+
     return ss.str();
 }
 
diff --git a/Tests/DiligentCoreTest/assets/shaders/SPIRV/SpecializationConstants.psh b/Tests/DiligentCoreTest/assets/shaders/SPIRV/SpecializationConstants.psh
@@ -0,0 +1,19 @@
+// Test shader for specialization constants reflection
+[[vk::constant_id(0)]] const bool   g_EnableFeature = true;
+[[vk::constant_id(1)]] const int    g_IntParam      = 42;
+[[vk::constant_id(2)]] const uint   g_UintParam     = 7;
+[[vk::constant_id(3)]] const float  g_FloatParam    = 1.0f;
+
+float4 main() : SV_Target
+{
+    float4 result = float4(0, 0, 0, 1);
+
+    if (g_EnableFeature)
+    {
+        result.x = float(g_IntParam);
+        result.y = float(g_UintParam);
+        result.z = g_FloatParam;
+    }
+
+    return result;
+}
diff --git a/Tests/DiligentCoreTest/src/ShaderTools/SPIRVShaderResourcesTest.cpp b/Tests/DiligentCoreTest/src/ShaderTools/SPIRVShaderResourcesTest.cpp
@@ -589,4 +589,70 @@ TEST_F(SPIRVShaderResourcesTest, MixedResources_DXC)
     TestMixedResources(SHADER_COMPILER_DXC);
 }
 
+
+struct SPIRVSpecConstRefAttribs
+{
+    const char* const            Name;
+    const uint32_t               SpecId;
+    const uint32_t               Size;
+    const SHADER_CODE_BASIC_TYPE BasicType;
+};
+
+void TestSpecializationConstants(SHADER_COMPILER Compiler)
+{
+    std::vector<unsigned int> SPIRV;
+    ASSERT_NO_FATAL_FAILURE(CompileSPIRV("SpecializationConstants.psh", Compiler, SHADER_TYPE_PIXEL, SHADER_SOURCE_LANGUAGE_HLSL, SPIRV));
+
+    if (::testing::Test::IsSkipped())
+        return;
+
+    SPIRVShaderResources::CreateInfo ResCI;
+    ResCI.ShaderType = SHADER_TYPE_PIXEL;
+    ResCI.Name       = "SpecConstants test";
+    SPIRVShaderResources Resources{
+        GetRawAllocator(),
+        SPIRV,
+        ResCI,
+    };
+
+    LOG_INFO_MESSAGE("SPIRV Resources:\n", Resources.DumpResources());
+
+    const std::vector<SPIRVSpecConstRefAttribs> RefSpecConstants = {
+        {"g_EnableFeature", 0, 4, SHADER_CODE_BASIC_TYPE_BOOL},
+        {"g_IntParam", 1, 4, SHADER_CODE_BASIC_TYPE_INT},
+        {"g_UintParam", 2, 4, SHADER_CODE_BASIC_TYPE_UINT},
+        {"g_FloatParam", 3, 4, SHADER_CODE_BASIC_TYPE_FLOAT},
+    };
+
+    const SPIRVShaderResources& ConstResources = Resources;
+    EXPECT_EQ(ConstResources.GetNumSpecConstants(), static_cast<Uint32>(RefSpecConstants.size()));
+
+    // Build a map from name to reference for order-independent matching
+    std::unordered_map<std::string, const SPIRVSpecConstRefAttribs*> RefMap;
+    for (const SPIRVSpecConstRefAttribs& Ref : RefSpecConstants)
+        RefMap[Ref.Name] = &Ref;
+
+    for (Uint32 i = 0; i < ConstResources.GetNumSpecConstants(); ++i)
+    {
+        const SPIRVSpecializationConstantAttribs& SC = ConstResources.GetSpecConstant(i);
+        const auto                                it = RefMap.find(SC.Name);
+        ASSERT_NE(it, RefMap.end()) << "Specialization constant '" << SC.Name << "' is not found in the reference list";
+
+        const auto* pRef = it->second;
+        EXPECT_EQ(SC.SpecId, pRef->SpecId) << SC.Name;
+        EXPECT_EQ(SC.Size, pRef->Size) << SC.Name;
+        EXPECT_EQ(SC.BasicType, pRef->BasicType) << SC.Name;
+    }
+}
+
+TEST_F(SPIRVShaderResourcesTest, SpecializationConstants_GLSLang)
+{
+    TestSpecializationConstants(SHADER_COMPILER_GLSLANG);
+}
+
+TEST_F(SPIRVShaderResourcesTest, SpecializationConstants_DXC)
+{
+    TestSpecializationConstants(SHADER_COMPILER_DXC);
+}
+
 } // namespace
