PSO refactoring for ray tracing

4 files changed, 106 insertions, 63 deletions

diff --git a/Graphics/GraphicsEngineD3DBase/include/D3DShaderResourceLoader.hpp b/Graphics/GraphicsEngineD3DBase/include/D3DShaderResourceLoader.hpp
index f63aa8ad..b12ef6ae 100644
--- a/Graphics/GraphicsEngineD3DBase/include/D3DShaderResourceLoader.hpp
+++ b/Graphics/GraphicsEngineD3DBase/include/D3DShaderResourceLoader.hpp
@@ -42,12 +42,13 @@ namespace Diligent
 
 struct D3DShaderResourceCounters
 {
-    Uint32 NumCBs      = 0;
-    Uint32 NumTexSRVs  = 0;
-    Uint32 NumTexUAVs  = 0;
-    Uint32 NumBufSRVs  = 0;
-    Uint32 NumBufUAVs  = 0;
-    Uint32 NumSamplers = 0;
+    Uint32 NumCBs          = 0;
+    Uint32 NumTexSRVs      = 0;
+    Uint32 NumTexUAVs      = 0;
+    Uint32 NumBufSRVs      = 0;
+    Uint32 NumBufUAVs      = 0;
+    Uint32 NumSamplers     = 0;
+    Uint32 NumAccelStructs = 0;
 };
 
 template <typename D3D_SHADER_DESC,
@@ -61,7 +62,8 @@ template <typename D3D_SHADER_DESC,
           typename TOnNewBuffUAV,
           typename TOnNewBuffSRV,
           typename TOnNewSampler,
