remap resource binding for DirectX 12

5 files changed, 158 insertions, 26 deletions

diff --git a/Graphics/GraphicsEngineD3D12/include/RootSignature.hpp b/Graphics/GraphicsEngineD3D12/include/RootSignature.hpp
index d74a7b0c..fe78a852 100644
--- a/Graphics/GraphicsEngineD3D12/include/RootSignature.hpp
+++ b/Graphics/GraphicsEngineD3D12/include/RootSignature.hpp
@@ -321,6 +321,7 @@ public:
                               const D3DShaderResourceAttribs& ShaderResAttribs,
                               SHADER_RESOURCE_VARIABLE_TYPE   VariableType,
                               D3D12_DESCRIPTOR_RANGE_TYPE     RangeType,
+                              Uint32&                         BindPoint,
                               Uint32&                         RootIndex,
                               Uint32&                         OffsetFromTableStart);
 
@@ -373,6 +374,27 @@ public:
         return m_RootParams.GetHash();
     }
 
+    // Note: sizeof(m_ImmutableSamplers) == 56 (MS compiler, release x64)
+    struct ImmutableSamplerAttribs
+    {
+        ImmutableSamplerDesc    SamplerDesc;
+        UINT                    ShaderRegister   = static_cast<UINT>(-1);
+        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 ImmutableSamplerAttribs* GetImmutableSamplers() const { return m_ImmutableSamplers.data(); }
+    size_t                         GetImmutableSamplerCount() const { return m_ImmutableSamplers.size(); }
+
 private:
 #ifdef DILIGENT_DEBUG
     void dbgVerifyRootParameters() const;
@@ -478,23 +500,10 @@ private:
     // This array contains the same data for Sampler root table
     std::array<Uint8, SHADER_RESOURCE_VARIABLE_TYPE_NUM_TYPES* MAX_SHADERS_IN_PIPELINE> m_SamplerRootTablesMap = {};
 
-    RootParamsManager m_RootParams;
+    std::array<Uint16, D3D12_DESCRIPTOR_RANGE_TYPE_SAMPLER + 1> m_NumResources = {};
 
-    struct ImmutableSamplerAttribs
-    {
-        ImmutableSamplerDesc    SamplerDesc;
-        UINT                    ShaderRegister   = static_cast<UINT>(-1);
-        UINT                    ArraySize        = 0;
-        UINT                    RegisterSpace    = 0;
-        D3D12_SHADER_VISIBILITY ShaderVisibility = static_cast<D3D12_SHADER_VISIBILITY>(-1);
+    RootParamsManager m_RootParams;
 
-        ImmutableSamplerAttribs() noexcept {}
-        ImmutableSamplerAttribs(const ImmutableSamplerDesc& SamDesc, D3D12_SHADER_VISIBILITY Visibility) noexcept :
-            SamplerDesc(SamDesc),
-            ShaderVisibility(Visibility)
-        {}
-    };
-    // Note: sizeof(m_ImmutableSamplers) == 56 (MS compiler, release x64)
     std::vector<ImmutableSamplerAttribs, STDAllocatorRawMem<ImmutableSamplerAttribs>> m_ImmutableSamplers;
 
     IMemoryAllocator& m_MemAllocator;
diff --git a/Graphics/GraphicsEngineD3D12/include/ShaderResourceLayoutD3D12.hpp b/Graphics/GraphicsEngineD3D12/include/ShaderResourceLayoutD3D12.hpp
index 58a8141d..6f71f8f8 100644
--- a/Graphics/GraphicsEngineD3D12/include/ShaderResourceLayoutD3D12.hpp
+++ b/Graphics/GraphicsEngineD3D12/include/ShaderResourceLayoutD3D12.hpp
@@ -190,6 +190,7 @@ public:
                       Uint32                           _SamplerId,
                       SHADER_RESOURCE_VARIABLE_TYPE    _VariableType,
                       CachedResourceType               _ResType,
+                      Uint32                           _BindPoint,
                       Uint32                           _RootIndex,
                       Uint32                           _OffsetFromTableStart) noexcept :
             // clang-format off
@@ -198,7 +199,8 @@ public:
             {
                 _StringPool,
                 _Attribs,
-                _SamplerId
+                _SamplerId,
+                _BindPoint
             },
             ResourceType        {static_cast<Uint32>(_ResType)     },
             VariableType        {static_cast<Uint32>(_VariableType)},
diff --git a/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp b/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
index 84453dfa..714643c3 100644
--- a/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
+++ b/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
@@ -38,6 +38,8 @@
 #include "StringTools.hpp"
 #include "ShaderVariableD3D12.hpp"
 #include "DynamicLinearAllocator.hpp"
