remap resource binding for DirectX 12

9 files changed, 441 insertions, 70 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);
         }
     };
 
diff --git a/Graphics/GraphicsEngineD3DBase/include/D3DShaderResourceLoader.hpp b/Graphics/GraphicsEngineD3DBase/include/D3DShaderResourceLoader.hpp
index b12ef6ae..5d5f847e 100644
--- a/Graphics/GraphicsEngineD3DBase/include/D3DShaderResourceLoader.hpp
+++ b/Graphics/GraphicsEngineD3DBase/include/D3DShaderResourceLoader.hpp
@@ -94,6 +94,9 @@ void LoadD3DShaderResources(TShaderReflection*  pShaderReflection,
         D3D_SHADER_INPUT_BIND_DESC BindingDesc = {};
         pShaderReflection->GetResourceBindingDesc(Res, &BindingDesc);
 
+        if (BindingDesc.BindPoint == UINT32_MAX)
+            BindingDesc.BindPoint = D3DShaderResourceAttribs::InvalidBindPoint;
+
         std::string Name(BindingDesc.Name);
 
         SkipCount = 1;
diff --git a/Graphics/GraphicsEngineD3DBase/include/ShaderResources.hpp b/Graphics/GraphicsEngineD3DBase/include/ShaderResources.hpp
index 926ce209..c5607ce3 100644
--- a/Graphics/GraphicsEngineD3DBase/include/ShaderResources.hpp
+++ b/Graphics/GraphicsEngineD3DBase/include/ShaderResources.hpp
@@ -146,20 +146,20 @@ public:
 #endif
     }
 
