D3D12 backend: reworked immutable sampler initialization

3 files changed, 89 insertions, 94 deletions

diff --git a/Graphics/GraphicsEngineD3D12/include/RootSignature.hpp b/Graphics/GraphicsEngineD3D12/include/RootSignature.hpp
index e29371a6..0ded449c 100644
--- a/Graphics/GraphicsEngineD3D12/include/RootSignature.hpp
+++ b/Graphics/GraphicsEngineD3D12/include/RootSignature.hpp
@@ -479,9 +479,7 @@ private:
 class RootSignatureBuilder
 {
 public:
-    explicit RootSignatureBuilder(RootSignature& RootSig);
-
-    void AllocateImmutableSamplers(const PipelineResourceLayoutDesc& ResourceLayout);
+    RootSignatureBuilder(RootSignature& RootSig, const PipelineResourceLayoutDesc& PipelineResLayout);
 
     void InitImmutableSampler(SHADER_TYPE                     ShaderType,
                               const char*                     SamplerName,
@@ -512,19 +510,28 @@ public:
 
     struct ImmutableSamplerAttribs
     {
-        ImmutableSamplerDesc    SamplerDesc;
-        UINT                    ShaderRegister   = ~UINT{0};
-        UINT                    ArraySize        = 0;
-        UINT                    RegisterSpace    = 0;
-        D3D12_SHADER_VISIBILITY ShaderVisibility = static_cast<D3D12_SHADER_VISIBILITY>(-1);
-        String                  Name;
-        SHADER_TYPE             ShaderType = SHADER_TYPE_UNKNOWN;
-
-        ImmutableSamplerAttribs() noexcept {}
-        ImmutableSamplerAttribs(const ImmutableSamplerDesc& SamDesc, D3D12_SHADER_VISIBILITY Visibility, SHADER_TYPE Stage) noexcept :
-            SamplerDesc(SamDesc),
-            ShaderVisibility(Visibility),
-            ShaderType{Stage}
+        const SamplerDesc& Desc;
+        const String       SamplerName;
+        const UINT         ShaderRegister;
+        const UINT         ArraySize;
+        const UINT         RegisterSpace;
+        const SHADER_TYPE  ShaderType;
+
+        ImmutableSamplerAttribs(
+            const SamplerDesc& _Desc,
+            const char*        _SamplerName,
+            const UINT         _ShaderRegister,
+            const UINT         _ArraySize,
+            const UINT         _RegisterSpace,
+            const SHADER_TYPE  _ShaderType) noexcept :
+            // clang-format off
+            Desc          {_Desc},
+            SamplerName   {_SamplerName},
+            ShaderRegister{_ShaderRegister},
+            ArraySize     {_ArraySize},
+            RegisterSpace {_RegisterSpace},
+            ShaderType    {_ShaderType}
+        // clang-format on
         {}
     };
     const ImmutableSamplerAttribs* GetImmutableSamplers() const { return m_ImmutableSamplers.data(); }
@@ -537,6 +544,8 @@ private:
 
     RootSignature& m_RootSig;
 
+    const PipelineResourceLayoutDesc& m_PipelineResLayout;
+
     // Resource counters for every descriptor range type used to assign bind points
     // for ray-tracing shaders.
     std::array<Uint32, D3D12_DESCRIPTOR_RANGE_TYPE_SAMPLER + 1> m_NumResources = {};
diff --git a/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp b/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
index cee5c851..8124ea22 100644
--- a/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
+++ b/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
@@ -319,19 +319,16 @@ TBindingMapPerStage ExtractResourceBindingMap(const RootSignatureBuilder&
         if (LayoutIdx < 0)
             continue;
 
-        if (ImtblSmplr.ShaderRegister == ~UINT{0})
-            continue; // Sampler has not been initialized
-
-        if (ImtblSmplr.Name.empty())
+        if (ImtblSmplr.SamplerName.empty())
         {
             UNEXPECTED("Immutable sampler name is empty");
             continue;
         }
 
         auto& BindingMap = BindingMapPerStage[ShaderIdx];
-        auto  Iter       = BindingMap.emplace(HashMapStringKey{ImtblSmplr.Name.c_str()}, ImtblSmplr.ShaderRegister).first;
+        auto  Iter       = BindingMap.emplace(HashMapStringKey{ImtblSmplr.SamplerName.c_str()}, ImtblSmplr.ShaderRegister).first;
         VERIFY(Iter->second == ImtblSmplr.ShaderRegister,
-               "Static sampler '", ImtblSmplr.Name, "' was assigned incosistent bind points in different resource layouts. This is a bug.");
+               "Static sampler '", ImtblSmplr.SamplerName, "' was assigned incosistent bind points in different resource layouts. This is a bug.");
     }
 
     return BindingMapPerStage;
@@ -413,8 +410,6 @@ void PipelineStateD3D12Impl::InitInternalObjects(const PSOCreateInfoType& Create
 
     InitializePipelineDesc(CreateInfo, MemPool);
 
-    RootSigBuilder.AllocateImmutableSamplers(CreateInfo.PSODesc.ResourceLayout);
-
     // It is important to construct all objects before initializing them because if an exception is thrown,
     // destructors will be called for all objects
 
@@ -430,7 +425,7 @@ PipelineStateD3D12Impl::PipelineStateD3D12Impl(IReferenceCounters*
 {
     try
     {
-        RootSignatureBuilder RootSigBuilder{m_RootSig};
+        RootSignatureBuilder RootSigBuilder{m_RootSig, CreateInfo.PSODesc.ResourceLayout};
         TShaderStages        ShaderStages;
         InitInternalObjects(CreateInfo, RootSigBuilder, ShaderStages);
 
@@ -636,7 +631,7 @@ PipelineStateD3D12Impl::PipelineStateD3D12Impl(IReferenceCounters*
 {
     try
     {
-        RootSignatureBuilder RootSigBuilder{m_RootSig};
+        RootSignatureBuilder RootSigBuilder{m_RootSig, CreateInfo.PSODesc.ResourceLayout};
         TShaderStages        ShaderStages;
         InitInternalObjects(CreateInfo, RootSigBuilder, ShaderStages);
 
@@ -693,7 +688,7 @@ PipelineStateD3D12Impl::PipelineStateD3D12Impl(IReferenceCounters*
     {
         LocalRootSignature   LocalRootSig{CreateInfo.pShaderRecordName, CreateInfo.RayTracingPipeline.ShaderRecordSize};
         TShaderStages        ShaderStages;
-        RootSignatureBuilder RootSigBuilder{m_RootSig};
+        RootSignatureBuilder RootSigBuilder{m_RootSig, CreateInfo.PSODesc.ResourceLayout};
         InitInternalObjects(CreateInfo, RootSigBuilder, ShaderStages, &LocalRootSig);
 
         auto* pd3d12Device = pDeviceD3D12->GetD3D12Device5();
diff --git a/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp b/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp
index bd8f86d7..e86c4405 100644
--- a/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp
+++ b/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp
@@ -173,8 +173,10 @@ size_t RootSignature::RootParamsManager::GetHash() const
 }
 
 
-RootSignatureBuilder::RootSignatureBuilder(RootSignature& RootSig) :
+RootSignatureBuilder::RootSignatureBuilder(RootSignature&                    RootSig,
+                                           const PipelineResourceLayoutDesc& PipelineResLayout) :
     m_RootSig{RootSig},
+    m_PipelineResLayout{PipelineResLayout},
     m_ImmutableSamplers(STD_ALLOCATOR_RAW_MEM(ImmutableSamplerAttribs, GetRawAllocator(), "Allocator for vector<ImmutableSamplerAttribs>"))
 {
 }
@@ -185,50 +187,64 @@ void RootSignatureBuilder::InitImmutableSampler(SHADER_TYPE
                                                 const char*                     SamplerSuffix,
                                                 const D3DShaderResourceAttribs& SamplerAttribs)
 {
-    auto ShaderVisibility = ShaderTypeToD3D12ShaderVisibility(ShaderType);
-    auto SamplerFound     = false;
-    for (auto& ImtblSmplr : m_ImmutableSamplers)
-    {
-        if (ImtblSmplr.ShaderVisibility == ShaderVisibility &&
-            StreqSuff(SamplerName, ImtblSmplr.SamplerDesc.SamplerOrTextureName, SamplerSuffix))
+    auto FindExistingSampler = [&]() -> const ImmutableSamplerAttribs* {
+        for (const auto& ImtblSmplr : m_ImmutableSamplers)
         {
-            if (ImtblSmplr.ShaderRegister == ~UINT{0})
+            if (ImtblSmplr.ShaderType == ShaderType && ImtblSmplr.SamplerName == 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;
-                }
+                return &ImtblSmplr;
             }
-            else
+        }
+        return nullptr;
+    };
+
+    if (ShaderType & RAY_TRACING_SHADER_TYPES)
+    {
+        if (const auto* pImtblSmplr = FindExistingSampler())
+        {
+            // The sampler has already been initialized when processing another
+            // shader from the same shader stage.
+            DEV_CHECK_ERR(pImtblSmplr->ArraySize == SamplerAttribs.BindCount, "Immutable sampler '", SamplerName,
+                          "' has already been found, but its previous array size (", pImtblSmplr->ArraySize,
+                          ") does not match the new array size (", SamplerAttribs.BindCount, ")");
+            return;
+        }
+    }
+    else
+    {
+        VERIFY(FindExistingSampler() == nullptr, "Multiple samplers with the same name in one stage are only expected in ray tracing shaders");
+    }
+
+    bool SamplerFound = false;
+    for (Uint32 SrcImtblSmplrId = 0; SrcImtblSmplrId < m_PipelineResLayout.NumImmutableSamplers; ++SrcImtblSmplrId)
+    {
+        const auto& ImtblSamDesc = m_PipelineResLayout.ImmutableSamplers[SrcImtblSmplrId];
+        if ((ImtblSamDesc.ShaderStages & ShaderType) != 0 &&
+            StreqSuff(SamplerName, ImtblSamDesc.SamplerOrTextureName, SamplerSuffix))
+        {
+            UINT ShaderRegister = SamplerAttribs.BindPoint;
+            if (ShaderType & RAY_TRACING_SHADER_TYPES)
             {
-                // Do not allocate another register if the sampler has already been processed
-                VERIFY(ShaderType & RAY_TRACING_SHADER_TYPES, "Multiple samplers with the same name can only be encountered in ray tracing shaders");
-                DEV_CHECK_ERR(ImtblSmplr.ArraySize == SamplerAttribs.BindCount, "Immutable sampler '", SamplerName,
-                              "' has already been found, but its previous array size (", ImtblSmplr.ArraySize,
-                              ") does not match the new array size (", SamplerAttribs.BindCount, ")");
+                // For ray tracing shaders, use the next available sampler register.
+                // The bindings will be remapped in the DXIL byte code.
+                ShaderRegister = m_NumResources[D3D12_DESCRIPTOR_RANGE_TYPE_SAMPLER];
+                m_NumResources[D3D12_DESCRIPTOR_RANGE_TYPE_SAMPLER] += SamplerAttribs.BindCount;
             }
 
+            m_ImmutableSamplers.emplace_back(
+                ImtblSamDesc.Desc,
+                SamplerName,
+                ShaderRegister,
+                SamplerAttribs.BindCount,
+                0, // Register space
+                ShaderType);
+
             SamplerFound = true;
             break;
         }
     }
 
-    if (!SamplerFound)
-    {
-        LOG_ERROR("Unable to find immutable sampler \'", SamplerName, '\'');
-    }
+    DEV_CHECK_ERR(SamplerFound, "Unable to find immutable sampler \'", SamplerName, '\'');
 }
 
 
@@ -319,27 +335,6 @@ void RootSignatureBuilder::AllocateResourceSlot(SHADER_TYPE
     }
 }
 
-
-void RootSignatureBuilder::AllocateImmutableSamplers(const PipelineResourceLayoutDesc& ResourceLayout)
-{
-    if (ResourceLayout.NumImmutableSamplers > 0)
-    {
-        m_ImmutableSamplers.reserve(ResourceLayout.NumImmutableSamplers);
-        for (Uint32 sam = 0; sam < ResourceLayout.NumImmutableSamplers; ++sam)
-        {
-            const auto& ImtblSamDesc = ResourceLayout.ImmutableSamplers[sam];
-            SHADER_TYPE ShaderStages = ImtblSamDesc.ShaderStages;
-            while (ShaderStages != 0)
-            {
-                auto Stage = ShaderStages & ~static_cast<SHADER_TYPE>(ShaderStages - 1);
-                m_ImmutableSamplers.emplace_back(ImtblSamDesc, ShaderTypeToD3D12ShaderVisibility(Stage), Stage);
-                ShaderStages &= ~Stage;
-            }
-        }
-    }
-}
-
-
 void RootSignatureBuilder::Finalize(ID3D12Device* pd3d12Device)
 {
     auto& RootParams          = m_RootSig.m_RootParams;
@@ -399,8 +394,7 @@ void RootSignatureBuilder::Finalize(ID3D12Device* pd3d12Device)
     UINT TotalD3D12StaticSamplers = 0;
     for (const auto& ImtblSam : m_ImmutableSamplers)
     {
-        if (ImtblSam.ShaderRegister != ~UINT{0})
-            TotalD3D12StaticSamplers += ImtblSam.ArraySize;
+        TotalD3D12StaticSamplers += ImtblSam.ArraySize;
     }
     rootSignatureDesc.NumStaticSamplers = TotalD3D12StaticSamplers;
     rootSignatureDesc.pStaticSamplers   = nullptr;
@@ -410,12 +404,9 @@ void RootSignatureBuilder::Finalize(ID3D12Device* pd3d12Device)
     {
         for (size_t s = 0; s < m_ImmutableSamplers.size(); ++s)
         {
-            const auto& ImtblSmplrDesc = m_ImmutableSamplers[s];
-            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)
+            const auto& ImtblSmplr = m_ImmutableSamplers[s];
+            const auto& SamDesc    = ImtblSmplr.Desc;
+            for (UINT ArrInd = 0; ArrInd < ImtblSmplr.ArraySize; ++ArrInd)
             {
                 D3D12StaticSamplers.emplace_back(
                     D3D12_STATIC_SAMPLER_DESC //
@@ -430,10 +421,10 @@ void RootSignatureBuilder::Finalize(ID3D12Device* pd3d12Device)
                         BorderColorToD3D12StaticBorderColor(SamDesc.BorderColor),
                         SamDesc.MinLOD,
                         SamDesc.MaxLOD,
-                        ImtblSmplrDesc.ShaderRegister + ArrInd,
-                        ImtblSmplrDesc.RegisterSpace,
-                        ImtblSmplrDesc.ShaderVisibility //
-                    }                                   //
+                        ImtblSmplr.ShaderRegister + ArrInd,
+                        ImtblSmplr.RegisterSpace,
+                        ShaderTypeToD3D12ShaderVisibility(ImtblSmplr.ShaderType) //
+                    }                                                            //
                 );
             }
         }