Ray tracing: eliminated the need for sacral knowledge about scratch b…

…uffer alignment. Now scratch buffer sizes get aligned under the hood to allow comfortable concatenation. Also scratch memory alignment exposed in DeviceDesc (+fixed typos)
NVIDIAGameWorks · Jul 17, 2024 · 11a601d · 11a601d
1 parent b5a2e90
commit 11a601d
Show file tree

Hide file tree

Showing 7 changed files with 20 additions and 17 deletions.
diff --git a/Include/NRIDescs.h b/Include/NRIDescs.h
@@ -1516,11 +1516,12 @@ NRI_STRUCT(DeviceDesc)
     uint32_t computeShaderWorkGroupMaxDim[3];
 
     // Ray tracing
-    uint64_t rayTracingShaderGroupIdentifierSize;
-    uint64_t rayTracingShaderTableAligment;
-    uint64_t rayTracingShaderTableMaxStride;
+    uint32_t rayTracingShaderGroupIdentifierSize;
+    uint32_t rayTracingShaderTableAlignment;
+    uint32_t rayTracingShaderTableMaxStride;
     uint32_t rayTracingShaderRecursionMaxDepth;
     uint32_t rayTracingGeometryObjectMaxNum;
+    uint32_t rayTracingScratchAlignment;
 
     // Mesh shaders
     uint32_t meshControlSharedMemoryMaxSize;