-          typename TOnNewTexSRV>
+          typename TOnNewTexSRV,
+          typename TOnNewAccelStruct>
 void LoadD3DShaderResources(TShaderReflection*  pShaderReflection,
                             THandleShaderDesc   HandleShaderDesc,
                             TOnResourcesCounted OnResourcesCounted,
@@ -70,7 +72,8 @@ void LoadD3DShaderResources(TShaderReflection*  pShaderReflection,
                             TOnNewBuffUAV       OnNewBuffUAV,
                             TOnNewBuffSRV       OnNewBuffSRV,
                             TOnNewSampler       OnNewSampler,
-                            TOnNewTexSRV        OnNewTexSRV)
+                            TOnNewTexSRV        OnNewTexSRV,
+                            TOnNewAccelStruct   OnNewAccelStruct)
 {
     D3D_SHADER_DESC shaderDesc = {};
     pShaderReflection->GetDesc(&shaderDesc);
@@ -173,6 +176,7 @@ void LoadD3DShaderResources(TShaderReflection*  pShaderReflection,
             case D3D_SIT_UAV_APPEND_STRUCTURED:         UNSUPPORTED( "Append structured buffers are not supported" );                          break;
             case D3D_SIT_UAV_CONSUME_STRUCTURED:        UNSUPPORTED( "Consume structured buffers are not supported" );                         break;
             case D3D_SIT_UAV_RWSTRUCTURED_WITH_COUNTER: UNSUPPORTED( "RW structured buffers with counter are not supported" );                 break;
+            case D3D_SIT_RTACCELERATIONSTRUCTURE:       ++RC.NumAccelStructs;                                                                  break;
             // clang-format on
             default: UNEXPECTED("Unexpected resource type");
         }
@@ -284,6 +288,12 @@ void LoadD3DShaderResources(TShaderReflection*  pShaderReflection,
                 break;
             }
 
+            case D3D_SIT_RTACCELERATIONSTRUCTURE:
+            {
+                OnNewAccelStruct(Res);
+                break;
+            }
+
             default:
             {
                 UNEXPECTED("Unexpected resource input type");
diff --git a/Graphics/GraphicsEngineD3DBase/include/ShaderResources.hpp b/Graphics/GraphicsEngineD3DBase/include/ShaderResources.hpp
index efbff703..cc24414e 100644
--- a/Graphics/GraphicsEngineD3DBase/include/ShaderResources.hpp
+++ b/Graphics/GraphicsEngineD3DBase/include/ShaderResources.hpp
@@ -280,20 +280,22 @@ public:
     ~ShaderResources();
 
     // clang-format off
-    Uint32 GetNumCBs()        const noexcept { return (m_TexSRVOffset   - 0);                }
-    Uint32 GetNumTexSRV()     const noexcept { return (m_TexUAVOffset   - m_TexSRVOffset);   }
-    Uint32 GetNumTexUAV()     const noexcept { return (m_BufSRVOffset   - m_TexUAVOffset);   }
-    Uint32 GetNumBufSRV()     const noexcept { return (m_BufUAVOffset   - m_BufSRVOffset);   }
-    Uint32 GetNumBufUAV()     const noexcept { return (m_SamplersOffset - m_BufUAVOffset);   }
-    Uint32 GetNumSamplers()   const noexcept { return (m_TotalResources - m_SamplersOffset); }
-    Uint32 GetTotalResources()const noexcept { return  m_TotalResources;                     }
-
-    const D3DShaderResourceAttribs& GetCB     (Uint32 n)const noexcept { return GetResAttribs(n, GetNumCBs(),                   0);   }
-    const D3DShaderResourceAttribs& GetTexSRV (Uint32 n)const noexcept { return GetResAttribs(n, GetNumTexSRV(),   m_TexSRVOffset);   }
-    const D3DShaderResourceAttribs& GetTexUAV (Uint32 n)const noexcept { return GetResAttribs(n, GetNumTexUAV(),   m_TexUAVOffset);   }
-    const D3DShaderResourceAttribs& GetBufSRV (Uint32 n)const noexcept { return GetResAttribs(n, GetNumBufSRV(),   m_BufSRVOffset);   }
-    const D3DShaderResourceAttribs& GetBufUAV (Uint32 n)const noexcept { return GetResAttribs(n, GetNumBufUAV(),   m_BufUAVOffset);   }
-    const D3DShaderResourceAttribs& GetSampler(Uint32 n)const noexcept { return GetResAttribs(n, GetNumSamplers(), m_SamplersOffset); }
+    Uint32 GetNumCBs()         const noexcept { return (m_TexSRVOffset       - 0);                    }
+    Uint32 GetNumTexSRV()      const noexcept { return (m_TexUAVOffset       - m_TexSRVOffset);       }
+    Uint32 GetNumTexUAV()      const noexcept { return (m_BufSRVOffset       - m_TexUAVOffset);       }
+    Uint32 GetNumBufSRV()      const noexcept { return (m_BufUAVOffset       - m_BufSRVOffset);       }
+    Uint32 GetNumBufUAV()      const noexcept { return (m_SamplersOffset     - m_BufUAVOffset);       }
+    Uint32 GetNumSamplers()    const noexcept { return (m_AccelStructsOffset - m_SamplersOffset);     }
+    Uint32 GetNumAccelStructs()const noexcept { return (m_TotalResources     - m_AccelStructsOffset); }
+    Uint32 GetTotalResources() const noexcept { return  m_TotalResources;                             }
+
+    const D3DShaderResourceAttribs& GetCB         (Uint32 n)const noexcept { return GetResAttribs(n, GetNumCBs(),                       0);       }
+    const D3DShaderResourceAttribs& GetTexSRV     (Uint32 n)const noexcept { return GetResAttribs(n, GetNumTexSRV(),       m_TexSRVOffset);       }
+    const D3DShaderResourceAttribs& GetTexUAV     (Uint32 n)const noexcept { return GetResAttribs(n, GetNumTexUAV(),       m_TexUAVOffset);       }
+    const D3DShaderResourceAttribs& GetBufSRV     (Uint32 n)const noexcept { return GetResAttribs(n, GetNumBufSRV(),       m_BufSRVOffset);       }
+    const D3DShaderResourceAttribs& GetBufUAV     (Uint32 n)const noexcept { return GetResAttribs(n, GetNumBufUAV(),       m_BufUAVOffset);       }
+    const D3DShaderResourceAttribs& GetSampler    (Uint32 n)const noexcept { return GetResAttribs(n, GetNumSamplers(),     m_SamplersOffset);     }
+    const D3DShaderResourceAttribs& GetAccelStruct(Uint32 n)const noexcept { return GetResAttribs(n, GetNumAccelStructs(), m_AccelStructsOffset); }
     // clang-format on
 
     const D3DShaderResourceAttribs& GetCombinedSampler(const D3DShaderResourceAttribs& TexSRV) const noexcept
@@ -317,13 +319,15 @@ public:
               typename THandleTexSRV,
               typename THandleTexUAV,
               typename THandleBufSRV,
-              typename THandleBufUAV>
-    void ProcessResources(THandleCB      HandleCB,
-                          THandleSampler HandleSampler,
-                          THandleTexSRV  HandleTexSRV,
-                          THandleTexUAV  HandleTexUAV,
-                          THandleBufSRV  HandleBufSRV,
-                          THandleBufUAV  HandleBufUAV) const
+              typename THandleBufUAV,
+              typename THandleAccelStruct>
+    void ProcessResources(THandleCB          HandleCB,
+                          THandleSampler     HandleSampler,
+                          THandleTexSRV      HandleTexSRV,
+                          THandleTexUAV      HandleTexUAV,
+                          THandleBufSRV      HandleBufSRV,
+                          THandleBufUAV      HandleBufUAV,
+                          THandleAccelStruct HandleAccelStruct) const
     {
         for (Uint32 n = 0; n < GetNumCBs(); ++n)
         {
@@ -360,6 +364,12 @@ public:
             const auto& BufUAV = GetBufUAV(n);
             HandleBufUAV(BufUAV, n);
         }
+
+        for (Uint32 n = 0; n < GetNumAccelStructs(); ++n)
+        {
+            const auto& AS = GetAccelStruct(n);
+            HandleAccelStruct(AS, n);
+        }
     }
 
     bool        IsCompatibleWith(const ShaderResources& Resources) const;
@@ -420,12 +430,13 @@ protected:
     }
 
     // clang-format off
-    D3DShaderResourceAttribs& GetCB(Uint32 n)      noexcept { return GetResAttribs(n, GetNumCBs(), 0); }
-    D3DShaderResourceAttribs& GetTexSRV(Uint32 n)  noexcept { return GetResAttribs(n, GetNumTexSRV(), m_TexSRVOffset); }
-    D3DShaderResourceAttribs& GetTexUAV(Uint32 n)  noexcept { return GetResAttribs(n, GetNumTexUAV(), m_TexUAVOffset); }
-    D3DShaderResourceAttribs& GetBufSRV(Uint32 n)  noexcept { return GetResAttribs(n, GetNumBufSRV(), m_BufSRVOffset); }
-    D3DShaderResourceAttribs& GetBufUAV(Uint32 n)  noexcept { return GetResAttribs(n, GetNumBufUAV(), m_BufUAVOffset); }
-    D3DShaderResourceAttribs& GetSampler(Uint32 n) noexcept { return GetResAttribs(n, GetNumSamplers(), m_SamplersOffset); }
+    D3DShaderResourceAttribs& GetCB(Uint32 n)         noexcept { return GetResAttribs(n, GetNumCBs(), 0); }
+    D3DShaderResourceAttribs& GetTexSRV(Uint32 n)     noexcept { return GetResAttribs(n, GetNumTexSRV(), m_TexSRVOffset); }
+    D3DShaderResourceAttribs& GetTexUAV(Uint32 n)     noexcept { return GetResAttribs(n, GetNumTexUAV(), m_TexUAVOffset); }
+    D3DShaderResourceAttribs& GetBufSRV(Uint32 n)     noexcept { return GetResAttribs(n, GetNumBufSRV(), m_BufSRVOffset); }
+    D3DShaderResourceAttribs& GetBufUAV(Uint32 n)     noexcept { return GetResAttribs(n, GetNumBufUAV(), m_BufUAVOffset); }
+    D3DShaderResourceAttribs& GetSampler(Uint32 n)    noexcept { return GetResAttribs(n, GetNumSamplers(), m_SamplersOffset); }
+    D3DShaderResourceAttribs& GetAccelStruct(Uint32 n)noexcept { return GetResAttribs(n, GetNumAccelStructs(), m_AccelStructsOffset); }
     // clang-format on
 
 private:
@@ -447,12 +458,13 @@ private:
 
     // Offsets in elements of D3DShaderResourceAttribs
     typedef Uint16 OffsetType;
-    OffsetType     m_TexSRVOffset   = 0;
-    OffsetType     m_TexUAVOffset   = 0;
-    OffsetType     m_BufSRVOffset   = 0;
-    OffsetType     m_BufUAVOffset   = 0;
-    OffsetType     m_SamplersOffset = 0;
-    OffsetType     m_TotalResources = 0;
+    OffsetType     m_TexSRVOffset       = 0;
+    OffsetType     m_TexUAVOffset       = 0;
+    OffsetType     m_BufSRVOffset       = 0;
+    OffsetType     m_BufUAVOffset       = 0;
+    OffsetType     m_SamplersOffset     = 0;
+    OffsetType     m_AccelStructsOffset = 0;
+    OffsetType     m_TotalResources     = 0;
 
     const SHADER_TYPE m_ShaderType;
 
@@ -469,7 +481,7 @@ void ShaderResources::Initialize(TShaderReflection*  pShaderReflection,
                                  const Char*         ShaderName,
                                  const Char*         CombinedSamplerSuffix)
 {
-    Uint32 CurrCB = 0, CurrTexSRV = 0, CurrTexUAV = 0, CurrBufSRV = 0, CurrBufUAV = 0, CurrSampler = 0;
+    Uint32 CurrCB = 0, CurrTexSRV = 0, CurrTexUAV = 0, CurrBufSRV = 0, CurrBufUAV = 0, CurrSampler = 0, CurrAS = 0;
 
     // Resource names pool is only needed to facilitate string allocation.
     StringPool ResourceNamesPool;
@@ -541,6 +553,13 @@ void ShaderResources::Initialize(TShaderReflection*  pShaderReflection,
             }
             ++CurrTexSRV;
             NewResHandler.OnNewTexSRV(*pNewTexSRV);
+        },
+
+        [&](const D3DShaderResourceAttribs& AccelStructAttribs) //
+        {
+            VERIFY_EXPR(AccelStructAttribs.GetInputType() == D3D_SIT_RTACCELERATIONSTRUCTURE);
+            auto* pNewAccelStruct = new (&GetAccelStruct(CurrAS++)) D3DShaderResourceAttribs{ResourceNamesPool, AccelStructAttribs};
+            NewResHandler.OnNewAccelStruct(*pNewAccelStruct);
         } //
     );
 
