Added support for local root signature & shader record.

Bug fix for ray tracing.

7 files changed, 51 insertions, 32 deletions

diff --git a/Graphics/GraphicsEngineVulkan/include/DeviceContextVkImpl.hpp b/Graphics/GraphicsEngineVulkan/include/DeviceContextVkImpl.hpp
index 39dca08d..705057e4 100644
--- a/Graphics/GraphicsEngineVulkan/include/DeviceContextVkImpl.hpp
+++ b/Graphics/GraphicsEngineVulkan/include/DeviceContextVkImpl.hpp
@@ -255,10 +255,10 @@ public:
     virtual void DILIGENT_CALL_TYPE Flush() override final;
 
     /// Implementation of IDeviceContext::BuildBLAS() in Vulkan backend.
-    virtual void DILIGENT_CALL_TYPE BuildBLAS(const BLASBuildAttribs& Attribs) override final;
+    virtual void DILIGENT_CALL_TYPE BuildBLAS(const BuildBLASAttribs& Attribs) override final;
 
     /// Implementation of IDeviceContext::BuildTLAS() in Vulkan backend.
-    virtual void DILIGENT_CALL_TYPE BuildTLAS(const TLASBuildAttribs& Attribs) override final;
+    virtual void DILIGENT_CALL_TYPE BuildTLAS(const BuildTLASAttribs& Attribs) override final;
 
     /// Implementation of IDeviceContext::CopyBLAS() in Vulkan backend.
     virtual void DILIGENT_CALL_TYPE CopyBLAS(const CopyBLASAttribs& Attribs) override final;
diff --git a/Graphics/GraphicsEngineVulkan/src/DeviceContextVkImpl.cpp b/Graphics/GraphicsEngineVulkan/src/DeviceContextVkImpl.cpp
index a76a901e..ce9c639c 100644
--- a/Graphics/GraphicsEngineVulkan/src/DeviceContextVkImpl.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/DeviceContextVkImpl.cpp
@@ -2792,7 +2792,7 @@ void DeviceContextVkImpl::ResolveTextureSubresource(ITexture*
                                  1, &ResolveRegion);
 }
 
-void DeviceContextVkImpl::BuildBLAS(const BLASBuildAttribs& Attribs)
+void DeviceContextVkImpl::BuildBLAS(const BuildBLASAttribs& Attribs)
 {
     if (!TDeviceContextBase::BuildBLAS(Attribs, 0))
         return;
@@ -2945,7 +2945,7 @@ void DeviceContextVkImpl::BuildBLAS(const BLASBuildAttribs& Attribs)
 #endif
 }
 
