fixed resource state transitions, some improvements for ray tracing

9 files changed, 88 insertions, 69 deletions

diff --git a/Graphics/GraphicsEngineD3D12/include/PipelineStateD3D12Impl.hpp b/Graphics/GraphicsEngineD3D12/include/PipelineStateD3D12Impl.hpp
index 6a27a029..f336c6ff 100644
--- a/Graphics/GraphicsEngineD3D12/include/PipelineStateD3D12Impl.hpp
+++ b/Graphics/GraphicsEngineD3D12/include/PipelineStateD3D12Impl.hpp
@@ -144,8 +144,11 @@ private:
 
     void Destruct();
 
-    CComPtr<ID3D12DeviceChild> m_pd3d12PSO;
-    RootSignature              m_RootSig;
+    void CreateLocalRootSignature(const RayTracingPipelineDesc& Desc);
+
+    CComPtr<ID3D12DeviceChild>   m_pd3d12PSO;
+    RootSignature                m_RootSig;
+    CComPtr<ID3D12RootSignature> m_LocalRootSignature;
 
     // Must be defined before default SRB
     SRBMemoryAllocator m_SRBMemAllocator;
diff --git a/Graphics/GraphicsEngineD3D12/include/RenderDeviceD3D12Impl.hpp b/Graphics/GraphicsEngineD3D12/include/RenderDeviceD3D12Impl.hpp
index ab72e22a..636175d5 100644
--- a/Graphics/GraphicsEngineD3D12/include/RenderDeviceD3D12Impl.hpp
+++ b/Graphics/GraphicsEngineD3D12/include/RenderDeviceD3D12Impl.hpp
@@ -178,10 +178,8 @@ public:
     ShaderVersion     GetMaxShaderModel() const;
     D3D_FEATURE_LEVEL GetD3DFeatureLevel() const;
 
-    static Uint32 GetShaderGroupHandleSize()
-    {
-        return D3D12_SHADER_IDENTIFIER_SIZE_IN_BYTES;
-    }
+    static Uint32 GetShaderGroupHandleSize() { return D3D12_SHADER_IDENTIFIER_SIZE_IN_BYTES; }
+    static Uint32 GetMaxShaderRecordStride() { return D3D12_RAYTRACING_MAX_SHADER_RECORD_STRIDE; }
 
 private:
     template <typename PSOCreateInfoType>
diff --git a/Graphics/GraphicsEngineD3D12/include/ShaderBindingTableD3D12Impl.hpp b/Graphics/GraphicsEngineD3D12/include/ShaderBindingTableD3D12Impl.hpp
index e866c6f2..daf84d14 100644
--- a/Graphics/GraphicsEngineD3D12/include/ShaderBindingTableD3D12Impl.hpp
+++ b/Graphics/GraphicsEngineD3D12/include/ShaderBindingTableD3D12Impl.hpp
@@ -54,10 +54,6 @@ public:
 
     virtual void DILIGENT_CALL_TYPE QueryInterface(const INTERFACE_ID& IID, IObject** ppInterface) override final;
 
-    virtual void DILIGENT_CALL_TYPE Verify() const override;
-
-    virtual void DILIGENT_CALL_TYPE Reset(const ShaderBindingTableDesc& Desc) override;
-
     virtual void DILIGENT_CALL_TYPE ResetHitGroups(Uint32 HitShadersPerInstance) override;
 
     virtual void DILIGENT_CALL_TYPE BindAll(const BindAllAttribs& Attribs) override;
@@ -70,9 +66,6 @@ public:
                                                                   D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE& CallableShaderBindingTable) override;
 
 private:
-    void ValidateDesc(const ShaderBindingTableDesc& Desc) const;
-
-private:
     RefCntAutoPtr<IBuffer> m_pBuffer;
 };
 