-    D3DShaderResourceAttribs(StringPool& NamesPool, const D3DShaderResourceAttribs& rhs, Uint32 SamplerId) noexcept :
+    D3DShaderResourceAttribs(StringPool& NamesPool, const D3DShaderResourceAttribs& rhs, Uint32 _SamplerId, Uint32 _BindPoint) noexcept :
         // clang-format off
         D3DShaderResourceAttribs
         {
             NamesPool.CopyString(rhs.Name),
-            rhs.BindPoint,
+            _BindPoint,
             rhs.BindCount,
             rhs.GetInputType(),
             rhs.GetSRVDimension(),
-            SamplerId
+            _SamplerId
         }
     // clang-format on
     {
-        VERIFY(SamplerId == InvalidSamplerId || (GetInputType() == D3D_SIT_TEXTURE && GetSRVDimension() != D3D_SRV_DIMENSION_BUFFER),
+        VERIFY(_SamplerId == InvalidSamplerId || (GetInputType() == D3D_SIT_TEXTURE && GetSRVDimension() != D3D_SRV_DIMENSION_BUFFER),
                "Only texture SRV can be assigned a valid texture sampler");
     }
 
@@ -552,7 +552,7 @@ void ShaderResources::Initialize(TShaderReflection*  pShaderReflection,
             VERIFY(CurrSampler == GetNumSamplers(), "All samplers must be initialized before texture SRVs");
 
             auto  SamplerId  = CombinedSamplerSuffix != nullptr ? FindAssignedSamplerId(TexAttribs, CombinedSamplerSuffix) : D3DShaderResourceAttribs::InvalidSamplerId;
-            auto* pNewTexSRV = new (&GetTexSRV(CurrTexSRV)) D3DShaderResourceAttribs{ResourceNamesPool, TexAttribs, SamplerId};
+            auto* pNewTexSRV = new (&GetTexSRV(CurrTexSRV)) D3DShaderResourceAttribs{ResourceNamesPool, TexAttribs, SamplerId, TexAttribs.BindPoint};
             if (SamplerId != D3DShaderResourceAttribs::InvalidSamplerId)
             {
                 GetSampler(SamplerId).SetTexSRVId(CurrTexSRV);
diff --git a/Graphics/ShaderTools/include/DXCompiler.hpp b/Graphics/ShaderTools/include/DXCompiler.hpp
index 6345ad7a..b996bdff 100644
--- a/Graphics/ShaderTools/include/DXCompiler.hpp
+++ b/Graphics/ShaderTools/include/DXCompiler.hpp
@@ -27,10 +27,14 @@
 
 #pragma once
 
+#include <array>
 #include <vector>
+#include <unordered_map>
 
+#include "Constants.h"
 #include "Shader.h"
 #include "DataBlob.h"
+#include "HashUtils.hpp"
 
 // defined in dxcapi.h
 struct DxcDefine;
@@ -80,6 +84,15 @@ public:
                          std::vector<uint32_t>*  pByteCode,
                          IDataBlob**             ppCompilerOutput) noexcept(false) = 0;
 
+    using TResourceBindingMap = std::unordered_map<HashMapStringKey, Uint32, HashMapStringKey::Hasher>;
+    using TBindingMapPerStage = std::array<TResourceBindingMap, MAX_SHADERS_IN_PIPELINE>;
+
+    virtual bool RemapResourceBinding(const TBindingMapPerStage& BindingMapPerStage,
+                                      const char*                EntryPoint,
+                                      const void*                pBytecode,
+                                      size_t                     BytecodeSize,
+                                      IDxcBlob**                 ppByteCodeBlob) = 0;
+
     // Attempts to extract shader reflection from the bytecode using DXC.
     virtual void GetD3D12ShaderReflection(IDxcBlob*                pShaderBytecode,
                                           ID3D12ShaderReflection** ppShaderReflection) = 0;
diff --git a/Graphics/ShaderTools/src/DXCompiler.cpp b/Graphics/ShaderTools/src/DXCompiler.cpp
index a92ac4c3..3d5bedb9 100644
--- a/Graphics/ShaderTools/src/DXCompiler.cpp
+++ b/Graphics/ShaderTools/src/DXCompiler.cpp
@@ -42,6 +42,8 @@
 #include "DataBlobImpl.hpp"
 #include "RefCntAutoPtr.hpp"
 #include "ShaderToolsCommon.hpp"
+#include "PlatformMisc.hpp"
+#include "GraphicsAccessories.hpp"
 
 #if D3D12_SUPPORTED
 #    include <d3d12shader.h>
@@ -92,6 +94,12 @@ public:
     virtual void GetD3D12ShaderReflection(IDxcBlob*                pShaderBytecode,
                                           ID3D12ShaderReflection** ppShaderReflection) override final;
 
+    virtual bool RemapResourceBinding(const TBindingMapPerStage& BindingMapPerStage,
+                                      const char*                EntryPoint,
+                                      const void*                pBytecode,
+                                      size_t                     BytecodeSize,
+                                      IDxcBlob**                 ppByteCodeBlob) override final;
+
 private:
     DxcCreateInstanceProc Load()
     {
@@ -133,6 +141,10 @@ private:
         return m_pCreateInstance;
     }
 
+    bool        ValidateAndSign(DxcCreateInstanceProc CreateInstance, IDxcLibrary* library, CComPtr<IDxcBlob>& compiled, IDxcBlob** ppBlobOut) const;
+    bool        PatchDXIL(const TResourceBindingMap& ResourceMap, String& DXIL) const;
+    SHADER_TYPE GetEntryShaderType(const String& EntryPoint, const String& DXIL) const;
+
 private:
     DxcCreateInstanceProc  m_pCreateInstance = nullptr;
     bool                   m_IsInitialized   = false;
@@ -327,55 +339,60 @@ bool DXCompilerImpl::Compile(const CompileAttribs& Attribs)
     // validate and sign
     if (m_Target == DXCompilerTarget::Direct3D12)
     {
-        CComPtr<IDxcValidator> validator;
-        hr = CreateInstance(CLSID_DxcValidator, IID_PPV_ARGS(&validator));
-        if (FAILED(hr))
-        {
-            LOG_ERROR("Failed to create  DXC Validator");
-            return false;
-        }
-
-        CComPtr<IDxcOperationResult> validationResult;
-        hr = validator->Validate(compiled, DxcValidatorFlags_InPlaceEdit, &validationResult);
+        return ValidateAndSign(CreateInstance, library, compiled, Attribs.ppBlobOut);
+    }
 
-        if (validationResult == nullptr || FAILED(hr))
-        {
-            LOG_ERROR("Failed to validate shader bytecode");
-            return false;
-        }
+    *Attribs.ppBlobOut = compiled.Detach();
+    return true;
+}
 
-        HRESULT status = E_FAIL;
-        validationResult->GetStatus(&status);
+bool DXCompilerImpl::ValidateAndSign(DxcCreateInstanceProc CreateInstance, IDxcLibrary* library, CComPtr<IDxcBlob>& compiled, IDxcBlob** ppBlobOut) const
+{
+    HRESULT                hr;
+    CComPtr<IDxcValidator> validator;
+    hr = CreateInstance(CLSID_DxcValidator, IID_PPV_ARGS(&validator));
+    if (FAILED(hr))
+    {
+        LOG_ERROR("Failed to create  DXC Validator");
+        return false;
+    }
 
-        if (SUCCEEDED(status))
-        {
-            CComPtr<IDxcBlob> validated;
-            hr = validationResult->GetResult(&validated);
-            if (FAILED(hr))
-                return false;
+    CComPtr<IDxcOperationResult> validationResult;
+    hr = validator->Validate(compiled, DxcValidatorFlags_InPlaceEdit, &validationResult);
 
-            *Attribs.ppBlobOut = validated ? validated.Detach() : compiled.Detach();
-            return true;
-        }
-        else
-        {
-            CComPtr<IDxcBlobEncoding> validationOutput;
-            CComPtr<IDxcBlobEncoding> validationOutputUtf8;
-            validationResult->GetErrorBuffer(&validationOutput);
-            library->GetBlobAsUtf8(validationOutput, &validationOutputUtf8);
+    if (validationResult == nullptr || FAILED(hr))
+    {
+        LOG_ERROR("Failed to validate shader bytecode");
+        return false;
+    }
 
-            size_t      ValidationMsgLen = validationOutputUtf8 ? validationOutputUtf8->GetBufferSize() : 0;
-            const char* ValidationMsg    = ValidationMsgLen > 0 ? static_cast<const char*>(validationOutputUtf8->GetBufferPointer()) : "";
+    HRESULT status = E_FAIL;
+    validationResult->GetStatus(&status);
 
-            LOG_ERROR("Shader validation failed: ", ValidationMsg);
+    if (SUCCEEDED(status))
+    {
+        CComPtr<IDxcBlob> validated;
+        hr = validationResult->GetResult(&validated);
+        if (FAILED(hr))
             return false;
-        }
+
+        *ppBlobOut = validated ? validated.Detach() : compiled.Detach();
+        return true;
     }
+    else
+    {
+        CComPtr<IDxcBlobEncoding> validationOutput;
+        CComPtr<IDxcBlobEncoding> validationOutputUtf8;
+        validationResult->GetErrorBuffer(&validationOutput);
+        library->GetBlobAsUtf8(validationOutput, &validationOutputUtf8);
 
-    *Attribs.ppBlobOut = compiled.Detach();
-    return true;
-}
+        size_t      ValidationMsgLen = validationOutputUtf8 ? validationOutputUtf8->GetBufferSize() : 0;
+        const char* ValidationMsg    = ValidationMsgLen > 0 ? static_cast<const char*>(validationOutputUtf8->GetBufferPointer()) : "";
 
+        LOG_ERROR("Shader validation failed: ", ValidationMsg);
+        return false;
+    }
+}
 
 #if D3D12_SUPPORTED
 class ShaderReflectionViaLibraryReflection final : public ID3D12ShaderReflection
