Some more updates to D3D12 PSO and ShaderResourceLayoutVk

2 files changed, 55 insertions, 71 deletions

diff --git a/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp b/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
index afdeb741..93c8e514 100644
--- a/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
+++ b/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
@@ -27,6 +27,8 @@
 
 #include "pch.h"
 #include <array>
+#include <sstream>
+
 #include "PipelineStateD3D12Impl.hpp"
 #include "ShaderD3D12Impl.hpp"
 #include "D3D12TypeConversions.hpp"
@@ -111,22 +113,10 @@ void BuildRTPipelineDescription(const RayTracingPipelineStateCreateInfo& CreateI
 {
 #define LOG_PSO_ERROR_AND_THROW(...) LOG_ERROR_AND_THROW("Description of ray tracing PSO '", CreateInfo.PSODesc.Name, "' is invalid: ", ##__VA_ARGS__)
 
-    Uint32 UnnamedShaderIndex = 0;
+    Uint32 UnnamedExportIndex = 0;
 
     std::unordered_map<IShader*, LPCWSTR> UniqueShaders;
 
-    const auto ShaderIndexToStr = [&TempPool](Uint32 Index) -> LPCWSTR {
-        const Uint32 Len = sizeof(Index) * 2;
-        auto* const  Dst = TempPool.Allocate<WCHAR>(Len + 1);
-        for (Uint32 i = 0; i < Len; ++i)
-        {
-            Dst[Len - 1 - i] = L"0123456789ABCDEF"[Index & 0x0F];
-            Index >>= 4;
-        }
-        Dst[Len] = 0;
-        return Dst;
-    };
-
     const auto AddDxilLib = [&](IShader* pShader, const char* Name) -> LPCWSTR {
         if (pShader == nullptr)
             return nullptr;
@@ -155,7 +145,11 @@ void BuildRTPipelineDescription(const RayTracingPipelineStateCreateInfo& CreateI
             if (Name != nullptr)
                 ExportDesc.Name = TempPool.CopyWString(Name);
             else
-                ExportDesc.Name = ShaderIndexToStr(++UnnamedShaderIndex);
+            {
+                std::stringstream ss;
+                ss << "__Shader_" << std::setfill('0') << std::setw(4) << UnnamedExportIndex++;
+                ExportDesc.Name = TempPool.CopyWString(ss.str());
+            }
 
             Subobjects.push_back({D3D12_STATE_SUBOBJECT_TYPE_DXIL_LIBRARY, &LibDesc});
             ShaderBlobs.push_back(pBlob);
@@ -207,7 +201,7 @@ void BuildRTPipelineDescription(const RayTracingPipelineStateCreateInfo& CreateI
 
     auto& PipelineConfig = *TempPool.Construct<D3D12_RAYTRACING_PIPELINE_CONFIG>();
     // For compatibility with Vulkan set minimal recursion depth to one, zero means no tracing of rays at all.
-    PipelineConfig.MaxTraceRecursionDepth = CreateInfo.RayTracingPipeline.MaxRecursionDepth + 1;
+    PipelineConfig.MaxTraceRecursionDepth = CreateInfo.RayTracingPipeline.MaxRecursionDepth;
     Subobjects.push_back({D3D12_STATE_SUBOBJECT_TYPE_RAYTRACING_PIPELINE_CONFIG, &PipelineConfig});
 
     auto& ShaderConfig                   = *TempPool.Construct<D3D12_RAYTRACING_SHADER_CONFIG>();
@@ -221,10 +215,9 @@ template <typename TNameToGroupIndexMap>
 void GetShaderIdentifiers(ID3D12DeviceChild*                       pSO,
                           const RayTracingPipelineStateCreateInfo& CreateInfo,
                           const TNameToGroupIndexMap&              NameToGroupIndex,
-                          Uint8*                                   ShaderData)
+                          Uint8*                                   ShaderData,
+                          Uint32                                   ShaderIdentifierSize)
 {
-    const Uint32 ShaderIdentifierSize = D3D12_SHADER_IDENTIFIER_SIZE_IN_BYTES;
-
     CComPtr<ID3D12StateObjectProperties> pStateObjectProperties;
 
     auto hr = pSO->QueryInterface(IID_PPV_ARGS(&pStateObjectProperties));
@@ -236,8 +229,10 @@ void GetShaderIdentifiers(ID3D12DeviceChild*                       pSO,
         const auto& GeneralShader = CreateInfo.pGeneralShaders[i];
 
         auto iter = NameToGroupIndex.find(GeneralShader.Name);
-        if (iter == NameToGroupIndex.end())
-            LOG_ERROR_AND_THROW("Failed to get shader group index for general shader group '", GeneralShader.Name, "'");
+        VERIFY(iter != NameToGroupIndex.end(),
+               "Can't find general shader '", GeneralShader.Name,
+               "'. This looks to be a bug as NameToGroupIndex is initialized by "
+               "CopyRTShaderGroupNames() that processes the same general shaders.");
 
         const auto* ShaderID = pStateObjectProperties->GetShaderIdentifier(WidenString(GeneralShader.Name).c_str());
         if (ShaderID == nullptr)
@@ -245,13 +240,16 @@ void GetShaderIdentifiers(ID3D12DeviceChild*                       pSO,
 
         std::memcpy(&ShaderData[ShaderIdentifierSize * iter->second], ShaderID, ShaderIdentifierSize);
     }
+
     for (Uint32 i = 0; i < CreateInfo.TriangleHitShaderCount; ++i)
     {
         const auto& TriHitShader = CreateInfo.pTriangleHitShaders[i];
 
         auto iter = NameToGroupIndex.find(TriHitShader.Name);
-        if (iter == NameToGroupIndex.end())
-            LOG_ERROR_AND_THROW("Failed to get shader group index for triangle hit group '", TriHitShader.Name, "'");
+        VERIFY(iter != NameToGroupIndex.end(),
+               "Can't find triangle hit group '", TriHitShader.Name,
+               "'. This looks to be a bug as NameToGroupIndex is initialized by "
+               "CopyRTShaderGroupNames() that processes the same hit groups.");
 
         const auto* ShaderID = pStateObjectProperties->GetShaderIdentifier(WidenString(TriHitShader.Name).c_str());
         if (ShaderID == nullptr)
@@ -259,13 +257,16 @@ void GetShaderIdentifiers(ID3D12DeviceChild*                       pSO,
 
         std::memcpy(&ShaderData[ShaderIdentifierSize * iter->second], ShaderID, ShaderIdentifierSize);
     }
+
     for (Uint32 i = 0; i < CreateInfo.ProceduralHitShaderCount; ++i)
     {
         const auto& ProcHitShader = CreateInfo.pProceduralHitShaders[i];
 
         auto iter = NameToGroupIndex.find(ProcHitShader.Name);
-        if (iter == NameToGroupIndex.end())
-            LOG_ERROR_AND_THROW("Failed to get shader group index for procedural hit shader group '", ProcHitShader.Name, "'");
+        VERIFY(iter != NameToGroupIndex.end(),
+               "Can't find procedural hit group '", ProcHitShader.Name,
+               "'. This looks to be a bug as NameToGroupIndex is initialized by "
+               "CopyRTShaderGroupNames() that processes the same hit groups.");
 
         const auto* ShaderID = pStateObjectProperties->GetShaderIdentifier(WidenString(ProcHitShader.Name).c_str());
         if (ShaderID == nullptr)
@@ -723,7 +724,8 @@ PipelineStateD3D12Impl::PipelineStateD3D12Impl(IReferenceCounters*
         if (FAILED(hr))
             LOG_ERROR_AND_THROW("Failed to create ray tracing state object");
 
-        GetShaderIdentifiers(m_pd3d12PSO, CreateInfo, m_pRayTracingPipelineData->NameToGroupIndex, m_pRayTracingPipelineData->ShaderHandles);
+        GetShaderIdentifiers(m_pd3d12PSO, CreateInfo, m_pRayTracingPipelineData->NameToGroupIndex,
+                             m_pRayTracingPipelineData->ShaderHandles, m_pRayTracingPipelineData->ShaderHandleSize);
 
         if (*m_Desc.Name != 0)
         {
diff --git a/Graphics/GraphicsEngineVulkan/src/ShaderResourceLayoutVk.cpp b/Graphics/GraphicsEngineVulkan/src/ShaderResourceLayoutVk.cpp
index 6c82350d..1700c492 100644
--- a/Graphics/GraphicsEngineVulkan/src/ShaderResourceLayoutVk.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/ShaderResourceLayoutVk.cpp
@@ -572,10 +572,6 @@ void ShaderResourceLayoutVk::Initialize(IRenderDevice*                    pRende
     // Mappings from resource name to its index, for every shader stage
     std::array<ResourceNameToIndex_t, MAX_SHADERS_IN_PIPELINE> ResourceNameToIndexArray;
 
-    const SHADER_RESOURCE_VARIABLE_TYPE* AllowedVarTypes = nullptr;
-    const Uint32                         NumAllowedTypes = 0;
-    const Uint32                         AllowedTypeBits = GetAllowedTypeBits(AllowedVarTypes, NumAllowedTypes);
-
     constexpr bool AllocateImmutableSamplers = true;
 
     std::vector<StringPool> stringPools;
@@ -584,8 +580,7 @@ void ShaderResourceLayoutVk::Initialize(IRenderDevice*                    pRende
     {
         stringPools.emplace_back(
             Layouts[s].AllocateMemory(ShaderStages[s].Shaders, LayoutDataAllocator, ResourceLayoutDesc,
-                                      AllowedVarTypes, NumAllowedTypes, ResourceNameToIndexArray[s],
-                                      AllocateImmutableSamplers));
+                                      nullptr, 0, ResourceNameToIndexArray[s], AllocateImmutableSamplers));
     }
 
     // Current resource index, for every variable type in every shader stage
@@ -597,20 +592,21 @@ void ShaderResourceLayoutVk::Initialize(IRenderDevice*                    pRende
     std::unordered_map<Uint32, std::pair<Uint32, Uint32>> dbgBindings_CacheOffsets;
 #endif
 
-    auto AddResource = [&](const Uint32                      ShaderInd,
-                           ShaderResourceLayoutVk&           ResLayout,
+    auto AddResource = [&](const Uint32                      ShaderStageInd,
                            const SPIRVShaderResources&       Resources,
                            const SPIRVShaderResourceAttribs& Attribs,
-                           ResourceNameToIndex_t&            ResourceNameToIndex,
                            std::vector<uint32_t>&            SPIRV) //
     {
-        const auto                          ShaderType = Resources.GetShaderType();
-        const SHADER_RESOURCE_VARIABLE_TYPE VarType    = FindShaderVariableType(ShaderType, Attribs, ResourceLayoutDesc, Resources.GetCombinedSamplerSuffix());
-        if (!IsAllowedType(VarType, AllowedTypeBits))
-            return;
+        auto& ResourceNameToIndex = ResourceNameToIndexArray[ShaderStageInd];
 
         auto ResIter = ResourceNameToIndex.find(HashMapStringKey{Attribs.Name});
-        VERIFY_EXPR(ResIter != ResourceNameToIndex.end());
+        VERIFY(ResIter != ResourceNameToIndex.end(), "Resource '", Attribs.Name,
+               "' is not found in ResourceNameToIndex map. This is a bug as the resource must have been processed by AllocateMemory and added to the map.");
+
+        const auto ShaderType = Resources.GetShaderType();
+        const auto VarType    = FindShaderVariableType(ShaderType, Attribs, ResourceLayoutDesc, Resources.GetCombinedSamplerSuffix());
+
+        auto& ResLayout = Layouts[ShaderStageInd];
 
         const VkResource* pResource = nullptr;
         if (ResIter->second == InvalidResourceIndex)
@@ -625,7 +621,7 @@ void ShaderResourceLayoutVk::Initialize(IRenderDevice*                    pRende
             {
                 // Separate samplers are enumerated before separate images, so the sampler
                 // assigned to this separate image must have already been created.
-                SamplerInd = FindAssignedSampler(ResLayout, Resources, Attribs, CurrResInd[ShaderInd][VarType], VarType);
+                SamplerInd = FindAssignedSampler(ResLayout, Resources, Attribs, CurrResInd[ShaderStageInd][VarType], VarType);
             }
 
             VkSampler vkImmutableSampler = VK_NULL_HANDLE;
@@ -640,7 +636,7 @@ void ShaderResourceLayoutVk::Initialize(IRenderDevice*                    pRende
 
                     // We reserve enough space for the maximum number of immutable samplers that may be used in the stage,
                     // but not all of them will necessarily be initialized.
-                    auto& ImmutableSampler = ResLayout.GetImmutableSampler(CurrImmutableSamplerInd[ShaderInd]++);
+                    auto& ImmutableSampler = ResLayout.GetImmutableSampler(CurrImmutableSamplerInd[ShaderStageInd]++);
                     VERIFY(!ImmutableSampler, "Immutable sampler has already been initialized. This is unexpected "
                                               "as all resources are deduplicated and should only be initialized once.");
                     const auto& ImmutableSamplerDesc = ResourceLayoutDesc.ImmutableSamplers[SrcImmutableSamplerInd].Desc;
@@ -664,13 +660,13 @@ void ShaderResourceLayoutVk::Initialize(IRenderDevice*                    pRende
             dbgBindings_CacheOffsets[DescriptorSet] = std::make_pair(Binding, CacheOffset);
 #endif
 
-            auto& ResInd    = CurrResInd[ShaderInd][VarType];
+            auto& ResInd    = CurrResInd[ShaderStageInd][VarType];
             ResIter->second = ResInd;
 
             pResource = ::new (&ResLayout.GetResource(VarType, ResInd++)) VkResource //
                 {
                     ResLayout,
-                    stringPools[ShaderInd].CopyString(Attribs.Name),
+                    stringPools[ShaderStageInd].CopyString(Attribs.Name),
                     Attribs.ArraySize,
                     Attribs.Type,
                     Attribs.GetResourceDimension(),
@@ -699,21 +695,15 @@ void ShaderResourceLayoutVk::Initialize(IRenderDevice*                    pRende
     // First process uniform buffers for ALL shader stages to make sure all UBs go first in every descriptor set
     for (size_t s = 0; s < ShaderStages.size(); ++s)
     {
-        auto& Shaders   = ShaderStages[s].Shaders;
-        auto& Layout    = Layouts[s];
-        auto& NameToIdx = ResourceNameToIndexArray[s];
+        auto& Shaders = ShaderStages[s].Shaders;
         for (size_t i = 0; i < Shaders.size(); ++i)
         {
             auto& SPIRV     = ShaderStages[s].SPIRVs[i];
             auto& Resources = *Shaders[i]->GetShaderResources();
             for (Uint32 n = 0; n < Resources.GetNumUBs(); ++n)
             {
-                const auto& UB      = Resources.GetUB(n);
-                auto        VarType = GetShaderVariableType(Resources.GetShaderType(), UB.Name, ResourceLayoutDesc);
-                if (IsAllowedType(VarType, AllowedTypeBits))
-                {
-                    AddResource(static_cast<Uint32>(s), Layout, Resources, UB, NameToIdx, SPIRV);
-                }
+                const auto& UB = Resources.GetUB(n);
+                AddResource(static_cast<Uint32>(s), Resources, UB, SPIRV);
             }
         }
     }
@@ -721,21 +711,15 @@ void ShaderResourceLayoutVk::Initialize(IRenderDevice*                    pRende
     // Second, process all storage buffers in all shader stages
     for (size_t s = 0; s < ShaderStages.size(); ++s)
     {
-        auto& Shaders   = ShaderStages[s].Shaders;
-        auto& Layout    = Layouts[s];
-        auto& NameToIdx = ResourceNameToIndexArray[s];
+        auto& Shaders = ShaderStages[s].Shaders;
         for (size_t i = 0; i < Shaders.size(); ++i)
         {
             auto& Resources = *Shaders[i]->GetShaderResources();
             auto& SPIRV     = ShaderStages[s].SPIRVs[i];
             for (Uint32 n = 0; n < Resources.GetNumSBs(); ++n)
             {
-                const auto& SB      = Resources.GetSB(n);
-                auto        VarType = GetShaderVariableType(Resources.GetShaderType(), SB.Name, ResourceLayoutDesc);
-                if (IsAllowedType(VarType, AllowedTypeBits))
-                {
-                    AddResource(static_cast<Uint32>(s), Layout, Resources, SB, NameToIdx, SPIRV);
-                }
+                const auto& SB = Resources.GetSB(n);
+                AddResource(static_cast<Uint32>(s), Resources, SB, SPIRV);
             }
         }
     }
@@ -743,9 +727,7 @@ void ShaderResourceLayoutVk::Initialize(IRenderDevice*                    pRende
     // Finally, process all other resource types
     for (size_t s = 0; s < ShaderStages.size(); ++s)
     {
-        auto& Layout    = Layouts[s];
-        auto& Shaders   = ShaderStages[s].Shaders;
-        auto& NameToIdx = ResourceNameToIndexArray[s];
+        auto& Shaders = ShaderStages[s].Shaders;
         for (size_t i = 0; i < Shaders.size(); ++i)
         {
             auto& Resources = *Shaders[i]->GetShaderResources();
@@ -765,37 +747,37 @@ void ShaderResourceLayoutVk::Initialize(IRenderDevice*                    pRende
                 [&](const SPIRVShaderResourceAttribs& Img, Uint32)
                 {
                     VERIFY_EXPR(Img.Type == SPIRVShaderResourceAttribs::ResourceType::StorageImage || Img.Type == SPIRVShaderResourceAttribs::ResourceType::StorageTexelBuffer);
-                    AddResource(static_cast<Uint32>(s), Layout, Resources, Img, NameToIdx, SPIRV);
+                    AddResource(static_cast<Uint32>(s), Resources, Img, SPIRV);
                 },
                 [&](const SPIRVShaderResourceAttribs& SmplImg, Uint32)
                 {
                     VERIFY_EXPR(SmplImg.Type == SPIRVShaderResourceAttribs::ResourceType::SampledImage || SmplImg.Type == SPIRVShaderResourceAttribs::ResourceType::UniformTexelBuffer);
-                    AddResource(static_cast<Uint32>(s), Layout, Resources, SmplImg, NameToIdx, SPIRV);
+                    AddResource(static_cast<Uint32>(s), Resources, SmplImg, SPIRV);
                 },
                 [&](const SPIRVShaderResourceAttribs& AC, Uint32)
                 {
                     VERIFY_EXPR(AC.Type == SPIRVShaderResourceAttribs::ResourceType::AtomicCounter);
-                    AddResource(static_cast<Uint32>(s), Layout, Resources, AC, NameToIdx, SPIRV);
+                    AddResource(static_cast<Uint32>(s), Resources, AC, SPIRV);
                 },
                 [&](const SPIRVShaderResourceAttribs& SepSmpl, Uint32)
                 {
                     VERIFY_EXPR(SepSmpl.Type == SPIRVShaderResourceAttribs::ResourceType::SeparateSampler);
-                    AddResource(static_cast<Uint32>(s), Layout, Resources, SepSmpl, NameToIdx, SPIRV);
+                    AddResource(static_cast<Uint32>(s), Resources, SepSmpl, SPIRV);
                 },
                 [&](const SPIRVShaderResourceAttribs& SepImg, Uint32)
                 {
                     VERIFY_EXPR(SepImg.Type == SPIRVShaderResourceAttribs::ResourceType::SeparateImage || SepImg.Type == SPIRVShaderResourceAttribs::ResourceType::UniformTexelBuffer);
-                    AddResource(static_cast<Uint32>(s), Layout, Resources, SepImg, NameToIdx, SPIRV);
+                    AddResource(static_cast<Uint32>(s), Resources, SepImg, SPIRV);
                 },
                 [&](const SPIRVShaderResourceAttribs& InputAtt, Uint32)
                 {
                     VERIFY_EXPR(InputAtt.Type == SPIRVShaderResourceAttribs::ResourceType::InputAttachment);
-                    AddResource(static_cast<Uint32>(s), Layout, Resources, InputAtt, NameToIdx, SPIRV);
+                    AddResource(static_cast<Uint32>(s), Resources, InputAtt, SPIRV);
                 },
                 [&](const SPIRVShaderResourceAttribs& AccelStruct, Uint32)
                 {
                     VERIFY_EXPR(AccelStruct.Type == SPIRVShaderResourceAttribs::ResourceType::AccelerationStructure);
-                    AddResource(static_cast<Uint32>(s), Layout, Resources, AccelStruct, NameToIdx, SPIRV);
+                    AddResource(static_cast<Uint32>(s), Resources, AccelStruct, SPIRV);
                 }
             );
             // clang-format on