+#include "DXCompiler.hpp"
+#include "dxc/dxcapi.h"
 
 namespace Diligent
 {
@@ -282,6 +284,90 @@ void GetShaderIdentifiers(ID3D12DeviceChild*                       pSO,
     }
 }
 
+void RemapResourceBinding(IDXCompiler*                                     compiler,
+                          const RootSignature&                             RootSig,
+                          const std::array<Int8, MAX_SHADERS_IN_PIPELINE>& ResourceLayoutIndex,
+                          const ShaderResourceLayoutD3D12*                 pResourceLayouts,
+                          const ShaderResourceLayoutD3D12*                 pStaticLayouts,
+                          Uint32                                           NumStages,
+                          std::vector<D3D12_STATE_SUBOBJECT>&              Subobjects,
+                          std::vector<CComPtr<IDxcBlob>>&                  ShaderBlobs) noexcept(false)
+{
+    IDXCompiler::TBindingMapPerStage BindingMapPerStage;
+    String                           EntryPoint;
+
+    const auto ExtractResources = [&](const ShaderResourceLayoutD3D12* pLayouts) //
+    {
+        for (Uint32 ShaderIdx = 0; ShaderIdx < ResourceLayoutIndex.size(); ++ShaderIdx)
+        {
+            const Int8 LayoutIdx = ResourceLayoutIndex[ShaderIdx];
+            if (LayoutIdx < 0)
+                continue;
+
+            auto&       BindingMap = BindingMapPerStage[ShaderIdx];
+            const auto& ResLayout  = pLayouts[LayoutIdx];
+            for (Uint32 v = 0; v < SHADER_RESOURCE_VARIABLE_TYPE_NUM_TYPES; ++v)
+            {
+                auto   VarType   = static_cast<SHADER_RESOURCE_VARIABLE_TYPE>(v);
+                Uint32 ResCount  = ResLayout.GetCbvSrvUavCount(VarType);
+                Uint32 SampCount = ResLayout.GetSamplerCount(VarType);
+
+                for (Uint32 i = 0; i < ResCount; ++i)
+                {
+                    const auto& Attribs = ResLayout.GetSrvCbvUav(VarType, i).Attribs;
+                    auto        Iter    = BindingMap.emplace(HashMapStringKey{Attribs.Name}, Attribs.BindPoint).first;
+                    VERIFY_EXPR(Iter->second == Attribs.BindPoint);
+                }
+                for (Uint32 i = 0; i < SampCount; ++i)
+                {
+                    const auto& Attribs = ResLayout.GetSampler(VarType, i).Attribs;
+                    auto        Iter    = BindingMap.emplace(HashMapStringKey{Attribs.Name}, Attribs.BindPoint).first;
+                    VERIFY_EXPR(Iter->second == Attribs.BindPoint);
+                }
+            }
+        }
+    };
+    ExtractResources(pResourceLayouts);
+    ExtractResources(pStaticLayouts);
+
+    for (size_t i = 0; i < RootSig.GetImmutableSamplerCount(); ++i)
+    {
+        const auto&  ImtblSmplr = RootSig.GetImmutableSamplers()[i];
+        const Uint32 ShaderIdx  = GetShaderTypePipelineIndex(ImtblSmplr.ShaderType, PIPELINE_TYPE_RAY_TRACING);
+        const Int8   LayoutIdx  = ResourceLayoutIndex[ShaderIdx];
+        if (LayoutIdx < 0)
+            continue;
+
+        auto& BindingMap = BindingMapPerStage[ShaderIdx];
+        BindingMap.emplace(HashMapStringKey{ImtblSmplr.Name.c_str()}, ImtblSmplr.ShaderRegister);
+    }
+
+    for (auto& SubObj : Subobjects)
+    {
+        if (SubObj.Type == D3D12_STATE_SUBOBJECT_TYPE_DXIL_LIBRARY)
+        {
+            auto& DxilLib = *reinterpret_cast<D3D12_DXIL_LIBRARY_DESC*>(const_cast<void*>(SubObj.pDesc));
+            VERIFY_EXPR(DxilLib.NumExports == 1);
+
+            const auto& Export = *DxilLib.pExports;
+            EntryPoint.resize(wcslen(Export.ExportToRename));
+            for (size_t i = 0; i < EntryPoint.size(); ++i)
+                EntryPoint[i] = static_cast<char>(Export.ExportToRename[i]);
+
+            CComPtr<IDxcBlob> pBlob;
+            compiler->RemapResourceBinding(BindingMapPerStage, EntryPoint.c_str(), DxilLib.DXILLibrary.pShaderBytecode, DxilLib.DXILLibrary.BytecodeLength, &pBlob);
+
+            if (!pBlob)
+                LOG_ERROR_AND_THROW("Failed to remap resource bindings");
+
+            DxilLib.DXILLibrary.pShaderBytecode = pBlob->GetBufferPointer();
+            DxilLib.DXILLibrary.BytecodeLength  = pBlob->GetBufferSize();
+
+            ShaderBlobs.push_back(pBlob);
+        }
+    }
+}
+
 } // namespace
 
 