@@ -748,4 +765,210 @@ void DXCompilerImpl::Compile(const ShaderCreateInfo& ShaderCI,
     }
 }
 
+bool DXCompilerImpl::RemapResourceBinding(const TBindingMapPerStage& BindingMapPerStage,
+                                          const char*                EntryPoint,
+                                          const void*                pBytecode,
+                                          size_t                     BytecodeSize,
+                                          IDxcBlob**                 ppByteCodeBlob)
+{
+    auto CreateInstance = GetCreateInstaceProc();
+
+    if (CreateInstance == nullptr)
+    {
+        LOG_ERROR("Failed to load DXCompiler");
+        return false;
+    }
+
+    HRESULT              hr;
+    CComPtr<IDxcLibrary> library;
+    hr = CreateInstance(CLSID_DxcLibrary, IID_PPV_ARGS(&library));
+    if (FAILED(hr))
+    {
+        LOG_ERROR("Failed to create DXC Library");
+        return false;
+    }
+
+    CComPtr<IDxcAssembler> assembler;
+    hr = CreateInstance(CLSID_DxcAssembler, IID_PPV_ARGS(&assembler));
+    if (FAILED(hr))
+    {
+        LOG_ERROR("Failed to create DXC assembler");
+        return false;
+    }
+
+    CComPtr<IDxcCompiler> compiler;
+    hr = CreateInstance(CLSID_DxcCompiler, IID_PPV_ARGS(&compiler));
+    if (FAILED(hr))
+    {
+        LOG_ERROR("Failed to create DXC Compiler");
+        return false;
+    }
+
+    CComPtr<IDxcBlobEncoding> srcBytecode;
+    hr = library->CreateBlobWithEncodingFromPinned(pBytecode, static_cast<Uint32>(BytecodeSize), 0, &srcBytecode);
+    if (FAILED(hr))
+    {
+        LOG_ERROR("Failed to create bytecode blob");
+        return false;
+    }
+
+    CComPtr<IDxcBlobEncoding> disasm;
+    hr = compiler->Disassemble(srcBytecode, &disasm);
+    if (FAILED(hr))
+    {
+        LOG_ERROR("Failed to disassemble bytecode");
+        return false;
+    }
+
+    String dxilAsm;
+    dxilAsm.assign(static_cast<const char*>(disasm->GetBufferPointer()), disasm->GetBufferSize());
+
+    SHADER_TYPE  shaderType  = GetEntryShaderType(EntryPoint, dxilAsm);
+    const Uint32 shaderIndex = GetShaderTypePipelineIndex(shaderType, PIPELINE_TYPE_RAY_TRACING);
+    const auto&  ResourceMap = BindingMapPerStage[shaderIndex];
+
+    if (!PatchDXIL(ResourceMap, dxilAsm))
+    {
+        LOG_ERROR("Failed to patch resource bindings");
+        return false;
+    }
+
+    CComPtr<IDxcBlobEncoding> patchedDisasm;
+    hr = library->CreateBlobWithEncodingFromPinned(dxilAsm.data(), static_cast<UINT32>(dxilAsm.size()), 0, &patchedDisasm);
+    if (FAILED(hr))
+    {
+        LOG_ERROR("Failed to create disassemble blob");
+        return false;
+    }
+
+    CComPtr<IDxcOperationResult> dxilResult;
+    hr = assembler->AssembleToContainer(patchedDisasm, &dxilResult);
+    if (FAILED(hr) || dxilResult == nullptr)
+    {
+        LOG_ERROR("Failed to create DXIL container");
+        return false;
+    }
+
+    HRESULT status = E_FAIL;
+    dxilResult->GetStatus(&status);
+
+    if (FAILED(status))
+    {
+        CComPtr<IDxcBlobEncoding> errorsBlob;
+        CComPtr<IDxcBlobEncoding> errorsBlobUtf8;
+        if (SUCCEEDED(dxilResult->GetErrorBuffer(&errorsBlob)) && SUCCEEDED(library->GetBlobAsUtf8(errorsBlob, &errorsBlobUtf8)))
+        {
+            String errorLog;
+            errorLog.assign(static_cast<const char*>(errorsBlobUtf8->GetBufferPointer()), errorsBlobUtf8->GetBufferSize());
+            LOG_ERROR_MESSAGE("Compilation message: ", errorLog);
+        }
+        else
+            LOG_ERROR("Failed to compile patched asm");
+
+        return false;
+    }
+
+    CComPtr<IDxcBlob> compiled;
+    hr = dxilResult->GetResult(static_cast<IDxcBlob**>(&compiled));
+    if (FAILED(hr))
+        return false;
+
+    return ValidateAndSign(CreateInstance, library, compiled, ppByteCodeBlob);
+}
+
+bool DXCompilerImpl::PatchDXIL(const TResourceBindingMap& ResourceMap, String& DXIL) const
+{
+    String     ResName;
+    char       BindPointStr[256];
+    const char Zero[] = "0";
+
+    for (auto& ResPair : ResourceMap)
+    {
+        // [res name], i32 [space], i32 [bind point]
+
+        const auto& Name      = ResPair.first;
+        const auto& BindPoint = ResPair.second;
+
+        ResName = String{"!\""} + Name.GetStr() + "\", ";
+
+        size_t pos = DXIL.find(ResName);
+        if (pos == String::npos)
+            continue;
+
+        Uint32 Part      = 0;
+        size_t PartStart = pos + ResName.length();
+
+        for (size_t i = PartStart; i < DXIL.size(); ++i)
+        {
+            const char c = DXIL[i];
+            if (c == ' ')
+            {
+                const char* str = &DXIL[PartStart];
+
+                if (Part == 0 || Part == 2)
+                {
+                    VERIFY_EXPR(std::memcmp(str, "i32", i - PartStart) == 0);
+                }
+                else if (Part == 1) // space
+                {
+                    DXIL.replace(PartStart, i - PartStart - 1, Zero);
+                    i = PartStart + strlen(Zero) + 1;
+                }
+                else if (Part == 3) // bind point
+                {
+                    _itoa_s(BindPoint, BindPointStr, 10);
+                    DXIL.replace(PartStart, i - PartStart - 1, BindPointStr);
+                    i = PartStart + strlen(BindPointStr) + 1;
+                }
+                else
+                    break;
+
+                PartStart = i + 1;
+                ++Part;
+            }
+        }
+    }
+    return true;
+}
+
+template <Uint32 S>
+bool ReverseCmp(const char* lhsRev, const char (&rhs)[S])
+{
+    const Uint32 count = S - 1;
+    const char*  lhs   = lhsRev - count;
+    return std::memcmp(lhs, rhs, count) == 0;
+}
+
+SHADER_TYPE DXCompilerImpl::GetEntryShaderType(const String& EntryPoint, const String& DXIL) const
+{
+    const String Pattern              = "void " + EntryPoint + "(";
+    const char   ShaderTypeStart[]    = "[shader(\\22";
+    const char   ShaderTypeEnd[]      = "\\22)]";
+    const char   RayGenShader[]       = "raygeneration";
+    const char   MissShader[]         = "miss";
+    const char   AnyHitShader[]       = "anyhit";
+    const char   ClosestHitShader[]   = "closesthit";
+    const char   IntersectionShader[] = "intersection";
+    const char   CallableShader[]     = "callable";
+
+    size_t pos = DXIL.find(Pattern);
+    if (pos == String::npos)
+        return SHADER_TYPE_UNKNOWN;
+
+    size_t endPos = DXIL.rfind(ShaderTypeEnd, pos);
+    if (endPos == String::npos)
+        return SHADER_TYPE_UNKNOWN;
+
+    const char* str = &DXIL[endPos];
+    // clang-format off
+    if (ReverseCmp(str, RayGenShader      )) return SHADER_TYPE_RAY_GEN;
+    if (ReverseCmp(str, MissShader        )) return SHADER_TYPE_RAY_MISS;
+    if (ReverseCmp(str, AnyHitShader      )) return SHADER_TYPE_RAY_ANY_HIT;
+    if (ReverseCmp(str, ClosestHitShader  )) return SHADER_TYPE_RAY_CLOSEST_HIT;
+    if (ReverseCmp(str, IntersectionShader)) return SHADER_TYPE_RAY_INTERSECTION;
+    if (ReverseCmp(str, CallableShader    )) return SHADER_TYPE_CALLABLE;
+    // clang-format on
+    return SHADER_TYPE_UNKNOWN;
+}
+
 } // namespace Diligent