Add PatchShaderConvertUniformBufferToPushConstant

hzqst · hzqst · commit b882f3eb3d0d · 2025-12-09T20:05:16.000+08:00
diff --git a/Graphics/GraphicsEngineVulkan/include/PipelineStateVkImpl.hpp b/Graphics/GraphicsEngineVulkan/include/PipelineStateVkImpl.hpp
@@ -133,7 +133,9 @@ class PipelineStateVkImpl final : public PipelineStateBase<EngineVkImplTraits>
     void InitializePipeline(const ComputePipelineStateCreateInfo& CreateInfo);
     void InitializePipeline(const RayTracingPipelineStateCreateInfo& CreateInfo);
 
-    void InitPushConstantInfo(const TShaderStages& ShaderStages, PushConstantInfoVk& PushConstant) noexcept(false);
+    void InitPushConstantInfoFromSignatures(PushConstantInfoVk& PushConstant) const noexcept(false);
+
+    void PatchShaderConvertUniformBufferToPushConstant(TShaderStages& ShaderStages) const noexcept(false);
 
     // TPipelineStateBase::Construct needs access to InitializePipeline
     friend TPipelineStateBase;
diff --git a/Graphics/GraphicsEngineVulkan/include/ShaderVkImpl.hpp b/Graphics/GraphicsEngineVulkan/include/ShaderVkImpl.hpp
@@ -110,6 +110,10 @@ class ShaderVkImpl final : public ShaderBase<EngineVkImplTraits>
         Size        = m_SPIRV.size() * sizeof(m_SPIRV[0]);
     }
 
+    void SetSPIRV(const std::vector<uint32_t>& SPIRV) noexcept(false);
+
+    void CreateSPIRVShaderResources() noexcept(false);
+
 private:
     void Initialize(const ShaderCreateInfo& ShaderCI,
                     const CreateInfo&       VkShaderCI) noexcept(false);
@@ -119,6 +123,8 @@ class ShaderVkImpl final : public ShaderBase<EngineVkImplTraits>
 
     std::string           m_EntryPoint;
     std::vector<uint32_t> m_SPIRV;
+
+    bool m_LoadConstantBufferReflection = false;
 };
 
 } // namespace Diligent
diff --git a/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp b/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
@@ -43,9 +43,7 @@
 #include "EngineMemory.h"
 #include "StringTools.hpp"
 
-#if !DILIGENT_NO_HLSL
-#    include "SPIRVTools.hpp"
-#endif
+#include "SPIRVTools.hpp"
 
 namespace Diligent
 {
@@ -854,33 +852,43 @@ void PipelineStateVkImpl::RemapOrVerifyShaderResources(
     }
 }
 