diff --git a/Graphics/GraphicsEngineD3D12/include/TopLevelASD3D12Impl.hpp b/Graphics/GraphicsEngineD3D12/include/TopLevelASD3D12Impl.hpp
index 8eccf530..c61ab368 100644
--- a/Graphics/GraphicsEngineD3D12/include/TopLevelASD3D12Impl.hpp
+++ b/Graphics/GraphicsEngineD3D12/include/TopLevelASD3D12Impl.hpp
@@ -33,6 +33,7 @@
 #include "TopLevelASD3D12.h"
 #include "RenderDeviceD3D12.h"
 #include "TopLevelASBase.hpp"
+#include "BottomLevelASD3D12Impl.hpp"
 #include "D3D12ResourceBase.hpp"
 #include "RenderDeviceD3D12Impl.hpp"
 
@@ -40,10 +41,10 @@ namespace Diligent
 {
 
 /// Top-level acceleration structure object implementation in Direct3D12 backend.
-class TopLevelASD3D12Impl final : public TopLevelASBase<ITopLevelASD3D12, RenderDeviceD3D12Impl>, public D3D12ResourceBase
+class TopLevelASD3D12Impl final : public TopLevelASBase<ITopLevelASD3D12, BottomLevelASD3D12Impl, RenderDeviceD3D12Impl>, public D3D12ResourceBase
 {
 public:
-    using TTopLevelASBase = TopLevelASBase<ITopLevelASD3D12, RenderDeviceD3D12Impl>;
+    using TTopLevelASBase = TopLevelASBase<ITopLevelASD3D12, BottomLevelASD3D12Impl, RenderDeviceD3D12Impl>;
 
     TopLevelASD3D12Impl(IReferenceCounters*          pRefCounters,
                         class RenderDeviceD3D12Impl* pDeviceD3D12,
diff --git a/Graphics/GraphicsEngineD3D12/interface/ShaderBindingTableD3D12.h b/Graphics/GraphicsEngineD3D12/interface/ShaderBindingTableD3D12.h
index 4b0835cd..aa50251a 100644
--- a/Graphics/GraphicsEngineD3D12/interface/ShaderBindingTableD3D12.h
+++ b/Graphics/GraphicsEngineD3D12/interface/ShaderBindingTableD3D12.h
@@ -65,6 +65,7 @@ DILIGENT_END_INTERFACE
 
 #if DILIGENT_C_INTERFACE
 
+#    define IShaderBindingTableD3D12_GetD3D12AddressRangeAndStride(This, ...) CALL_IFACE_METHOD(ShaderBindingTableD3D12, GetD3D12AddressRangeAndStride, This, __VA_ARGS__)
 
 #endif
 
diff --git a/Graphics/GraphicsEngineD3D12/src/DeviceContextD3D12Impl.cpp b/Graphics/GraphicsEngineD3D12/src/DeviceContextD3D12Impl.cpp
index 3cd896ac..0adcaa59 100644
--- a/Graphics/GraphicsEngineD3D12/src/DeviceContextD3D12Impl.cpp
+++ b/Graphics/GraphicsEngineD3D12/src/DeviceContextD3D12Impl.cpp
@@ -2198,8 +2198,6 @@ void DeviceContextD3D12Impl::TransitionOrVerifyTLASState(CommandContext&
                                                          RESOURCE_STATE                 RequiredState,
                                                          const char*                    OperationName)
 {
-    // AZ TODO: transit BLAS state too?
-
     if (TransitionMode == RESOURCE_STATE_TRANSITION_MODE_TRANSITION)
     {
         if (TLAS.IsInKnownState() && !TLAS.CheckState(RequiredState))
@@ -2210,6 +2208,11 @@ void DeviceContextD3D12Impl::TransitionOrVerifyTLASState(CommandContext&
     {
         DvpVerifyTLASState(TLAS, RequiredState, OperationName);
     }
+
+    if (RequiredState & RESOURCE_STATE_RAY_TRACING)
+    {
+        TLAS.CheckBLASVersion();
+    }
 #endif
 }
 
@@ -2318,6 +2321,8 @@ void DeviceContextD3D12Impl::BuildBLAS(const BLASBuildAttribs& Attribs)
             d3d12Tris.VertexBuffer.StartAddress  = pVB->GetGPUAddress() + SrcTris.VertexOffset;
             d3d12Tris.VertexBuffer.StrideInBytes = SrcTris.VertexStride;
 
+            TransitionOrVerifyBufferState(CmdCtx, *pVB, Attribs.GeometryTransitionMode, RESOURCE_STATE_BUILD_AS_READ, OpName);
+
             if (SrcTris.pIndexBuffer)
             {
                 auto* const pIB       = ValidatedCast<BufferD3D12Impl>(SrcTris.pIndexBuffer);
@@ -2389,6 +2394,10 @@ void DeviceContextD3D12Impl::BuildBLAS(const BLASBuildAttribs& Attribs)
 
     CmdCtx.AsGraphicsContext4().BuildRaytracingAccelerationStructure(Desc, 0, nullptr);
     ++m_State.NumCommands;
+
+#ifdef DILIGENT_DEVELOPMENT
+    pBLASD12->UpdateVersion();
+#endif
 }
 
 void DeviceContextD3D12Impl::BuildTLAS(const TLASBuildAttribs& Attribs)
@@ -2401,7 +2410,6 @@ void DeviceContextD3D12Impl::BuildTLAS(const TLASBuildAttribs& Attribs)
     auto* pTLASD12      = ValidatedCast<TopLevelASD3D12Impl>(Attribs.pTLAS);
     auto* pScratchD12   = ValidatedCast<BufferD3D12Impl>(Attribs.pScratchBuffer);
     auto* pInstancesD12 = ValidatedCast<BufferD3D12Impl>(Attribs.pInstanceBuffer);
-    //auto& TLASDesc      = pTLASD12->GetDesc();
 
     auto&       CmdCtx = GetCmdContext();
     const char* OpName = "Build TopLevelAS (DeviceContextD3D12Impl::BuildTLAS)";
@@ -2423,7 +2431,7 @@ void DeviceContextD3D12Impl::BuildTLAS(const TLASBuildAttribs& Attribs)
             auto* const pBLASD12  = ValidatedCast<BottomLevelASD3D12Impl>(Inst.pBLAS);
 
             static_assert(sizeof(d3d12Inst.Transform) == sizeof(Inst.Transform), "size mismatch");
-            std::memcpy(&d3d12Inst.Transform, Inst.Transform, sizeof(d3d12Inst.Transform));
+            std::memcpy(&d3d12Inst.Transform, Inst.Transform.data, sizeof(d3d12Inst.Transform));
 
             d3d12Inst.InstanceID                          = Inst.CustomId;
             d3d12Inst.InstanceContributionToHitGroupIndex = pTLASD12->GetInstanceDesc(Inst.InstanceName).ContributionToHitGroupIndex; // AZ TODO: optimize
@@ -2458,7 +2466,20 @@ void DeviceContextD3D12Impl::CopyBLAS(const CopyBLASAttribs& Attribs)
     if (!TDeviceContextBase::CopyBLAS(Attribs, 0))
         return;
 
-    // AZ TODO
+    auto* pSrcD3D12 = ValidatedCast<BottomLevelASD3D12Impl>(Attribs.pSrc);
+    auto* pDstD3D12 = ValidatedCast<BottomLevelASD3D12Impl>(Attribs.pDst);
+    auto& CmdCtx    = GetCmdContext();
+
+    const char* OpName = "Copy BottomLevelAS (DeviceContextD3D12Impl::CopyBLAS)";
+    TransitionOrVerifyBLASState(CmdCtx, *pSrcD3D12, Attribs.TransitionMode, RESOURCE_STATE_BUILD_AS_READ, OpName);
+    TransitionOrVerifyBLASState(CmdCtx, *pDstD3D12, Attribs.TransitionMode, RESOURCE_STATE_BUILD_AS_WRITE, OpName);
+
+    CmdCtx.AsGraphicsContext4().CopyRaytracingAccelerationStructure(pSrcD3D12->GetGPUAddress(), pDstD3D12->GetGPUAddress(), D3D12_RAYTRACING_ACCELERATION_STRUCTURE_COPY_MODE_CLONE);
+    ++m_State.NumCommands;
+
+#ifdef DILIGENT_DEVELOPMENT
+    pDstD3D12->UpdateVersion();
+#endif
 }
 
 void DeviceContextD3D12Impl::CopyTLAS(const CopyTLASAttribs& Attribs)
@@ -2466,7 +2487,18 @@ void DeviceContextD3D12Impl::CopyTLAS(const CopyTLASAttribs& Attribs)
     if (!TDeviceContextBase::CopyTLAS(Attribs, 0))
         return;
 
-    // AZ TODO
+    auto* pSrcD3D12 = ValidatedCast<TopLevelASD3D12Impl>(Attribs.pSrc);
+    auto* pDstD3D12 = ValidatedCast<TopLevelASD3D12Impl>(Attribs.pDst);
+    auto& CmdCtx    = GetCmdContext();
+
+    pDstD3D12->CopyInstancceData(*pSrcD3D12);
+
+    const char* OpName = "Copy BottomLevelAS (DeviceContextD3D12Impl::CopyTLAS)";
+    TransitionOrVerifyTLASState(CmdCtx, *pSrcD3D12, Attribs.TransitionMode, RESOURCE_STATE_BUILD_AS_READ, OpName);
+    TransitionOrVerifyTLASState(CmdCtx, *pDstD3D12, Attribs.TransitionMode, RESOURCE_STATE_BUILD_AS_WRITE, OpName);
+
+    CmdCtx.AsGraphicsContext4().CopyRaytracingAccelerationStructure(pSrcD3D12->GetGPUAddress(), pDstD3D12->GetGPUAddress(), D3D12_RAYTRACING_ACCELERATION_STRUCTURE_COPY_MODE_CLONE);
+    ++m_State.NumCommands;
 }
 
 void DeviceContextD3D12Impl::TraceRays(const TraceRaysAttribs& Attribs)
diff --git a/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp b/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
index 405f6d4b..b1fe19fe 100644
--- a/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
+++ b/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
@@ -224,7 +224,7 @@ void BuildRTPipelineDescription(const RayTracingPipelineStateCreateInfo& CreateI
 }
 
 template <typename TNameToGroupIndexMap>
-void GetShaderIdentifiers(ID3D12StateObject*                       pSO,
+void GetShaderIdentifiers(ID3D12DeviceChild*                       pSO,
                           const RayTracingPipelineStateCreateInfo& CreateInfo,
                           const TNameToGroupIndexMap&              NameToGroupIndex,
                           Uint8*                                   ShaderData)
@@ -625,6 +625,8 @@ PipelineStateD3D12Impl::PipelineStateD3D12Impl(IReferenceCounters*
 {
     try
     {
+        CreateLocalRootSignature(CreateInfo.RayTracingPipeline);
+
         TShaderStages                      ShaderStages;
         std::vector<D3D12_STATE_SUBOBJECT> Subobjects;
         DynamicLinearAllocator             TempPool{GetRawAllocator(), 4 << 10};
@@ -640,21 +642,21 @@ PipelineStateD3D12Impl::PipelineStateD3D12Impl(IReferenceCounters*
         D3D12_GLOBAL_ROOT_SIGNATURE GlobalRoot = {m_RootSig.GetD3D12RootSignature()};
         Subobjects.push_back({D3D12_STATE_SUBOBJECT_TYPE_GLOBAL_ROOT_SIGNATURE, &GlobalRoot});
 
+        D3D12_LOCAL_ROOT_SIGNATURE LocalRoot = {m_LocalRootSignature};
+        if (m_LocalRootSignature)
+            Subobjects.push_back({D3D12_STATE_SUBOBJECT_TYPE_LOCAL_ROOT_SIGNATURE, &LocalRoot});
+
         D3D12_STATE_OBJECT_DESC RTPipelineDesc = {};
         RTPipelineDesc.Type                    = D3D12_STATE_OBJECT_TYPE_RAYTRACING_PIPELINE;
         RTPipelineDesc.NumSubobjects           = static_cast<UINT>(Subobjects.size());
         RTPipelineDesc.pSubobjects             = Subobjects.data();
 
-        CComPtr<ID3D12StateObject> pSO;
-
         auto    pd3d12Device = pDeviceD3D12->GetD3D12Device5();
-        HRESULT hr           = pd3d12Device->CreateStateObject(&RTPipelineDesc, IID_PPV_ARGS(&pSO));
+        HRESULT hr           = pd3d12Device->CreateStateObject(&RTPipelineDesc, IID_PPV_ARGS(&m_pd3d12PSO));
         if (FAILED(hr))
             LOG_ERROR_AND_THROW("Failed to create ray tracing state object");
 
-        m_pd3d12PSO = pSO;
-
-        GetShaderIdentifiers(pSO, CreateInfo, m_pRayTracingPipelineData->NameToGroupIndex, m_pRayTracingPipelineData->Shaders);
+        GetShaderIdentifiers(m_pd3d12PSO, CreateInfo, m_pRayTracingPipelineData->NameToGroupIndex, m_pRayTracingPipelineData->Shaders);
 
         if (*m_Desc.Name != 0)
         {
@@ -672,6 +674,35 @@ PipelineStateD3D12Impl::PipelineStateD3D12Impl(IReferenceCounters*
     }
 }
 
+void PipelineStateD3D12Impl::CreateLocalRootSignature(const RayTracingPipelineDesc& Desc)
+{
+    // AZ TODO
+    /*if (Desc.ShaderRecordSize == 0)
+        return;
+
+    D3D12_ROOT_SIGNATURE_DESC d3d12RootSignatureDesc = {};
+    D3D12_ROOT_PARAMETER      d3d12Params            = {};
+
+    d3d12Params.ParameterType            = D3D12_ROOT_PARAMETER_TYPE_32BIT_CONSTANTS;
+    d3d12Params.ShaderVisibility         = D3D12_SHADER_VISIBILITY_ALL;
+    d3d12Params.Constants.Num32BitValues = Desc.ShaderRecordSize / 4;
+    d3d12Params.Constants.RegisterSpace  = Desc.LocalRootRegisterSpace;
+    d3d12Params.Constants.ShaderRegister = 0;
+
+    d3d12RootSignatureDesc.Flags         = D3D12_ROOT_SIGNATURE_FLAG_LOCAL_ROOT_SIGNATURE;
+    d3d12RootSignatureDesc.NumParameters = 1;
+    d3d12RootSignatureDesc.pParameters   = &d3d12Params;
+
+    CComPtr<ID3DBlob> signature;
+    auto              hr = D3D12SerializeRootSignature(&d3d12RootSignatureDesc, D3D_ROOT_SIGNATURE_VERSION_1, &signature, nullptr);
+    CHECK_D3D_RESULT_THROW(hr, "Failed to serialize root signature");
+
+    auto pd3d12Device = GetDevice()->GetD3D12Device();
+
+    hr = pd3d12Device->CreateRootSignature(0, signature->GetBufferPointer(), signature->GetBufferSize(), IID_PPV_ARGS(&m_LocalRootSignature));
+    CHECK_D3D_RESULT_THROW(hr, "Failed to create root signature");*/
+}
+
 PipelineStateD3D12Impl::~PipelineStateD3D12Impl()
 {
     Destruct();
diff --git a/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp b/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp
index d0db2201..e6af12c8 100644
--- a/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp
+++ b/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp
@@ -705,7 +705,7 @@ __forceinline void TransitionResource(CommandContext&                     Ctx,
         {
             VERIFY(RangeType == D3D12_DESCRIPTOR_RANGE_TYPE_SRV, "Unexpected descriptor range type");
             auto* pTLASD3D12 = Res.pObject.RawPtr<TopLevelASD3D12Impl>();
-            if (pTLASD3D12->IsInKnownState() && !pTLASD3D12->CheckState(RESOURCE_STATE_RAY_TRACING))
+            if (pTLASD3D12->IsInKnownState())
                 Ctx.TransitionResource(pTLASD3D12, RESOURCE_STATE_RAY_TRACING);
         }
         break;
diff --git a/Graphics/GraphicsEngineD3D12/src/ShaderBindingTableD3D12Impl.cpp b/Graphics/GraphicsEngineD3D12/src/ShaderBindingTableD3D12Impl.cpp
index d352caf0..d71d98b7 100644
--- a/Graphics/GraphicsEngineD3D12/src/ShaderBindingTableD3D12Impl.cpp
+++ b/Graphics/GraphicsEngineD3D12/src/ShaderBindingTableD3D12Impl.cpp
@@ -44,9 +44,6 @@ ShaderBindingTableD3D12Impl::ShaderBindingTableD3D12Impl(IReferenceCounters*
                                                          bool                          bIsDeviceInternal) :
     TShaderBindingTableBase{pRefCounters, pDeviceD3D12, Desc, bIsDeviceInternal}
 {
-    ValidateDesc(Desc);
-
-    m_ShaderRecordStride = m_Desc.ShaderRecordSize + D3D12_SHADER_IDENTIFIER_SIZE_IN_BYTES;
 }
 
 ShaderBindingTableD3D12Impl::~ShaderBindingTableD3D12Impl()
@@ -55,43 +52,6 @@ ShaderBindingTableD3D12Impl::~ShaderBindingTableD3D12Impl()
 
 IMPLEMENT_QUERY_INTERFACE(ShaderBindingTableD3D12Impl, IID_ShaderBindingTableD3D12, TShaderBindingTableBase)
 
-void ShaderBindingTableD3D12Impl::ValidateDesc(const ShaderBindingTableDesc& Desc) const
-{
-    if (Desc.ShaderRecordSize + D3D12_SHADER_IDENTIFIER_SIZE_IN_BYTES > D3D12_RAYTRACING_MAX_SHADER_RECORD_STRIDE)
-    {
-        LOG_ERROR_AND_THROW("Description of Shader binding table '", (Desc.Name ? Desc.Name : ""),
-                            "' is invalid: ShaderRecordSize is too big, max size is: ", D3D12_RAYTRACING_MAX_SHADER_RECORD_STRIDE - D3D12_SHADER_IDENTIFIER_SIZE_IN_BYTES);
-    }
-}
-
-void ShaderBindingTableD3D12Impl::Verify() const
-{
-    // AZ TODO
-}
-
-void ShaderBindingTableD3D12Impl::Reset(const ShaderBindingTableDesc& Desc)
-{
-    m_RayGenShaderRecord.clear();
-    m_MissShadersRecord.clear();
-    m_CallableShadersRecord.clear();
-    m_HitGroupsRecord.clear();
-    m_Changed = true;
-
-    try
-    {
-        ValidateShaderBindingTableDesc(Desc);
-        ValidateDesc(Desc);
-    }
-    catch (const std::runtime_error&)
-    {
-        // AZ TODO
-        return;
-    }
-
-    m_Desc               = Desc;
-    m_ShaderRecordStride = m_Desc.ShaderRecordSize + D3D12_SHADER_IDENTIFIER_SIZE_IN_BYTES;
-}
-
 void ShaderBindingTableD3D12Impl::ResetHitGroups(Uint32 HitShadersPerInstance)
 {
     // AZ TODO