@@ -562,12 +581,13 @@ void ShaderResources::Initialize(TShaderReflection*  pShaderReflection,
 
     VERIFY_EXPR(ResourceNamesPool.GetRemainingSize() == 0);
     // clang-format off
-    VERIFY(CurrCB      == GetNumCBs(),      "Not all CBs are initialized which will cause a crash when ~D3DShaderResourceAttribs() is called");
-    VERIFY(CurrTexSRV  == GetNumTexSRV(),   "Not all Tex SRVs are initialized which will cause a crash when ~D3DShaderResourceAttribs() is called" );
-    VERIFY(CurrTexUAV  == GetNumTexUAV(),   "Not all Tex UAVs are initialized which will cause a crash when ~D3DShaderResourceAttribs() is called" );
-    VERIFY(CurrBufSRV  == GetNumBufSRV(),   "Not all Buf SRVs are initialized which will cause a crash when ~D3DShaderResourceAttribs() is called" );
-    VERIFY(CurrBufUAV  == GetNumBufUAV(),   "Not all Buf UAVs are initialized which will cause a crash when ~D3DShaderResourceAttribs() is called" );
-    VERIFY(CurrSampler == GetNumSamplers(), "Not all Samplers are initialized which will cause a crash when ~D3DShaderResourceAttribs() is called" );
+    VERIFY(CurrCB      == GetNumCBs(),          "Not all CBs are initialized which will cause a crash when ~D3DShaderResourceAttribs() is called");
+    VERIFY(CurrTexSRV  == GetNumTexSRV(),       "Not all Tex SRVs are initialized which will cause a crash when ~D3DShaderResourceAttribs() is called" );
+    VERIFY(CurrTexUAV  == GetNumTexUAV(),       "Not all Tex UAVs are initialized which will cause a crash when ~D3DShaderResourceAttribs() is called" );
+    VERIFY(CurrBufSRV  == GetNumBufSRV(),       "Not all Buf SRVs are initialized which will cause a crash when ~D3DShaderResourceAttribs() is called" );
+    VERIFY(CurrBufUAV  == GetNumBufUAV(),       "Not all Buf UAVs are initialized which will cause a crash when ~D3DShaderResourceAttribs() is called" );
+    VERIFY(CurrSampler == GetNumSamplers(),     "Not all Samplers are initialized which will cause a crash when ~D3DShaderResourceAttribs() is called" );
+    VERIFY(CurrAS      == GetNumAccelStructs(), "Not all Accel Structs are initialized which will cause a crash when ~D3DShaderResourceAttribs() is called" );
     // clang-format on
 }
 