diff --git a/Source/D3D12/AccelerationStructureD3D12.cpp b/Source/D3D12/AccelerationStructureD3D12.cpp
@@ -57,11 +57,11 @@ void AccelerationStructureD3D12::GetMemoryInfo(MemoryDesc& memoryDesc) const {
 }
 
 uint64_t AccelerationStructureD3D12::GetUpdateScratchBufferSize() const {
-    return m_PrebuildInfo.UpdateScratchDataSizeInBytes;
+    return Align(m_PrebuildInfo.UpdateScratchDataSizeInBytes, D3D12_RAYTRACING_ACCELERATION_STRUCTURE_BYTE_ALIGNMENT);
 }
 
 uint64_t AccelerationStructureD3D12::GetBuildScratchBufferSize() const {
-    return m_PrebuildInfo.ScratchDataSizeInBytes;
+    return Align(m_PrebuildInfo.ScratchDataSizeInBytes, D3D12_RAYTRACING_ACCELERATION_STRUCTURE_BYTE_ALIGNMENT);
 }
 
 Result AccelerationStructureD3D12::BindMemory(Memory* memory, uint64_t offset) {

diff --git a/Source/D3D12/DeviceD3D12.cpp b/Source/D3D12/DeviceD3D12.cpp
@@ -525,10 +525,11 @@ void DeviceD3D12::FillDesc(bool enableDrawParametersEmulation) {
 
     if (m_Desc.isRayTracingSupported) {
         m_Desc.rayTracingShaderGroupIdentifierSize = D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT;
-        m_Desc.rayTracingShaderTableAligment = D3D12_RAYTRACING_SHADER_TABLE_BYTE_ALIGNMENT;
-        m_Desc.rayTracingShaderTableMaxStride = std::numeric_limits<uint64_t>::max();
+        m_Desc.rayTracingShaderTableAlignment = D3D12_RAYTRACING_SHADER_TABLE_BYTE_ALIGNMENT;
+        m_Desc.rayTracingShaderTableMaxStride = std::numeric_limits<uint32_t>::max();
         m_Desc.rayTracingShaderRecursionMaxDepth = D3D12_RAYTRACING_MAX_DECLARABLE_TRACE_RECURSION_DEPTH;
         m_Desc.rayTracingGeometryObjectMaxNum = (1 << 24) - 1;
+        m_Desc.rayTracingScratchAlignment = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_BYTE_ALIGNMENT;
     }
 
     if (m_Desc.isMeshShaderSupported) {

diff --git a/Source/Shared/HelperDataUpload.cpp b/Source/Shared/HelperDataUpload.cpp
@@ -5,7 +5,7 @@
 using namespace nri;
 
 constexpr uint32_t BARRIERS_PER_PASS = 256;
-constexpr uint64_t COPY_ALIGMENT = 16;
+constexpr uint64_t COPY_ALIGNMENT = 16;
 
 static void DoTransition(const CoreInterface& NRI, CommandBuffer* commandBuffer, bool isInitial, const TextureUploadDesc* textureUploadDescs, uint32_t textureDataDescNum) {
     TextureBarrierDesc textureBarriers[BARRIERS_PER_PASS];
@@ -84,7 +84,7 @@ Result HelperDataUpload::UploadData(
         m_UploadBufferSize = std::max(m_UploadBufferSize, contentSize);
     }
 
-    m_UploadBufferSize = Align(m_UploadBufferSize, COPY_ALIGMENT);
+    m_UploadBufferSize = Align(m_UploadBufferSize, COPY_ALIGNMENT);
 
     Result result = Create();
     if (result == Result::SUCCESS) {
@@ -278,13 +278,13 @@ bool HelperDataUpload::CopyTextureContent(const TextureUploadDesc& textureUpload
 
             NRI.CmdUploadBufferToTexture(*m_CommandBuffer, *textureUploadDesc.texture, dstRegion, *m_UploadBuffer, srcDataLayout);
 
-            m_UploadBufferOffset = Align(m_UploadBufferOffset + contentSize, COPY_ALIGMENT);
+            m_UploadBufferOffset = Align(m_UploadBufferOffset + contentSize, COPY_ALIGNMENT);
         }
         mipOffset = 0;
     }
     arrayOffset = 0;
 
-    m_UploadBufferOffset = Align(m_UploadBufferOffset, COPY_ALIGMENT);
+    m_UploadBufferOffset = Align(m_UploadBufferOffset, COPY_ALIGNMENT);
 
     return true;
 }
@@ -328,7 +328,7 @@ bool HelperDataUpload::CopyBufferContent(const BufferUploadDesc& bufferUploadDes
         return false;
 
     bufferContentOffset = 0;
-    m_UploadBufferOffset = Align(m_UploadBufferOffset, COPY_ALIGMENT);
+    m_UploadBufferOffset = Align(m_UploadBufferOffset, COPY_ALIGNMENT);
 
     return true;
 }
diff --git a/Source/VK/AccelerationStructureVK.h b/Source/VK/AccelerationStructureVK.h
@@ -38,11 +38,11 @@ struct AccelerationStructureVK {
     //================================================================================================================
 
     inline uint64_t GetUpdateScratchBufferSize() const {
-        return m_UpdateScratchSize;
+        return Align(m_UpdateScratchSize, m_Device.GetDesc().rayTracingScratchAlignment);
     }
 
     inline uint64_t GetBuildScratchBufferSize() const {
-        return m_BuildScratchSize;
+        return Align(m_BuildScratchSize, m_Device.GetDesc().rayTracingScratchAlignment);
     }
 
     inline VkDeviceAddress GetNativeHandle() const {

diff --git a/Source/VK/DeviceVK.cpp b/Source/VK/DeviceVK.cpp
@@ -770,10 +770,11 @@ Result DeviceVK::Create(const DeviceCreationDesc& deviceCreationDesc, const Devi
         m_Desc.computeShaderWorkGroupMaxDim[2] = limits.maxComputeWorkGroupSize[2];
 
         m_Desc.rayTracingShaderGroupIdentifierSize = rayTracingProps.shaderGroupHandleSize;
+        m_Desc.rayTracingShaderTableAlignment = rayTracingProps.shaderGroupBaseAlignment;
+        m_Desc.rayTracingShaderTableMaxStride = rayTracingProps.maxShaderGroupStride;
         m_Desc.rayTracingShaderRecursionMaxDepth = rayTracingProps.maxRayRecursionDepth;
         m_Desc.rayTracingGeometryObjectMaxNum = (uint32_t)accelerationStructureProps.maxGeometryCount;
-        m_Desc.rayTracingShaderTableAligment = rayTracingProps.shaderGroupBaseAlignment;
-        m_Desc.rayTracingShaderTableMaxStride = rayTracingProps.maxShaderGroupStride;
+        m_Desc.rayTracingScratchAlignment = accelerationStructureProps.minAccelerationStructureScratchOffsetAlignment;
 
         m_Desc.meshControlSharedMemoryMaxSize = meshShaderProps.maxTaskSharedMemorySize;
         m_Desc.meshControlWorkGroupInvocationMaxNum = meshShaderProps.maxTaskWorkGroupInvocations;

diff --git a/Source/Validation/CommandBufferVal.cpp b/Source/Validation/CommandBufferVal.cpp
@@ -583,7 +583,7 @@ void CommandBufferVal::WriteAccelerationStructureSize(
 }
 
 void CommandBufferVal::DispatchRays(const DispatchRaysDesc& dispatchRaysDesc) {
-    uint64_t align = m_Device.GetDesc().rayTracingShaderTableAligment;
+    uint64_t align = m_Device.GetDesc().rayTracingShaderTableAlignment;
     RETURN_ON_FAILURE(&m_Device, m_IsRecordingStarted, ReturnVoid(), "CmdDispatchRays: the command buffer must be in the recording state");
     RETURN_ON_FAILURE(&m_Device, !m_IsRenderPass, ReturnVoid(), "CmdDispatchRays: must be called outside of 'CmdBeginRendering/CmdEndRendering'");
     RETURN_ON_FAILURE(&m_Device, dispatchRaysDesc.raygenShader.buffer != nullptr, ReturnVoid(), "CmdDispatchRays: 'dispatchRaysDesc.raygenShader.buffer' is NULL");