D3D12 backend: fixed few potential issues with immutable sample intialization in ray tracing pipelines

3 files changed, 32 insertions, 13 deletions

diff --git a/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp b/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
index 93c8e514..cee5c851 100644
--- a/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
+++ b/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
@@ -132,7 +132,7 @@ void BuildRTPipelineDescription(const RayTracingPipelineStateCreateInfo& CreateI
 
             CComPtr<IDxcBlob> pBlob;
             if (!compiler->RemapResourceBindings(BindingMap, reinterpret_cast<IDxcBlob*>(pShaderD3D12->GetShaderByteCode()), &pBlob))
-                LOG_ERROR_AND_THROW("Failed to remap resource bindings");
+                LOG_ERROR_AND_THROW("Failed to remap resource bindings in shader '", pShaderD3D12->GetDesc().Name, "'.");
 
             LibDesc.DXILLibrary.BytecodeLength  = pBlob->GetBufferSize();
             LibDesc.DXILLibrary.pShaderBytecode = pBlob->GetBufferPointer();
@@ -200,7 +200,7 @@ void BuildRTPipelineDescription(const RayTracingPipelineStateCreateInfo& CreateI
     constexpr Uint32 DefaultPayloadSize = sizeof(float) * 8;
 
     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;
     Subobjects.push_back({D3D12_STATE_SUBOBJECT_TYPE_RAYTRACING_PIPELINE_CONFIG, &PipelineConfig});
 
@@ -298,7 +298,7 @@ TBindingMapPerStage ExtractResourceBindingMap(const RootSignatureBuilder&
             for (Uint32 i = 0; i < TotalResCount; ++i)
             {
                 const auto& Attribs = ResLayout.GetResource(i).Attribs;
-                VERIFY_EXPR(Attribs.Name != nullptr && strlen(Attribs.Name) > 0);
+                VERIFY_EXPR(Attribs.Name != nullptr && Attribs.Name[0] != '\0');
 
                 auto Iter = BindingMap.emplace(HashMapStringKey{Attribs.Name}, Attribs.BindPoint).first;
                 VERIFY(Iter->second == Attribs.BindPoint,
@@ -319,6 +319,9 @@ TBindingMapPerStage ExtractResourceBindingMap(const RootSignatureBuilder&
         if (LayoutIdx < 0)
             continue;
 
+        if (ImtblSmplr.ShaderRegister == ~UINT{0})
+            continue; // Sampler has not been initialized
+
         if (ImtblSmplr.Name.empty())
         {
             UNEXPECTED("Immutable sampler name is empty");
@@ -817,7 +820,7 @@ void PipelineStateD3D12Impl::InitResourceLayouts(const PipelineStateCreateInfo&
 
         m_ResourceLayoutIndex[ShaderInd] = static_cast<Int8>(s);
 
-        // Initialize general layout
+        // Initialize all-resources layout
         m_pShaderResourceLayouts[s].Initialize(
             m_Desc.PipelineType,
             ResourceLayout,
diff --git a/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp b/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp
index f6e9f2c6..bd8f86d7 100644
--- a/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp
+++ b/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp
@@ -194,16 +194,22 @@ void RootSignatureBuilder::InitImmutableSampler(SHADER_TYPE
         {
             if (ImtblSmplr.ShaderRegister == ~UINT{0})
             {
-                ImtblSmplr.ShaderRegister = SamplerAttribs.BindPoint;
-                ImtblSmplr.ArraySize      = SamplerAttribs.BindCount;
-                ImtblSmplr.RegisterSpace  = 0;
-                ImtblSmplr.Name           = SamplerName;
+                ImtblSmplr.ArraySize     = SamplerAttribs.BindCount;
+                ImtblSmplr.RegisterSpace = 0;
+                ImtblSmplr.Name          = SamplerName;
 
                 if (ShaderType & RAY_TRACING_SHADER_TYPES)
                 {
+                    // For ray tracing shaders, use the next available sampler register.
+                    // The bindings will be remapped in the DXIL byte code.
                     ImtblSmplr.ShaderRegister = m_NumResources[D3D12_DESCRIPTOR_RANGE_TYPE_SAMPLER];
                     m_NumResources[D3D12_DESCRIPTOR_RANGE_TYPE_SAMPLER] += SamplerAttribs.BindCount;
                 }
+                else
+                {
+                    // Use original bind point.
+                    ImtblSmplr.ShaderRegister = SamplerAttribs.BindPoint;
+                }
             }
             else
             {
@@ -249,6 +255,7 @@ void RootSignatureBuilder::AllocateResourceSlot(SHADER_TYPE
     }
     else
     {
+        // Use original bind point
         BindPoint = ShaderResAttribs.BindPoint;
     }
 
@@ -391,7 +398,10 @@ void RootSignatureBuilder::Finalize(ID3D12Device* pd3d12Device)
 
     UINT TotalD3D12StaticSamplers = 0;
     for (const auto& ImtblSam : m_ImmutableSamplers)
-        TotalD3D12StaticSamplers += ImtblSam.ArraySize;
+    {
+        if (ImtblSam.ShaderRegister != ~UINT{0})
+            TotalD3D12StaticSamplers += ImtblSam.ArraySize;
+    }
     rootSignatureDesc.NumStaticSamplers = TotalD3D12StaticSamplers;
     rootSignatureDesc.pStaticSamplers   = nullptr;
     std::vector<D3D12_STATIC_SAMPLER_DESC, STDAllocatorRawMem<D3D12_STATIC_SAMPLER_DESC>> D3D12StaticSamplers(STD_ALLOCATOR_RAW_MEM(D3D12_STATIC_SAMPLER_DESC, GetRawAllocator(), "Allocator for vector<D3D12_STATIC_SAMPLER_DESC>"));
@@ -401,7 +411,10 @@ void RootSignatureBuilder::Finalize(ID3D12Device* pd3d12Device)
         for (size_t s = 0; s < m_ImmutableSamplers.size(); ++s)
         {
             const auto& ImtblSmplrDesc = m_ImmutableSamplers[s];
-            const auto& SamDesc        = ImtblSmplrDesc.SamplerDesc.Desc;
+            if (ImtblSmplrDesc.ShaderRegister == ~UINT{0})
+                continue; // The sampler has not been initialized
+
+            const auto& SamDesc = ImtblSmplrDesc.SamplerDesc.Desc;
             for (UINT ArrInd = 0; ArrInd < ImtblSmplrDesc.ArraySize; ++ArrInd)
             {
                 D3D12StaticSamplers.emplace_back(
diff --git a/Graphics/GraphicsEngineD3D12/src/ShaderResourceLayoutD3D12.cpp b/Graphics/GraphicsEngineD3D12/src/ShaderResourceLayoutD3D12.cpp
index f81896d7..89f303c6 100644
--- a/Graphics/GraphicsEngineD3D12/src/ShaderResourceLayoutD3D12.cpp
+++ b/Graphics/GraphicsEngineD3D12/src/ShaderResourceLayoutD3D12.cpp
@@ -241,7 +241,10 @@ void ShaderResourceLayoutD3D12::Initialize(PIPELINE_TYPE
             const auto VarType = ShaderRes.FindVariableType(Attribs, ResourceLayout);
 
             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 should never happen as "
+                   "all resources are added to the map when they are being counted.");
 
             if (ResIter->second == InvalidResourceIndex)
             {
@@ -317,8 +320,8 @@ void ShaderResourceLayoutD3D12::Initialize(PIPELINE_TYPE
                 if (ImtblSamplerInd >= 0)
                 {
                     // Note that there may be multiple immutable samplers with the same name in different shaders that
-                    // are assigned to different registers. This is not a problem, because InitImmutableSampler allocates new
-                    // register only first time the sampler is encountered.
+                    // are assigned to different registers. InitImmutableSampler() handles this by allocating new
+                    // register only first time the sampler is encountered. All bindings will be remapped afterwards.
                     RootSgnBldr.InitImmutableSampler(ShaderRes.GetShaderType(), Sam.Name, ShaderRes.GetCombinedSamplerSuffix(), Sam);
                 }
                 else