diff --git a/Graphics/GraphicsEngineD3DBase/include/ShaderVariableD3DBase.hpp b/Graphics/GraphicsEngineD3DBase/include/ShaderVariableD3DBase.hpp
index 2497b17d..78525f30 100644
--- a/Graphics/GraphicsEngineD3DBase/include/ShaderVariableD3DBase.hpp
+++ b/Graphics/GraphicsEngineD3DBase/include/ShaderVariableD3DBase.hpp
@@ -87,8 +87,8 @@ protected:
 };
 
 
-template <typename BufferViewImplType>
-bool VerifyBufferViewModeD3D(BufferViewImplType* pViewD3D11, const D3DShaderResourceAttribs& Attribs, const char* ShaderName)
+template <typename BufferViewImplType, typename AttribsType>
+bool VerifyBufferViewModeD3D(BufferViewImplType* pViewD3D11, const AttribsType& Attribs, const char* ShaderName)
 {
     if (pViewD3D11 == nullptr)
         return true;
diff --git a/Graphics/GraphicsEngineD3DBase/src/ShaderResources.cpp b/Graphics/GraphicsEngineD3DBase/src/ShaderResources.cpp
index b975b506..79d2cb22 100644
--- a/Graphics/GraphicsEngineD3DBase/src/ShaderResources.cpp
+++ b/Graphics/GraphicsEngineD3DBase/src/ShaderResources.cpp
@@ -74,23 +74,25 @@ void ShaderResources::AllocateMemory(IMemoryAllocator&                Allocator,
     };
 
     // clang-format off
-    auto CBOffset    = AdvanceOffset(ResCounters.NumCBs);       (void)CBOffset; // To suppress warning
-    m_TexSRVOffset   = AdvanceOffset(ResCounters.NumTexSRVs);
-    m_TexUAVOffset   = AdvanceOffset(ResCounters.NumTexUAVs);
-    m_BufSRVOffset   = AdvanceOffset(ResCounters.NumBufSRVs);
-    m_BufUAVOffset   = AdvanceOffset(ResCounters.NumBufUAVs);
-    m_SamplersOffset = AdvanceOffset(ResCounters.NumSamplers);
-    m_TotalResources = AdvanceOffset(0);
+    auto CBOffset        = AdvanceOffset(ResCounters.NumCBs);       (void)CBOffset; // To suppress warning
+    m_TexSRVOffset       = AdvanceOffset(ResCounters.NumTexSRVs);
+    m_TexUAVOffset       = AdvanceOffset(ResCounters.NumTexUAVs);
+    m_BufSRVOffset       = AdvanceOffset(ResCounters.NumBufSRVs);
+    m_BufUAVOffset       = AdvanceOffset(ResCounters.NumBufUAVs);
+    m_SamplersOffset     = AdvanceOffset(ResCounters.NumSamplers);
+    m_AccelStructsOffset = AdvanceOffset(ResCounters.NumAccelStructs);
+    m_TotalResources     = AdvanceOffset(0);
 
     auto AlignedResourceNamesPoolSize = Align(ResourceNamesPoolSize, sizeof(void*));
     auto MemorySize = m_TotalResources * sizeof(D3DShaderResourceAttribs) + AlignedResourceNamesPoolSize * sizeof(char);
 
-    VERIFY_EXPR(GetNumCBs()     == ResCounters.NumCBs);
-    VERIFY_EXPR(GetNumTexSRV()  == ResCounters.NumTexSRVs);
-    VERIFY_EXPR(GetNumTexUAV()  == ResCounters.NumTexUAVs);
-    VERIFY_EXPR(GetNumBufSRV()  == ResCounters.NumBufSRVs);
-    VERIFY_EXPR(GetNumBufUAV()  == ResCounters.NumBufUAVs);
-    VERIFY_EXPR(GetNumSamplers()== ResCounters.NumSamplers);
+    VERIFY_EXPR(GetNumCBs()         == ResCounters.NumCBs);
+    VERIFY_EXPR(GetNumTexSRV()      == ResCounters.NumTexSRVs);
+    VERIFY_EXPR(GetNumTexUAV()      == ResCounters.NumTexUAVs);
+    VERIFY_EXPR(GetNumBufSRV()      == ResCounters.NumBufSRVs);
+    VERIFY_EXPR(GetNumBufUAV()      == ResCounters.NumBufUAVs);
+    VERIFY_EXPR(GetNumSamplers()    == ResCounters.NumSamplers);
+    VERIFY_EXPR(GetNumAccelStructs()== ResCounters.NumAccelStructs);
     // clang-format on
 
     if (MemorySize)
@@ -208,6 +210,12 @@ D3DShaderResourceCounters ShaderResources::CountResources(const PipelineResource
             auto VarType = FindVariableType(BufUAV, ResourceLayout);
             if (IsAllowedType(VarType, AllowedTypeBits))
                 ++Counters.NumBufUAVs;
+        },
+        [&](const D3DShaderResourceAttribs& AccelStruct, Uint32) //
+        {
+            auto VarType = FindVariableType(AccelStruct, ResourceLayout);
+            if (IsAllowedType(VarType, AllowedTypeBits))
+                ++Counters.NumAccelStructs;
         } //
     );
 
@@ -400,6 +408,11 @@ bool ShaderResources::IsCompatibleWith(const ShaderResources& Res) const
         {
             if (!BufUAV.IsCompatibleWith(Res.GetBufUAV(n)))
                 IsCompatible = false;
+        },
+        [&](const D3DShaderResourceAttribs& AccelStruct, Uint32 n) //
+        {
+            if (!AccelStruct.IsCompatibleWith(Res.GetAccelStruct(n)))
+                IsCompatible = false;
         } //
     );
     return IsCompatible;