@@ -650,6 +736,11 @@ PipelineStateD3D12Impl::PipelineStateD3D12Impl(IReferenceCounters*
         if (LocalRoot.pLocalRootSignature)
             Subobjects.push_back({D3D12_STATE_SUBOBJECT_TYPE_LOCAL_ROOT_SIGNATURE, &LocalRoot});
 
+        std::vector<CComPtr<IDxcBlob>> ShaderBlobs;
+        RemapResourceBinding(pDeviceD3D12->GetDxCompiler(), m_RootSig, m_ResourceLayoutIndex,
+                             &m_pShaderResourceLayouts[0], &m_pShaderResourceLayouts[GetNumShaderStages()], GetNumShaderStages(),
+                             Subobjects, ShaderBlobs);
+
         D3D12_STATE_OBJECT_DESC RTPipelineDesc = {};
         RTPipelineDesc.Type                    = D3D12_STATE_OBJECT_TYPE_RAYTRACING_PIPELINE;
         RTPipelineDesc.NumSubobjects           = static_cast<UINT>(Subobjects.size());
diff --git a/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp b/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp
index 2ac15372..cbe9d7bd 100644
--- a/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp
+++ b/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp
@@ -39,6 +39,7 @@
 
 namespace Diligent
 {
+static constexpr auto RayTracingMask = SHADER_TYPE_RAY_GEN | SHADER_TYPE_RAY_MISS | SHADER_TYPE_RAY_CLOSEST_HIT | SHADER_TYPE_RAY_ANY_HIT | SHADER_TYPE_RAY_INTERSECTION | SHADER_TYPE_CALLABLE;
 
 RootSignature::RootParamsManager::RootParamsManager(IMemoryAllocator& MemAllocator) :
     m_MemAllocator{MemAllocator},
@@ -218,6 +219,13 @@ void RootSignature::InitImmutableSampler(SHADER_TYPE                     ShaderT
             ImtblSmplr.ShaderRegister = SamplerAttribs.BindPoint;
             ImtblSmplr.ArraySize      = SamplerAttribs.BindCount;
             ImtblSmplr.RegisterSpace  = 0;
+            ImtblSmplr.Name           = SamplerName;
+
+            if (ShaderType & RayTracingMask)
+            {
+                ImtblSmplr.ShaderRegister = m_NumResources[D3D12_DESCRIPTOR_RANGE_TYPE_SAMPLER];
+                m_NumResources[D3D12_DESCRIPTOR_RANGE_TYPE_SAMPLER] += SamplerAttribs.BindCount;
+            }
 
             SamplerFound = true;
             break;
@@ -236,11 +244,22 @@ void RootSignature::AllocateResourceSlot(SHADER_TYPE                     ShaderT
                                          const D3DShaderResourceAttribs& ShaderResAttribs,
                                          SHADER_RESOURCE_VARIABLE_TYPE   VariableType,
                                          D3D12_DESCRIPTOR_RANGE_TYPE     RangeType,
+                                         Uint32&                         BindPoint,           // in/out parameter
                                          Uint32&                         RootIndex,           // Output parameter
                                          Uint32&                         OffsetFromTableStart // Output parameter
 )
 {
     const auto ShaderVisibility = ShaderTypeToD3D12ShaderVisibility(ShaderType);
+
+    // update resource binding for ray tracing
+    if (ShaderType & RayTracingMask)
+    {
+        BindPoint = m_NumResources[RangeType];
+        m_NumResources[RangeType] += ShaderResAttribs.BindCount;
+    }
+    else
+        BindPoint = ShaderResAttribs.BindPoint;
+
     if (RangeType == D3D12_DESCRIPTOR_RANGE_TYPE_CBV && ShaderResAttribs.BindCount == 1)
     {
         // Allocate single CBV directly in the root signature
@@ -250,7 +269,7 @@ void RootSignature::AllocateResourceSlot(SHADER_TYPE                     ShaderT
         OffsetFromTableStart = 0;
 
         // Add new root view to existing root parameters
-        m_RootParams.AddRootView(D3D12_ROOT_PARAMETER_TYPE_CBV, RootIndex, ShaderResAttribs.BindPoint, ShaderVisibility, VariableType);
+        m_RootParams.AddRootView(D3D12_ROOT_PARAMETER_TYPE_CBV, RootIndex, BindPoint, ShaderVisibility, VariableType);
     }
     else
     {
@@ -293,7 +312,7 @@ void RootSignature::AllocateResourceSlot(SHADER_TYPE                     ShaderT
         Uint32 NewDescriptorRangeIndex = d3d12RootParam.DescriptorTable.NumDescriptorRanges - 1;
         CurrParam.SetDescriptorRange(NewDescriptorRangeIndex,
                                      RangeType,                  // Range type (CBV, SRV, UAV or SAMPLER)
-                                     ShaderResAttribs.BindPoint, // Shader register
+                                     BindPoint,                  // Shader register
                                      ShaderResAttribs.BindCount, // Number of registers used (1 for non-array resources)
                                      0,                          // Register space. Always 0 for now
                                      OffsetFromTableStart        // Offset in descriptors from the table start
@@ -376,11 +395,11 @@ void RootSignature::AllocateImmutableSamplers(const PipelineResourceLayoutDesc&
         for (Uint32 sam = 0; sam < ResourceLayout.NumImmutableSamplers; ++sam)
         {
             const auto& ImtblSamDesc = ResourceLayout.ImmutableSamplers[sam];
-            Uint32      ShaderStages = ImtblSamDesc.ShaderStages;
+            SHADER_TYPE ShaderStages = ImtblSamDesc.ShaderStages;
             while (ShaderStages != 0)
             {
-                auto Stage = ShaderStages & ~(ShaderStages - 1);
-                m_ImmutableSamplers.emplace_back(ImtblSamDesc, ShaderTypeToD3D12ShaderVisibility(static_cast<SHADER_TYPE>(Stage)));
+                auto Stage = ShaderStages & ~static_cast<SHADER_TYPE>(ShaderStages - 1);
+                m_ImmutableSamplers.emplace_back(ImtblSamDesc, ShaderTypeToD3D12ShaderVisibility(Stage), Stage);
                 ShaderStages &= ~Stage;
             }
         }
@@ -475,8 +494,8 @@ void RootSignature::Finalize(ID3D12Device* pd3d12Device)
         rootSignatureDesc.pStaticSamplers = D3D12StaticSamplers.data();
 
         // Release immutable samplers array, we no longer need it
-        std::vector<ImmutableSamplerAttribs, STDAllocatorRawMem<ImmutableSamplerAttribs>> EmptySamplers(STD_ALLOCATOR_RAW_MEM(ImmutableSamplerAttribs, GetRawAllocator(), "Allocator for vector<ImmutableSamplerAttribs>"));
-        m_ImmutableSamplers.swap(EmptySamplers);
+        //std::vector<ImmutableSamplerAttribs, STDAllocatorRawMem<ImmutableSamplerAttribs>> EmptySamplers(STD_ALLOCATOR_RAW_MEM(ImmutableSamplerAttribs, GetRawAllocator(), "Allocator for vector<ImmutableSamplerAttribs>"));
+        //m_ImmutableSamplers.swap(EmptySamplers);
 
         VERIFY_EXPR(D3D12StaticSamplers.size() == TotalD3D12StaticSamplers);
     }
diff --git a/Graphics/GraphicsEngineD3D12/src/ShaderResourceLayoutD3D12.cpp b/Graphics/GraphicsEngineD3D12/src/ShaderResourceLayoutD3D12.cpp
index 5923797c..30665698 100644
--- a/Graphics/GraphicsEngineD3D12/src/ShaderResourceLayoutD3D12.cpp
+++ b/Graphics/GraphicsEngineD3D12/src/ShaderResourceLayoutD3D12.cpp
@@ -91,7 +91,9 @@ StringPool ShaderResourceLayoutD3D12::AllocateMemory(IMemoryAllocator&
                                                      size_t                                                             StringPoolSize)
 {
     m_CbvSrvUavOffsets[0] = 0;
-    for (SHADER_RESOURCE_VARIABLE_TYPE VarType = SHADER_RESOURCE_VARIABLE_TYPE_STATIC; VarType < SHADER_RESOURCE_VARIABLE_TYPE_NUM_TYPES; VarType = static_cast<SHADER_RESOURCE_VARIABLE_TYPE>(VarType + 1))
+    for (SHADER_RESOURCE_VARIABLE_TYPE VarType = SHADER_RESOURCE_VARIABLE_TYPE_STATIC;
+         VarType < SHADER_RESOURCE_VARIABLE_TYPE_NUM_TYPES;
+         VarType = static_cast<SHADER_RESOURCE_VARIABLE_TYPE>(VarType + 1))
     {
         VERIFY(m_CbvSrvUavOffsets[VarType] + CbvSrvUavCount[VarType] <= std::numeric_limits<Uint16>::max(), "Offset is not representable in 16 bits");
         m_CbvSrvUavOffsets[VarType + 1] = static_cast<Uint16>(m_CbvSrvUavOffsets[VarType] + CbvSrvUavCount[VarType]);
@@ -99,7 +101,9 @@ StringPool ShaderResourceLayoutD3D12::AllocateMemory(IMemoryAllocator&
     }
 
     m_SamplersOffsets[0] = m_CbvSrvUavOffsets[SHADER_RESOURCE_VARIABLE_TYPE_NUM_TYPES];
-    for (SHADER_RESOURCE_VARIABLE_TYPE VarType = SHADER_RESOURCE_VARIABLE_TYPE_STATIC; VarType < SHADER_RESOURCE_VARIABLE_TYPE_NUM_TYPES; VarType = static_cast<SHADER_RESOURCE_VARIABLE_TYPE>(VarType + 1))
+    for (SHADER_RESOURCE_VARIABLE_TYPE VarType = SHADER_RESOURCE_VARIABLE_TYPE_STATIC;
+         VarType < SHADER_RESOURCE_VARIABLE_TYPE_NUM_TYPES;
+         VarType = static_cast<SHADER_RESOURCE_VARIABLE_TYPE>(VarType + 1))
     {
         VERIFY(m_SamplersOffsets[VarType] + SamplerCount[VarType] <= std::numeric_limits<Uint16>::max(), "Offset is not representable in 16 bits");
         m_SamplersOffsets[VarType + 1] = static_cast<Uint16>(m_SamplersOffsets[VarType] + SamplerCount[VarType]);
@@ -254,13 +258,15 @@ void ShaderResourceLayoutD3D12::Initialize(ID3D12Device*
         {
             Uint32 RootIndex = D3D12Resource::InvalidRootIndex;
             Uint32 Offset    = D3D12Resource::InvalidOffset;
+            Uint32 BindPoint = D3DShaderResourceAttribs::InvalidBindPoint;
 
             D3D12_DESCRIPTOR_RANGE_TYPE DescriptorRangeType = GetDescriptorRangeType(ResType);
 
             if (pRootSig)
             {
-                pRootSig->AllocateResourceSlot(GetShaderType(), PipelineType, Attribs, VarType, DescriptorRangeType, RootIndex, Offset);
+                pRootSig->AllocateResourceSlot(GetShaderType(), PipelineType, Attribs, VarType, DescriptorRangeType, BindPoint, RootIndex, Offset);
                 VERIFY(RootIndex <= D3D12Resource::MaxRootIndex, "Root index excceeds allowed limit");
+                VERIFY(BindPoint <= D3DShaderResourceAttribs::MaxBindPoint, "Bind point excceeds allowed limit");
             }
             else
             {
@@ -277,6 +283,7 @@ void ShaderResourceLayoutD3D12::Initialize(ID3D12Device*
 
                 RootIndex = DescriptorRangeType;
                 Offset    = Attribs.BindPoint;
+                BindPoint = Attribs.BindPoint;
                 // Resources in the static resource cache are indexed by the bind point
                 StaticResCacheTblSizes[RootIndex] = std::max(StaticResCacheTblSizes[RootIndex], Offset + Attribs.BindCount);
             }
@@ -297,6 +304,7 @@ void ShaderResourceLayoutD3D12::Initialize(ID3D12Device*
                     SamplerId,
                     VarType,
                     ResType,
+                    BindPoint,
                     RootIndex,
                     Offset //
                 };
@@ -307,7 +315,10 @@ void ShaderResourceLayoutD3D12::Initialize(ID3D12Device*
             auto& ExistingRes = GetResource(ResIter->second);
             VERIFY_EXPR(ExistingRes.VariableType == VarType);
             VERIFY_EXPR(ExistingRes.Attribs.GetInputType() == Attribs.GetInputType());
+            VERIFY_EXPR(ExistingRes.Attribs.GetSRVDimension() == Attribs.GetSRVDimension());
             VERIFY_EXPR(ExistingRes.Attribs.BindCount == Attribs.BindCount);
+            VERIFY_EXPR(ExistingRes.Attribs.BindPoint == Attribs.BindPoint ||
+                        Attribs.BindPoint == D3DShaderResourceAttribs::InvalidBindPoint);
         }
     };