-void DeviceContextVkImpl::BuildTLAS(const TLASBuildAttribs& Attribs)
+void DeviceContextVkImpl::BuildTLAS(const BuildTLASAttribs& Attribs)
 {
     if (!TDeviceContextBase::BuildTLAS(Attribs, 0))
         return;
@@ -2967,14 +2967,13 @@ void DeviceContextVkImpl::BuildTLAS(const TLASBuildAttribs& Attribs)
 
     // copy instance data into instance buffer
     {
-        size_t Size             = Attribs.InstanceCount * sizeof(VkAccelerationStructureInstanceKHR);
-        auto   TmpSpace         = m_UploadHeap.Allocate(Size, 16);
-        void*  pMappedInstances = TmpSpace.CPUAddress;
+        size_t Size     = Attribs.InstanceCount * sizeof(VkAccelerationStructureInstanceKHR);
+        auto   TmpSpace = m_UploadHeap.Allocate(Size, 16);
 
         for (Uint32 i = 0; i < Attribs.InstanceCount; ++i)
         {
             const auto& Inst     = Attribs.pInstances[i];
-            auto&       vkASInst = static_cast<VkAccelerationStructureInstanceKHR*>(pMappedInstances)[i];
+            auto&       vkASInst = static_cast<VkAccelerationStructureInstanceKHR*>(TmpSpace.CPUAddress)[i];
             auto* const pBLASVk  = ValidatedCast<BottomLevelASVkImpl>(Inst.pBLAS);
 
             static_assert(sizeof(vkASInst.transform) == sizeof(Inst.Transform), "size mismatch");
diff --git a/Graphics/GraphicsEngineVulkan/src/PipelineLayout.cpp b/Graphics/GraphicsEngineVulkan/src/PipelineLayout.cpp
index c24ed496..3840b739 100644
--- a/Graphics/GraphicsEngineVulkan/src/PipelineLayout.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/PipelineLayout.cpp
@@ -45,19 +45,19 @@ class ResourceTypeToVkDescriptorType
 public:
     ResourceTypeToVkDescriptorType()
     {
-        static_assert(SPIRVShaderResourceAttribs::ResourceType::NumResourceTypes == 12, "Please add the corresponding decriptor type");
-        m_Map[SPIRVShaderResourceAttribs::ResourceType::UniformBuffer]         = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC;
-        m_Map[SPIRVShaderResourceAttribs::ResourceType::ROStorageBuffer]       = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC;
-        m_Map[SPIRVShaderResourceAttribs::ResourceType::RWStorageBuffer]       = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC;
-        m_Map[SPIRVShaderResourceAttribs::ResourceType::UniformTexelBuffer]    = VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER;
-        m_Map[SPIRVShaderResourceAttribs::ResourceType::StorageTexelBuffer]    = VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER;
-        m_Map[SPIRVShaderResourceAttribs::ResourceType::StorageImage]          = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE;
-        m_Map[SPIRVShaderResourceAttribs::ResourceType::SampledImage]          = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
-        m_Map[SPIRVShaderResourceAttribs::ResourceType::AtomicCounter]         = VK_DESCRIPTOR_TYPE_MAX_ENUM; // atomic counter doesn't exist in Vulkan
-        m_Map[SPIRVShaderResourceAttribs::ResourceType::SeparateImage]         = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE;
-        m_Map[SPIRVShaderResourceAttribs::ResourceType::SeparateSampler]       = VK_DESCRIPTOR_TYPE_SAMPLER;
-        m_Map[SPIRVShaderResourceAttribs::ResourceType::InputAttachment]       = VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT;
-        m_Map[SPIRVShaderResourceAttribs::ResourceType::AccelerationStructure] = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR;
+        static_assert(Uint32{SPIRVShaderResourceAttribs::ResourceType::NumResourceTypes} == 12, "Please add the corresponding decriptor type");
+        m_Map[Uint32{SPIRVShaderResourceAttribs::ResourceType::UniformBuffer}]         = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC;
+        m_Map[Uint32{SPIRVShaderResourceAttribs::ResourceType::ROStorageBuffer}]       = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC;
+        m_Map[Uint32{SPIRVShaderResourceAttribs::ResourceType::RWStorageBuffer}]       = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC;
+        m_Map[Uint32{SPIRVShaderResourceAttribs::ResourceType::UniformTexelBuffer}]    = VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER;
+        m_Map[Uint32{SPIRVShaderResourceAttribs::ResourceType::StorageTexelBuffer}]    = VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER;
+        m_Map[Uint32{SPIRVShaderResourceAttribs::ResourceType::StorageImage}]          = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE;
+        m_Map[Uint32{SPIRVShaderResourceAttribs::ResourceType::SampledImage}]          = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
+        m_Map[Uint32{SPIRVShaderResourceAttribs::ResourceType::AtomicCounter}]         = VK_DESCRIPTOR_TYPE_MAX_ENUM; // atomic counter doesn't exist in Vulkan
+        m_Map[Uint32{SPIRVShaderResourceAttribs::ResourceType::SeparateImage}]         = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE;
+        m_Map[Uint32{SPIRVShaderResourceAttribs::ResourceType::SeparateSampler}]       = VK_DESCRIPTOR_TYPE_SAMPLER;
+        m_Map[Uint32{SPIRVShaderResourceAttribs::ResourceType::InputAttachment}]       = VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT;
+        m_Map[Uint32{SPIRVShaderResourceAttribs::ResourceType::AccelerationStructure}] = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR;
     }
 
     VkDescriptorType operator[](SPIRVShaderResourceAttribs::ResourceType ResType) const
@@ -66,7 +66,7 @@ public:
     }
 
 private:
-    std::array<VkDescriptorType, SPIRVShaderResourceAttribs::ResourceType::NumResourceTypes> m_Map = {};
+    std::array<VkDescriptorType, Uint32{SPIRVShaderResourceAttribs::ResourceType::NumResourceTypes}> m_Map = {};
 };
 
 VkDescriptorType PipelineLayout::GetVkDescriptorType(SPIRVShaderResourceAttribs::ResourceType Type)
diff --git a/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp b/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
index 1cac2b6f..a9dcc5cb 100644
--- a/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
@@ -761,9 +761,6 @@ PipelineStateVkImpl::PipelineStateVkImpl(IReferenceCounters*
         const auto& LogicalDevice         = GetDevice()->GetLogicalDevice();
         const auto  ShaderGroupHandleSize = pDeviceVk->GetShaderGroupHandleSize();
 
-        if (LogicalDevice.GetEnabledExtFeatures().RayTracing.rayTracing == VK_FALSE)
-            LOG_ERROR_AND_THROW("Ray tracing is not supported by this device");
-
         std::vector<VkPipelineShaderStageCreateInfo>      vkShaderStages;
         std::vector<VulkanUtilities::ShaderModuleWrapper> ShaderModules;
         std::vector<VkRayTracingShaderGroupCreateInfoKHR> ShaderGroups;
diff --git a/Graphics/GraphicsEngineVulkan/src/ShaderResourceCacheVk.cpp b/Graphics/GraphicsEngineVulkan/src/ShaderResourceCacheVk.cpp
index 0054315f..2e1039ba 100644
--- a/Graphics/GraphicsEngineVulkan/src/ShaderResourceCacheVk.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/ShaderResourceCacheVk.cpp
@@ -156,7 +156,7 @@ void ShaderResourceCacheVk::TransitionResources(DeviceContextVkImpl* pCtxVkImpl)
     for (Uint32 res = 0; res < m_TotalResources; ++res)
     {
         auto& Res = pResources[res];
-        static_assert(SPIRVShaderResourceAttribs::ResourceType::NumResourceTypes == 12, "Please handle the new resource type below");
+        static_assert(Uint32{SPIRVShaderResourceAttribs::ResourceType::NumResourceTypes} == 12, "Please handle the new resource type below");
         switch (Res.Type)
         {
             case SPIRVShaderResourceAttribs::ResourceType::UniformBuffer:
diff --git a/Graphics/GraphicsEngineVulkan/src/ShaderResourceLayoutVk.cpp b/Graphics/GraphicsEngineVulkan/src/ShaderResourceLayoutVk.cpp
index cab26f83..7763b480 100644
--- a/Graphics/GraphicsEngineVulkan/src/ShaderResourceLayoutVk.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/ShaderResourceLayoutVk.cpp
@@ -300,7 +300,7 @@ void ShaderResourceLayoutVk::InitializeStaticResourceLayout(const std::vector<co
                         // For immutable separate samplers we allocate VkResource instances, but they are never exposed to the app
                     }
 
-                    Uint32 Binding       = Attribs.Type;
+                    Uint32 Binding       = Uint32{Attribs.Type};
                     Uint32 DescriptorSet = 0;
                     Uint32 CacheOffset   = StaticResCacheSize;
                     StaticResCacheSize += Attribs.ArraySize;
@@ -1161,7 +1161,7 @@ void ShaderResourceLayoutVk::VkResource::BindResource(IDeviceObject* pObj, Uint3
 
     if (pObj)
     {
-        static_assert(SPIRVShaderResourceAttribs::ResourceType::NumResourceTypes == 12, "Please handle the new resource type below");
+        static_assert(Uint32{SPIRVShaderResourceAttribs::ResourceType::NumResourceTypes} == 12, "Please handle the new resource type below");
         switch (Type)
         {
             case SPIRVShaderResourceAttribs::ResourceType::UniformBuffer:
@@ -1265,7 +1265,7 @@ void ShaderResourceLayoutVk::InitializeStaticResources(const ShaderResourceLayou
         // Get resource attributes
         const auto& DstRes = GetResource(SHADER_RESOURCE_VARIABLE_TYPE_STATIC, r);
         const auto& SrcRes = SrcLayout.GetResource(SHADER_RESOURCE_VARIABLE_TYPE_STATIC, r);
-        VERIFY(SrcRes.Binding == SrcRes.Type, "Unexpected binding");
+        VERIFY(SrcRes.Binding == Uint32{SrcRes.Type}, "Unexpected binding");
         VERIFY(SrcRes.ArraySize == DstRes.ArraySize, "Inconsistent array size");
 
         if (DstRes.Type == SPIRVShaderResourceAttribs::ResourceType::SeparateSampler &&
@@ -1412,7 +1412,7 @@ void ShaderResourceLayoutVk::CommitDynamicResources(const ShaderResourceCacheVk&
         WriteDescrSetIt->descriptorType = PipelineLayout::GetVkDescriptorType(Res.Type);
 
         // For every resource type, try to batch as many descriptor updates as we can
-        static_assert(SPIRVShaderResourceAttribs::ResourceType::NumResourceTypes == 12, "Please handle the new resource type below");
+        static_assert(Uint32{SPIRVShaderResourceAttribs::ResourceType::NumResourceTypes} == 12, "Please handle the new resource type below");
         switch (Res.Type)
         {
             case SPIRVShaderResourceAttribs::ResourceType::UniformBuffer:
diff --git a/Graphics/GraphicsEngineVulkan/src/VulkanUtilities/VulkanRayTracingKHRviaNV.cpp b/Graphics/GraphicsEngineVulkan/src/VulkanUtilities/VulkanRayTracingKHRviaNV.cpp
index 5b5434d5..451a7cc2 100644
--- a/Graphics/GraphicsEngineVulkan/src/VulkanUtilities/VulkanRayTracingKHRviaNV.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/VulkanUtilities/VulkanRayTracingKHRviaNV.cpp
@@ -55,14 +55,33 @@ constexpr VkDeviceAddress                     g_BufferMask                 = 0xF
 PFN_vkCreateBuffer              Origin_vkCreateBuffer              = nullptr;
 PFN_vkDestroyBuffer             Origin_vkDestroyBuffer             = nullptr;
 PFN_vkGetBufferDeviceAddressKHR Origin_vkGetBufferDeviceAddressKHR = nullptr;
+PFN_vkAllocateMemory            Origin_vkAllocateMemory            = nullptr;
+
 
 VKAPI_ATTR VkResult VKAPI_CALL Wrap_vkCreateBuffer(VkDevice                     device,
                                                    const VkBufferCreateInfo*    pCreateInfo,
                                                    const VkAllocationCallbacks* pAllocator,
                                                    VkBuffer*                    pBuffer)
 {
-    const_cast<VkBufferCreateInfo*>(pCreateInfo)->usage &= ~VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT;
-    return Origin_vkCreateBuffer(device, pCreateInfo, pAllocator, pBuffer);
+    VkBufferCreateInfo CreateInfo = *pCreateInfo;
+    CreateInfo.usage &= ~VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT;
+    return Origin_vkCreateBuffer(device, &CreateInfo, pAllocator, pBuffer);
+}
+
+VKAPI_ATTR VkResult VKAPI_PTR Wrap_vkAllocateMemory(VkDevice device, const VkMemoryAllocateInfo* pAllocateInfo, const VkAllocationCallbacks* pAllocator, VkDeviceMemory* pMemory)
+{
+    VkMemoryAllocateInfo AllocInfo = *pAllocateInfo;
+
+    for (auto* pNext = static_cast<VkBaseOutStructure*>(const_cast<void*>(AllocInfo.pNext)); pNext;)
+    {
+        // remove VkMemoryAllocateFlagsInfo because VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT is removed from buffer create info.
+        if (pNext->sType == VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_FLAGS_INFO)
+            pNext->pNext = pNext->pNext;
+
+        pNext = pNext->pNext;
+    }
+
+    return Origin_vkAllocateMemory(device, &AllocInfo, pAllocator, pMemory);
 }
 
 VKAPI_ATTR void VKAPI_CALL Wrap_vkDestroyBuffer(VkDevice                     device,
@@ -509,6 +528,8 @@ VKAPI_ATTR VkResult VKAPI_CALL Redirect_vkCreateRayTracingPipelinesKHR(VkDevice
 
 void EnableRayTracingKHRviaNV()
 {
+    LOG_WARNING_MESSAGE("This is fallback implementation, you should use VK_KHR_ray_tracing instead");
+
     vkCreateAccelerationStructureKHR                = &Redirect_vkCreateAccelerationStructureKHR;
     vkGetAccelerationStructureMemoryRequirementsKHR = &Redirect_vkGetAccelerationStructureMemoryRequirementsKHR;
     vkBindAccelerationStructureMemoryKHR            = &Redirect_vkBindAccelerationStructureMemoryKHR;
@@ -523,8 +544,10 @@ void EnableRayTracingKHRviaNV()
     Origin_vkGetBufferDeviceAddressKHR = vkGetBufferDeviceAddressKHR;
     Origin_vkCreateBuffer              = vkCreateBuffer;
     Origin_vkDestroyBuffer             = vkDestroyBuffer;
+    Origin_vkAllocateMemory            = vkAllocateMemory;
     vkCreateBuffer                     = &Wrap_vkCreateBuffer;
     vkDestroyBuffer                    = &Wrap_vkDestroyBuffer;
+    vkAllocateMemory                   = Wrap_vkAllocateMemory;
     vkGetBufferDeviceAddressKHR        = &Wrap_vkGetBufferDeviceAddressKHR;
     vkGetBufferDeviceAddress           = &Wrap_vkGetBufferDeviceAddressKHR;
     vkGetBufferDeviceAddressEXT        = &Wrap_vkGetBufferDeviceAddressKHR;