Add stuffs for vk push constants

Author: hzqst

Graphics/GraphicsAccessories/src/GraphicsAccessories.cpp

@@ -1637,7 +1637,7 @@ String GetPipelineResourceFlagsString(PIPELINE_RESOURCE_FLAGS Flags, bool GetFul
 
         PIPELINE_RESOURCE_FLAGS Flag = ExtractLSB(Flags);
 
-        static_assert(PIPELINE_RESOURCE_FLAG_LAST == (1u << 5), "Please update the switch below to handle the new pipeline resource flag.");
+        static_assert(PIPELINE_RESOURCE_FLAG_LAST == (1u << 6), "Please update the switch below to handle the new pipeline resource flag.");
         switch (Flag)
         {
             case PIPELINE_RESOURCE_FLAG_NO_DYNAMIC_BUFFERS:
@@ -1664,6 +1664,10 @@ String GetPipelineResourceFlagsString(PIPELINE_RESOURCE_FLAGS Flags, bool GetFul
                 Str.append(GetFullName ? "PIPELINE_RESOURCE_FLAG_GENERAL_INPUT_ATTACHMENT" : "GENERAL_INPUT_ATTACHMENT");
                 break;
 
+            case PIPELINE_RESOURCE_FLAG_VULKAN_PUSH_CONSTANT:
+                Str.append(GetFullName ? "PIPELINE_RESOURCE_FLAG_VULKAN_PUSH_CONSTANT" : "VULKAN_PUSH_CONSTANT");
+                break;
+
             default:
                 UNEXPECTED("Unexpected pipeline resource flag");
         }
@@ -1812,7 +1816,7 @@ PIPELINE_RESOURCE_FLAGS GetValidPipelineResourceFlags(SHADER_RESOURCE_TYPE Resou
     switch (ResourceType)
     {
         case SHADER_RESOURCE_TYPE_CONSTANT_BUFFER:
-            return PIPELINE_RESOURCE_FLAG_NO_DYNAMIC_BUFFERS | PIPELINE_RESOURCE_FLAG_RUNTIME_ARRAY | PIPELINE_RESOURCE_FLAG_INLINE_CONSTANTS;
+            return PIPELINE_RESOURCE_FLAG_NO_DYNAMIC_BUFFERS | PIPELINE_RESOURCE_FLAG_RUNTIME_ARRAY | PIPELINE_RESOURCE_FLAG_INLINE_CONSTANTS | PIPELINE_RESOURCE_FLAG_VULKAN_PUSH_CONSTANT;
 
         case SHADER_RESOURCE_TYPE_TEXTURE_SRV:
             return PIPELINE_RESOURCE_FLAG_COMBINED_SAMPLER | PIPELINE_RESOURCE_FLAG_RUNTIME_ARRAY;

Graphics/GraphicsEngine/interface/PipelineResourceSignature.h

@@ -148,7 +148,19 @@ DILIGENT_TYPED_ENUM(PIPELINE_RESOURCE_FLAGS, Uint8)
     /// \note This flag is only valid in Vulkan.
     PIPELINE_RESOURCE_FLAG_GENERAL_INPUT_ATTACHMENT = 1u << 5,
 
-    PIPELINE_RESOURCE_FLAG_LAST = PIPELINE_RESOURCE_FLAG_GENERAL_INPUT_ATTACHMENT
+    /// Internal flag indicating that this inline constant is a Vulkan push constant
+    /// (declared with push_constant storage class in SPIR-V).
+    /// This flag is automatically set by the engine during shader reflection and should not
+    /// be set manually by the application.
+    ///
+    /// When this flag is set along with PIPELINE_RESOURCE_FLAG_INLINE_CONSTANTS, the engine
+    /// will use vkCmdPushConstants instead of a dynamic uniform buffer to update the constants.
+    ///
+    /// \note This flag is only valid in Vulkan and is always combined with
+    ///       PIPELINE_RESOURCE_FLAG_INLINE_CONSTANTS.
+    PIPELINE_RESOURCE_FLAG_VULKAN_PUSH_CONSTANT = 1u << 6,
+
+    PIPELINE_RESOURCE_FLAG_LAST = PIPELINE_RESOURCE_FLAG_VULKAN_PUSH_CONSTANT
 };
 DEFINE_FLAG_ENUM_OPERATORS(PIPELINE_RESOURCE_FLAGS);
 

Graphics/GraphicsEngine/src/PipelineResourceSignatureBase.cpp

@@ -113,10 +113,10 @@ void ValidatePipelineResourceSignatureDesc(const PipelineResourceSignatureDesc&
 
         if ((Res.Flags & PIPELINE_RESOURCE_FLAG_INLINE_CONSTANTS) != 0)
         {
-            if (Res.Flags != PIPELINE_RESOURCE_FLAG_INLINE_CONSTANTS)
+            if ((Res.Flags & (PIPELINE_RESOURCE_FLAG_INLINE_CONSTANTS | PIPELINE_RESOURCE_FLAG_VULKAN_PUSH_CONSTANT)) != Res.Flags)
             {
                 LOG_PRS_ERROR_AND_THROW("Incorrect Desc.Resources[", i, "].Flags (", GetPipelineResourceFlagsString(Res.Flags),
-                                        "). INLINE_CONSTANTS flag cannot be combined with other flags.");
+                                        "). INLINE_CONSTANTS and VULKAN_PUSH_CONSTANT flags cannot be combined with other flags.");
             }
 
             if (Res.ArraySize > 64)

Graphics/GraphicsEngineVulkan/include/DeviceContextVkImpl.hpp

@@ -321,6 +321,9 @@ class DeviceContextVkImpl final : public DeviceContextNextGenBase<EngineVkImplTr
     /// Implementation of IDeviceContextVk::GetVkCommandBuffer().
     virtual VkCommandBuffer DILIGENT_CALL_TYPE GetVkCommandBuffer() override final;
 
+    /// Implementation of IDeviceContextVk::SetPushConstants().
+    virtual void DILIGENT_CALL_TYPE SetPushConstants(const void* pData, Uint32 Offset, Uint32 Size) override final;
+
     // Transitions BLAS state from OldState to NewState, and optionally updates internal state.
     // If OldState == RESOURCE_STATE_UNKNOWN, internal BLAS state is used as old state.
     void TransitionBLASState(BottomLevelASVkImpl& BLAS,
@@ -514,11 +517,19 @@ class DeviceContextVkImpl final : public DeviceContextNextGenBase<EngineVkImplTr
         /// Current graphics PSO uses no depth/render targets.
         bool NullRenderTargets = false;
 
+        /// Flag indicating if push constants have been updated since last draw/dispatch
+        bool PushConstantsDirty = false;
+
         Uint32 NumCommands = 0;
 
         VkPipelineBindPoint vkPipelineBindPoint = VK_PIPELINE_BIND_POINT_MAX_ENUM;
     } m_State;
 
+    // Push constants data storage (max size is typically 128-256 bytes, use 256 for safety)
+    static constexpr Uint32 MAX_PUSH_CONSTANTS_SIZE = 256;
+    std::array<Uint8, MAX_PUSH_CONSTANTS_SIZE> m_PushConstantsData = {};
+    Uint32 m_PushConstantsDataSize = 0; // Actual size of valid push constants data
+
     // Graphics/mesh, compute, ray tracing
     static constexpr Uint32 NUM_PIPELINE_BIND_POINTS = 3;
 
@@ -558,6 +569,9 @@ class DeviceContextVkImpl final : public DeviceContextNextGenBase<EngineVkImplTr
 
     void UpdateInlineConstantBuffers(ResourceBindInfo& BindInfo);
 
+    // Commits push constants to the command buffer if they are dirty
+    void CommitPushConstants();
+
 #ifdef DILIGENT_DEVELOPMENT
     void DvpValidateCommittedShaderResources(ResourceBindInfo& BindInfo);
 #endif

Graphics/GraphicsEngineVulkan/include/PipelineLayoutVk.hpp

@@ -40,14 +40,24 @@ namespace Diligent
 class RenderDeviceVkImpl;
 class PipelineResourceSignatureVkImpl;
 
+/// Push constant information extracted from shaders
+struct PushConstantInfoVk
+{
+    VkShaderStageFlags StageFlags = 0;
+    Uint32             Size       = 0;
+};
+
 /// Implementation of the Diligent::PipelineLayoutVk class
 class PipelineLayoutVk
 {
 public:
     PipelineLayoutVk();
     ~PipelineLayoutVk();
 
-    void Create(RenderDeviceVkImpl* pDeviceVk, RefCntAutoPtr<PipelineResourceSignatureVkImpl> ppSignatures[], Uint32 SignatureCount) noexcept(false);
+    void Create(RenderDeviceVkImpl*                             pDeviceVk,
+                RefCntAutoPtr<PipelineResourceSignatureVkImpl>  ppSignatures[],
+                Uint32                                          SignatureCount,
+                const PushConstantInfoVk&                       PushConstant = {}) noexcept(false);
     void Release(RenderDeviceVkImpl* pDeviceVkImpl, Uint64 CommandQueueMask);
 
     VkPipelineLayout GetVkPipelineLayout() const { return m_VkPipelineLayout; }
@@ -59,6 +69,15 @@ class PipelineLayoutVk
         return m_FirstDescrSetIndex[Index];
     }
 
+    // Returns true if this pipeline layout has push constants
+    bool HasPushConstants() const { return m_PushConstantSize > 0; }
+
+    // Returns the size of push constants in bytes
+    Uint32 GetPushConstantSize() const { return m_PushConstantSize; }
+
+    // Returns the shader stage flags for push constants
+    VkShaderStageFlags GetPushConstantStageFlags() const { return m_PushConstantStageFlags; }
+
 private:
     VulkanUtilities::PipelineLayoutWrapper m_VkPipelineLayout;
 
@@ -70,6 +89,12 @@ class PipelineLayoutVk
     // (Maximum is MAX_RESOURCE_SIGNATURES * 2)
     Uint8 m_DescrSetCount = 0;
 
+    // Push constant size in bytes
+    Uint32 m_PushConstantSize = 0;
+
+    // Shader stages that use push constants
+    VkShaderStageFlags m_PushConstantStageFlags = 0;
+
 #ifdef DILIGENT_DEBUG
     Uint32 m_DbgMaxBindIndex = 0;
 #endif

Graphics/GraphicsEngineVulkan/include/PipelineResourceSignatureVkImpl.hpp

@@ -60,14 +60,16 @@ struct ImmutableSamplerAttribsVk
 ASSERT_SIZEOF(ImmutableSamplerAttribsVk, 8, "The struct is used in serialization and must be tightly packed");
 
 /// Inline constant buffer attributes for Vulkan backend.
-/// Since Vulkan only supports one push constant block per pipeline, inline constants
-/// are emulated using dynamic uniform buffers (similar to D3D11 backend).
+/// Inline constants can be either:
+/// 1. True push constants (from SPIR-V push_constant storage class) - use vkCmdPushConstants
+/// 2. Emulated inline constants - use dynamic uniform buffers (similar to D3D11 backend)
 struct InlineConstantBufferAttribsVk
 {
-    Uint32                      DescrSet     = 0; // Descriptor set index
-    Uint32                      BindingIndex = 0; // Binding index within the descriptor set
-    Uint32                      NumConstants = 0; // Number of 32-bit constants
-    RefCntAutoPtr<BufferVkImpl> pBuffer;          // Internal dynamic uniform buffer
+    Uint32                      DescrSet         = 0;     // Descriptor set index
+    Uint32                      BindingIndex     = 0;     // Binding index within the descriptor set
+    Uint32                      NumConstants     = 0;     // Number of 32-bit constants
+    bool                        IsPushConstant = false; // True if this is a Vulkan push constant (not emulated)
+    RefCntAutoPtr<BufferVkImpl> pBuffer;                  // Internal dynamic uniform buffer (null for true push constants)
 };
 
 struct PipelineResourceSignatureInternalDataVk : PipelineResourceSignatureInternalData<PipelineResourceAttribsVk, ImmutableSamplerAttribsVk>
@@ -156,6 +158,11 @@ class PipelineResourceSignatureVkImpl final : public PipelineResourceSignatureBa
     // Returns the inline constant buffer attributes
     const InlineConstantBufferAttribsVk* GetInlineConstantBuffers() const { return m_InlineConstantBuffers.get(); }
 
+    // Adds PIPELINE_RESOURCE_FLAG_VULKAN_PUSH_CONSTANT flag to a resource if it has
+    // PIPELINE_RESOURCE_FLAG_INLINE_CONSTANTS but is missing the Vulkan push constant flag.
+    // This is called during pipeline state creation when shader reflection detects a push constant.
+    void UpdatePushConstantFlags(Uint32 ResIndex);
+
 #ifdef DILIGENT_DEVELOPMENT
     /// Verifies committed resource using the SPIRV resource attributes from the PSO.
     bool DvpValidateCommittedResource(const DeviceContextVkImpl*        pDeviceCtx,

Graphics/GraphicsEngineVulkan/include/PipelineStateVkImpl.hpp

@@ -133,6 +133,8 @@ 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);
+
     // TPipelineStateBase::Construct needs access to InitializePipeline
     friend TPipelineStateBase;
 

Graphics/GraphicsEngineVulkan/include/VulkanUtilities/CommandBuffer.hpp

@@ -504,6 +504,18 @@ class CommandBuffer
         vkCmdBindDescriptorSets(m_VkCmdBuffer, pipelineBindPoint, layout, firstSet, descriptorSetCount, pDescriptorSets, dynamicOffsetCount, pDynamicOffsets);
     }
 
+    __forceinline void PushConstants(VkPipelineLayout   layout,
+                                     VkShaderStageFlags stageFlags,
+                                     uint32_t           offset,
+                                     uint32_t           size,
+                                     const void*        pValues)
+    {
+        VERIFY_EXPR(m_VkCmdBuffer != VK_NULL_HANDLE);
+        VERIFY_EXPR(pValues != nullptr);
+        VERIFY_EXPR(size > 0);
+        vkCmdPushConstants(m_VkCmdBuffer, layout, stageFlags, offset, size, pValues);
+    }
+
     __forceinline void CopyBuffer(VkBuffer            srcBuffer,
                                   VkBuffer            dstBuffer,
                                   uint32_t            regionCount,

Graphics/GraphicsEngineVulkan/interface/DeviceContextVk.h

@@ -84,6 +84,25 @@ DILIGENT_BEGIN_INTERFACE(IDeviceContextVk, IDeviceContext)
     /// calling IDeviceContext::InvalidateState() and then manually restore all required states via
     /// appropriate Diligent API calls.
     VIRTUAL VkCommandBuffer METHOD(GetVkCommandBuffer)(THIS) PURE;
+
+    /// Sets push constant data for the currently bound pipeline.
+
+    /// \param [in] pData  - Pointer to the push constant data.
+    /// \param [in] Offset - Offset in bytes from the start of the push constant block.
+    /// \param [in] Size   - Size of the data to set in bytes.
+    ///
+    /// \remarks This method sets push constant data that will be used in subsequent draw or dispatch calls.
+    ///          The pipeline state must be set before calling this method.
+    ///          Push constants are a Vulkan feature that allows small amounts of data to be passed to
+    ///          shaders efficiently without using descriptor sets. In shaders, push constants are declared
+    ///          using the layout(push_constant) qualifier in GLSL or [[vk::push_constant]] attribute in HLSL.
+    ///
+    /// \note    The total size of push constants is limited by the device (typically 128-256 bytes).
+    ///          Check VkPhysicalDeviceLimits::maxPushConstantsSize for the exact limit.
+    VIRTUAL void METHOD(SetPushConstants)(THIS_
+                                          const void* pData,
+                                          Uint32      Offset,
+                                          Uint32      Size) PURE;
 };
 DILIGENT_END_INTERFACE
 
@@ -95,6 +114,7 @@ DILIGENT_END_INTERFACE
 
 #    define IDeviceContextVk_TransitionImageLayout(This, ...) CALL_IFACE_METHOD(DeviceContextVk, TransitionImageLayout, This, __VA_ARGS__)
 #    define IDeviceContextVk_BufferMemoryBarrier(This, ...)   CALL_IFACE_METHOD(DeviceContextVk, BufferMemoryBarrier,   This, __VA_ARGS__)
+#    define IDeviceContextVk_SetPushConstants(This, ...)      CALL_IFACE_METHOD(DeviceContextVk, SetPushConstants,      This, __VA_ARGS__)
 
 // clang-format on
 

Graphics/GraphicsEngineVulkan/src/DeviceContextVkImpl.cpp

@@ -390,6 +390,16 @@ void DeviceContextVkImpl::SetPipelineState(IPipelineState* pPipelineState)
 
     // Reserve space to store all dynamic buffer offsets
     m_DynamicBufferOffsets.resize(TotalDynamicOffsetCount);
+
+    // Mark push constants as dirty when PSO changes, since the new PSO might have
+    // a different pipeline layout. If the new PSO has push constants, they need to
+    // be re-committed even if the data hasn't changed.
+    if (Layout.HasPushConstants())
+    {
+        m_State.PushConstantsDirty = true;
+        // Ensure we have initialized push constants data for the new PSO
+        // (user should call SetPushConstants before draw/dispatch)
+    }
 }
 
 DeviceContextVkImpl::ResourceBindInfo& DeviceContextVkImpl::GetBindInfo(PIPELINE_TYPE Type)
@@ -433,6 +443,54 @@ void DeviceContextVkImpl::UpdateInlineConstantBuffers(ResourceBindInfo& BindInfo
     }
 }
 
+void DeviceContextVkImpl::SetPushConstants(const void* pData, Uint32 Offset, Uint32 Size)
+{
+    DEV_CHECK_ERR(pData != nullptr || Size == 0, "pData must not be null when Size is non-zero");
+    DEV_CHECK_ERR(Offset + Size <= MAX_PUSH_CONSTANTS_SIZE,
+                  "Push constant data range [", Offset, ", ", Offset + Size, ") exceeds the maximum supported size (", MAX_PUSH_CONSTANTS_SIZE, ")");
+
+    // Note: This function may be called from UpdateInlineConstantBuffers during draw preparation,
+    // at which point the pipeline state is guaranteed to be bound. When called from user code,
+    // we validate that the pipeline has push constants.
+    if (m_pPipelineState != nullptr)
+    {
+        const PipelineLayoutVk& Layout = m_pPipelineState->GetPipelineLayout();
+        if (Layout.HasPushConstants())
+        {
+            DEV_CHECK_ERR(Offset + Size <= Layout.GetPushConstantSize(),
+                          "Push constant data range [", Offset, ", ", Offset + Size, ") exceeds the push constant block size (", Layout.GetPushConstantSize(), ")");
+        }
+    }
+
+    if (Size > 0)
+    {
+        memcpy(m_PushConstantsData.data() + Offset, pData, Size);
+        m_PushConstantsDataSize = std::max(m_PushConstantsDataSize, Offset + Size);
+        m_State.PushConstantsDirty = true;
+    }
+}
+
+void DeviceContextVkImpl::CommitPushConstants()
+{
+    if (!m_State.PushConstantsDirty)
+        return;
+
+    VERIFY_EXPR(m_pPipelineState != nullptr);
+    const PipelineLayoutVk& Layout = m_pPipelineState->GetPipelineLayout();
+
+    if (!Layout.HasPushConstants())
+        return;
+
+    const Uint32             Size       = Layout.GetPushConstantSize();
+    const VkShaderStageFlags StageFlags = Layout.GetPushConstantStageFlags();
+    const VkPipelineLayout   vkLayout   = Layout.GetVkPipelineLayout();
+
+    VERIFY_EXPR(Size <= m_PushConstantsData.size());
+
+    m_CommandBuffer.PushConstants(vkLayout, StageFlags, 0, Size, m_PushConstantsData.data());
+    m_State.PushConstantsDirty = false;
+}
+
 void DeviceContextVkImpl::CommitDescriptorSets(ResourceBindInfo& BindInfo, Uint32 CommitSRBMask)
 {
     VERIFY(CommitSRBMask != 0, "This method should not be called when there is nothing to commit");
@@ -791,6 +849,10 @@ void DeviceContextVkImpl::PrepareForDraw(DRAW_FLAGS Flags)
     {
         CommitDescriptorSets(BindInfo, CommitMask);
     }
+
+    // Commit push constants if dirty
+    CommitPushConstants();
+
 #ifdef DILIGENT_DEVELOPMENT
     // Must be called after CommitDescriptorSets as it needs SetInfo.BaseInd
     DvpValidateCommittedShaderResources(BindInfo);
@@ -1094,6 +1156,9 @@ void DeviceContextVkImpl::PrepareForDispatchCompute()
         CommitDescriptorSets(BindInfo, CommitMask);
     }
 
+    // Commit push constants if dirty
+    CommitPushConstants();
+
 #ifdef DILIGENT_DEVELOPMENT
     // Must be called after CommitDescriptorSets as it needs SetInfo.BaseInd
     DvpValidateCommittedShaderResources(BindInfo);
@@ -1114,6 +1179,9 @@ void DeviceContextVkImpl::PrepareForRayTracing()
         CommitDescriptorSets(BindInfo, CommitMask);
     }
 
+    // Commit push constants if dirty
+    CommitPushConstants();
+
 #ifdef DILIGENT_DEVELOPMENT
     // Must be called after CommitDescriptorSets as it needs SetInfo.BaseInd
     DvpValidateCommittedShaderResources(BindInfo);