-void PipelineStateVkImpl::InitPushConstantInfo(const TShaderStages& ShaderStages, PushConstantInfoVk& PushConstant) noexcept(false)
+void PipelineStateVkImpl::InitPushConstantInfoFromSignatures(PushConstantInfoVk& PushConstant) const noexcept(false)
 {
-    for (const ShaderStageInfo& Stage : ShaderStages)
+    // Iterate through all signatures to find resources marked as VULKAN_PUSH_CONSTANT
+    for (Uint32 s = 0; s < m_SignatureCount; ++s)
     {
-        for (const ShaderVkImpl* pShader : Stage.Shaders)
+        const PipelineResourceSignatureVkImpl* pSignature = m_Signatures[s];
+        if (pSignature == nullptr)
+            continue;
+
+        for (Uint32 r = 0; r < pSignature->GetTotalResourceCount(); ++r)
         {
-            const auto& pShaderResources = pShader->GetShaderResources();
-            if (pShaderResources && pShaderResources->GetNumPushConstants() > 0)
+            const PipelineResourceDesc& ResDesc = pSignature->GetResourceDesc(r);
+
+            // Check if this resource is marked as a Vulkan push constant
+            if ((ResDesc.Flags & PIPELINE_RESOURCE_FLAG_VULKAN_PUSH_CONSTANT) != 0)
             {
-                // There should be at most one push constant block per shader
-                const SPIRVShaderResourceAttribs& PCAttribs = pShaderResources->GetPushConstant(0);
-                const Uint32                      PCSize    = PCAttribs.BufferStaticSize;
+                // For push constants, ArraySize contains the number of 32-bit constants
+                const Uint32 PCSize = ResDesc.ArraySize * sizeof(Uint32);
 
                 if (PushConstant.Size == 0)
                 {
-                    // First shader with push constants - record the size
+                    // First push constant resource - record the size
                     PushConstant.Size = PCSize;
                 }
                 else if (PushConstant.Size != PCSize)
                 {
-                    // Multiple shaders with different push constant sizes
-                    // This is allowed in Vulkan - take the maximum size
+                    // Multiple push constant resources with different sizes
+                    // Take the maximum size (Vulkan allows only one push constant block per pipeline)
                     PushConstant.Size = std::max(PushConstant.Size, PCSize);
                 }
 
-                // Add this shader stage to the stage flags
-                PushConstant.StageFlags |= ShaderTypeToVkShaderStageFlagBit(Stage.Type);
+                // Add shader stages to the stage flags
+                for (SHADER_TYPE ShaderStages = ResDesc.ShaderStages; ShaderStages != SHADER_TYPE_UNKNOWN;)
+                {
+                    const SHADER_TYPE ShaderType = ExtractLSB(ShaderStages);
+                    PushConstant.StageFlags |= ShaderTypeToVkShaderStageFlagBit(ShaderType);
+                }
             }
         }
     }
@@ -900,13 +908,19 @@ void PipelineStateVkImpl::InitPipelineLayout(const PipelineStateCreateInfo& Crea
     DvpValidateResourceLimits();
 #endif
 
-    // Extract push constant information from shaders
-    // Vulkan allows only one push constant block per pipeline, but it can be accessed from multiple shader stages
+    // Extract push constant information from signatures
+    // Vulkan allows only one push constant block per pipeline, but it can be accessed from multiple shader stages.
+    // We use resource attributes from m_Signatures instead of from shader stages because we may patch SPIRV
+    // later according to m_Signatures definitions (e.g., converting uniform buffers to push constants).
     PushConstantInfoVk PushConstant;
-    InitPushConstantInfo(ShaderStages, PushConstant);
+    InitPushConstantInfoFromSignatures(PushConstant);
 
     m_PipelineLayout.Create(GetDevice(), m_Signatures, m_SignatureCount, PushConstant);
 
+    // Check if any resource in shader is a uniform buffer but marked as VULKAN_PUSH_CONSTANT in signatures.
+    // If so, convert the uniform buffer to push constant in SPIRV bytecode.
+    PatchShaderConvertUniformBufferToPushConstant(ShaderStages);
+
     const bool RemapResources = (CreateInfo.Flags & PSO_CREATE_FLAG_DONT_REMAP_SHADER_RESOURCES) == 0;
     const bool VerifyBindings = !RemapResources && ((InternalFlags & PSO_CREATE_INTERNAL_FLAG_NO_SHADER_REFLECTION) == 0);
     if (RemapResources || VerifyBindings)
@@ -933,6 +947,87 @@ void PipelineStateVkImpl::InitPipelineLayout(const PipelineStateCreateInfo& Crea
     }
 }
 
+void PipelineStateVkImpl::PatchShaderConvertUniformBufferToPushConstant(TShaderStages& ShaderStages) const noexcept(false)
+{
+    // Build a set of resource names that need to be converted to push constants
+    std::unordered_set<std::string> PushConstantNames;
+    for (Uint32 s = 0; s < m_SignatureCount; ++s)
+    {
+        const PipelineResourceSignatureVkImpl* pSignature = m_Signatures[s];
+        if (pSignature == nullptr)
+            continue;
+
+        for (Uint32 r = 0; r < pSignature->GetTotalResourceCount(); ++r)
+        {
+            const PipelineResourceDesc& ResDesc = pSignature->GetResourceDesc(r);
+            if ((ResDesc.Flags & PIPELINE_RESOURCE_FLAG_VULKAN_PUSH_CONSTANT) != 0)
+            {
+                PushConstantNames.insert(ResDesc.Name);
+            }
+        }
+    }
+
+    if (PushConstantNames.empty())
+        return;
+
+    // For each shader stage, check if any uniform buffer needs to be patched
+    for (ShaderStageInfo& Stage : ShaderStages)
+    {
+        for (size_t i = 0; i < Stage.Shaders.size(); ++i)
+        {
+            ShaderVkImpl*               pShader = const_cast<ShaderVkImpl*>(Stage.Shaders[i]);
+
+            bool ShouldPatchUniformBuffer = false;
+            std::string PatchingUniformBufferName;
+            {
+                const SPIRVShaderResources* pRes    = pShader->GetShaderResources().get();
+
+                if (pRes == nullptr)
+                    continue;
+
+                // Check each uniform buffer in the shader
+                for (Uint32 ub = 0; ub < pRes->GetNumUBs(); ++ub)
+                {
+                    const SPIRVShaderResourceAttribs& UBAttribs = pRes->GetUB(ub);
+
+                    // If this uniform buffer is marked as push constant in the signature,
+                    // convert it in the SPIRV bytecode
+                    if (PushConstantNames.count(UBAttribs.Name) > 0)
+                    {
+                        PatchingUniformBufferName = UBAttribs.Name;
+                        ShouldPatchUniformBuffer = true;
+                        break;
+                    }
+                }
+            }
+
+            if(ShouldPatchUniformBuffer)
+            {
+                const std::vector<uint32_t>& SPIRV   = Stage.SPIRVs[i];
+                std::vector<uint32_t> ConvertedSPIRV = PatchSPIRVConvertUniformBufferToPushConstant(
+                    SPIRV,
+                    SPV_ENV_MAX, // Auto-detect target environment
+                    PatchingUniformBufferName);
+
+                if (!ConvertedSPIRV.empty())
+                {
+                    Stage.SPIRVs[i] = ConvertedSPIRV;
+
+                    pShader->SetSPIRV(ConvertedSPIRV);
+
+                    //Reconstruct shader resources from SPIRV
+                    pShader->CreateSPIRVShaderResources();
+                }
+                else
+                {
+                    LOG_ERROR_MESSAGE("Failed to convert uniform buffer '", PatchingUniformBufferName,
+                                    "' to push constant in shader '", pShader->GetDesc().Name, "'");
+                }
+            }
+        }
+    }
+}
+
 template <typename PSOCreateInfoType>
 PipelineStateVkImpl::TShaderStages PipelineStateVkImpl::InitInternalObjects(
     const PSOCreateInfoType&                           CreateInfo,
diff --git a/Graphics/GraphicsEngineVulkan/src/ShaderVariableManagerVk.cpp b/Graphics/GraphicsEngineVulkan/src/ShaderVariableManagerVk.cpp
@@ -667,30 +667,17 @@ void ShaderVariableManagerVk::SetInlineConstants(Uint32      ResIndex,
     const bool IsPushConstant = (ResDesc.Flags & PIPELINE_RESOURCE_FLAG_VULKAN_PUSH_CONSTANT) != 0;
     if (IsPushConstant)
     {
-        const InlineConstantBufferAttribsVk* pInlineCBAttribs   = m_pSignature->GetInlineConstantBufferAttribs();
-        const Uint32                         NumInlineCBAttribs = m_pSignature->GetNumInlineConstantBufferAttribs();
-
-        // For both Signature's static cache and SRB cache, push constant data is stored
-        // in the resource cache via GetPushConstantDataPtr()
-        for (Uint32 i = 0; i < NumInlineCBAttribs; ++i)
+        // Get the data pointer from the resource cache
+        void* pPushConstantData = m_ResourceCache.GetPushConstantDataPtr(ResIndex);
+        if (pPushConstantData != nullptr)
         {
-            const InlineConstantBufferAttribsVk& InlineCBAttr = pInlineCBAttribs[i];
-            if (!InlineCBAttr.IsPushConstant)
-                continue;
-
-            if (InlineCBAttr.ResIndex == ResIndex)
-            {
-                // Get the data pointer from the resource cache
-                void* pPushConstantData = m_ResourceCache.GetPushConstantDataPtr(InlineCBAttr.ResIndex);
-                if (pPushConstantData != nullptr)
-                {
-                    Uint32* pDstConstants = reinterpret_cast<Uint32*>(pPushConstantData);
-                    memcpy(pDstConstants + FirstConstant, pConstants, NumConstants * sizeof(Uint32));
-                }
-                return;
-            }
+            Uint32* pDstConstants = reinterpret_cast<Uint32*>(pPushConstantData);
+            memcpy(pDstConstants + FirstConstant, pConstants, NumConstants * sizeof(Uint32));
+        }
+        else
+        {
+            UNEXPECTED("Push constant buffer not found");
         }
-        UNEXPECTED("Push constant buffer not found");
     }
     else
     {
diff --git a/Graphics/GraphicsEngineVulkan/src/ShaderVkImpl.cpp b/Graphics/GraphicsEngineVulkan/src/ShaderVkImpl.cpp
@@ -229,37 +229,17 @@ void ShaderVkImpl::Initialize(const ShaderCreateInfo& ShaderCI,
 
     // We cannot create shader module here because resource bindings are assigned when
     // pipeline state is created
-
+    m_LoadConstantBufferReflection = ShaderCI.LoadConstantBufferReflection;
     // Load shader resources
     if (!m_SPIRV.empty())
     {
         if ((ShaderCI.CompileFlags & SHADER_COMPILE_FLAG_SKIP_REFLECTION) == 0)
         {
-            IMemoryAllocator& Allocator = GetRawAllocator();
-
-            std::unique_ptr<void, STDDeleterRawMem<void>> pRawMem{
-                ALLOCATE(Allocator, "Memory for SPIRVShaderResources", SPIRVShaderResources, 1),
-                STDDeleterRawMem<void>(Allocator),
-            };
-            const bool LoadShaderInputs = m_Desc.ShaderType == SHADER_TYPE_VERTEX;
-            new (pRawMem.get()) SPIRVShaderResources // May throw
-                {
-                    Allocator,
-                    m_SPIRV,
-                    m_Desc,
-                    m_Desc.UseCombinedTextureSamplers ? m_Desc.CombinedSamplerSuffix : nullptr,
-                    LoadShaderInputs,
-                    ShaderCI.LoadConstantBufferReflection,
-                    m_EntryPoint //
-                };
+            CreateSPIRVShaderResources();
+
             VERIFY_EXPR(ShaderCI.ByteCode != nullptr || m_EntryPoint == ShaderCI.EntryPoint ||
                         (m_EntryPoint == "main" && (ShaderCI.CompileFlags & SHADER_COMPILE_FLAG_HLSL_TO_SPIRV_VIA_GLSL) != 0));
-            m_pShaderResources.reset(static_cast<SPIRVShaderResources*>(pRawMem.release()), STDDeleterRawMem<SPIRVShaderResources>(Allocator));
 
-            if (LoadShaderInputs && m_pShaderResources->IsHLSLSource())
-            {
-                m_pShaderResources->MapHLSLVertexShaderInputs(m_SPIRV);
-            }
         }
         else
         {
@@ -270,6 +250,38 @@ void ShaderVkImpl::Initialize(const ShaderCreateInfo& ShaderCI,
     }
 }
 
+void ShaderVkImpl::SetSPIRV(const std::vector<uint32_t>& SPIRV) noexcept(false)
+{
+    m_SPIRV = SPIRV;
+}
+
+void ShaderVkImpl::CreateSPIRVShaderResources() noexcept(false)
+{
+    IMemoryAllocator& Allocator = GetRawAllocator();
+
+    std::unique_ptr<void, STDDeleterRawMem<void>> pRawMem{
+        ALLOCATE(Allocator, "Memory for SPIRVShaderResources", SPIRVShaderResources, 1),
+        STDDeleterRawMem<void>(Allocator),
+    };
+    const bool LoadShaderInputs = m_Desc.ShaderType == SHADER_TYPE_VERTEX;
+    new (pRawMem.get()) SPIRVShaderResources // May throw
+        {
+            Allocator,
+            m_SPIRV,
+            m_Desc,
+            m_Desc.UseCombinedTextureSamplers ? m_Desc.CombinedSamplerSuffix : nullptr,
+            LoadShaderInputs,
+            m_LoadConstantBufferReflection,
+            m_EntryPoint //
+        };
+
+    m_pShaderResources.reset(static_cast<SPIRVShaderResources*>(pRawMem.release()), STDDeleterRawMem<SPIRVShaderResources>(Allocator));
+
+    if (LoadShaderInputs && m_pShaderResources->IsHLSLSource())
+    {
+        m_pShaderResources->MapHLSLVertexShaderInputs(m_SPIRV);
+    }
+}
 
 ShaderVkImpl::ShaderVkImpl(IReferenceCounters*     pRefCounters,
                            RenderDeviceVkImpl*     pRenderDeviceVk,
