Merge pull request #861 from Devsh-Graphics-Programming/tlas_blas_tracking

devshgraphicsprogramming · web-flow · commit 4e431838b0e3 · 2025-04-01T12:13:57.000+02:00
TLAS tracking BLAS they use doing build
diff --git a/include/nbl/video/IDeferredOperation.h b/include/nbl/video/IDeferredOperation.h
@@ -30,7 +30,15 @@ class IDeferredOperation : public IBackendObject
         {
             const auto retval = execute_impl();
             if (retval==STATUS::COMPLETED || retval==STATUS::_ERROR)
+            {
+                std::lock_guard lock(m_completionMutex);
+                if (m_callback)
+                {
+                    m_callback(this);
+                    m_callback = {};
+                }
                 m_resourceTracking.clear();
+            }
             return retval;
         }
 
@@ -66,6 +74,8 @@ class IDeferredOperation : public IBackendObject
     private:
         friend class ILogicalDevice;
         // when we improve allocators, etc. we'll stop using STL containers here
+        std::mutex m_completionMutex;
+        std::function<void(IDeferredOperation*)> m_callback;
         core::vector<core::smart_refctd_ptr<const IReferenceCounted>> m_resourceTracking;
 };
 
diff --git a/include/nbl/video/IGPUAccelerationStructure.h b/include/nbl/video/IGPUAccelerationStructure.h
@@ -400,8 +400,9 @@ class IGPUTopLevelAccelerationStructure : public asset::ITopLevelAccelerationStr
 				template<typename T> requires nbl::is_any_of_v<T,std::conditional_t<std::is_same_v<BufferType,IGPUBuffer>,uint32_t,BuildRangeInfo>,BuildRangeInfo>
 				inline uint32_t valid(const T& buildRangeInfo) const
 				{
+					uint32_t retval = trackedBLASes.size();
 					if constexpr (std::is_same_v<T,uint32_t>)
-						return valid<BuildRangeInfo>({.instanceCount=buildRangeInfo,.instanceByteOffset=0});
+						retval += valid<BuildRangeInfo>({.instanceCount=buildRangeInfo,.instanceByteOffset=0});
 					else
 					{
 						if (IGPUAccelerationStructure::BuildInfo<BufferType>::invalid(srcAS,dstAS))
@@ -444,8 +445,9 @@ class IGPUTopLevelAccelerationStructure : public asset::ITopLevelAccelerationStr
 						#endif
 
 						// destination, scratch and instanceData are required, source is optional
-						return Base::isUpdate ? 4u:3u;
+						retval += Base::isUpdate ? 4u:3u;
 					}
+					return retval;
 				}
 
 				inline core::smart_refctd_ptr<const IReferenceCounted>* fillTracking(core::smart_refctd_ptr<const IReferenceCounted>* oit) const
@@ -457,6 +459,9 @@ class IGPUTopLevelAccelerationStructure : public asset::ITopLevelAccelerationStr
 
 					*(oit++) = core::smart_refctd_ptr<const IReferenceCounted>(instanceData.buffer);
 
+					for (const auto& blas : trackedBLASes)
+						*(oit++) = core::smart_refctd_ptr<const IReferenceCounted>(blas);
+
 					return oit;
 				}
 
@@ -470,6 +475,8 @@ class IGPUTopLevelAccelerationStructure : public asset::ITopLevelAccelerationStr
 				//	+ an array of `PolymorphicInstance` if our `SCreationParams::flags.hasFlags(MOTION_BIT)`, otherwise
 				//	+ an array of `StaticInstance`
 				asset::SBufferBinding<const BufferType> instanceData = {};
+				// [optional] Provide info about what BLAS references to hold onto after the build. For performance make sure the list is compact (without repeated elements).
+				std::span<const IGPUBottomLevelAccelerationStructure*> trackedBLASes = {};
 		};
 		using DeviceBuildInfo = BuildInfo<IGPUBuffer>;
 		using HostBuildInfo = BuildInfo<asset::ICPUBuffer>;
@@ -545,11 +552,71 @@ class IGPUTopLevelAccelerationStructure : public asset::ITopLevelAccelerationStr
 		using HostPolymorphicInstance = PolymorphicInstance<IGPUBottomLevelAccelerationStructure::host_op_ref_t>;
 		static_assert(sizeof(DevicePolymorphicInstance)==sizeof(HostPolymorphicInstance));
 
+		//
+		using build_ver_t = uint32_t;
+		// this gets called when execution is sure to happen 100%, e.g. not during command recording but during submission
+		inline build_ver_t registerNextBuildVer()
+		{
+			return m_pendingBuildVer++;
+		}
+		// 
+		using blas_smart_ptr_t = core::smart_refctd_ptr<const IGPUBottomLevelAccelerationStructure>;
+		// returns number of tracked BLASes if `tracked==nullptr` otherwise writes `*count` tracked BLASes from `first` into `*tracked`
+		inline build_ver_t getTrackedBLASes(uint32_t* count, blas_smart_ptr_t* tracked, const uint32_t first=0) const
+		{
+			if (!count)
+				return 0;
+			// stop multiple threads messing with us
+			std::lock_guard lk(m_trackingLock);
+			const uint32_t toWrite = std::min<uint32_t>(std::max<uint32_t>(m_trackedBLASes.size(),first)-first,tracked ? (*count):0xffFFffFFu);
+			*count = toWrite;
+			if (tracked && toWrite)
+			{
+				auto it = m_trackedBLASes.begin();
+				// cmon its an unordered map, iterator should have operator +=
+				for (auto i=0; i<first; i++)
+					it++;
+				for (auto i=0; i<toWrite; i++)
+					*(tracked++) = *(it++);
+			}
+			return m_completedBuildVer;
+		}
+		// Useful if TLAS got built externally as well, returns if there were no later builds that preempted us setting the result here
+		template<typename Iterator>
+		inline bool setTrackedBLASes(const Iterator begin, const Iterator end, const build_ver_t buildVer)
+		{
+			// stop multiple threads messing with us
+			std::lock_guard lk(m_trackingLock);
+			// stop out of order callbacks
+			if (buildVer<=m_completedBuildVer)
+				return false;
+			m_completedBuildVer = buildVer;
+			// release already tracked BLASes
+			m_trackedBLASes.clear();
+			// sanity check, TODO: this should be an atomic_max on the `m_pendingBuildVer`
+			if (m_completedBuildVer>m_pendingBuildVer)
+				m_pendingBuildVer = m_completedBuildVer;
+			// now fill the contents
+			m_trackedBLASes.insert(begin,end);
+			return true;
+		}
+		// a little utility to make sure nothing from this build version and before gets tracked
+		inline bool clearTrackedBLASes(const build_ver_t buildVer)
+		{
+			return setTrackedBLASes<const blas_smart_ptr_t*>(nullptr,nullptr,buildVer);
+		}
+
 	protected:
 		inline IGPUTopLevelAccelerationStructure(core::smart_refctd_ptr<const ILogicalDevice>&& dev, SCreationParams&& params)
-			: asset::ITopLevelAccelerationStructure<IGPUAccelerationStructure>(std::move(dev),std::move(params)), m_maxInstanceCount(params.maxInstanceCount) {}
+			: asset::ITopLevelAccelerationStructure<IGPUAccelerationStructure>(std::move(dev),std::move(params)),
+			m_maxInstanceCount(params.maxInstanceCount),m_trackedBLASes() {}
 
 		const uint32_t m_maxInstanceCount;
+		// TODO: maybe replace with new readers/writers lock
+		mutable std::mutex m_trackingLock;
+		std::atomic<build_ver_t> m_pendingBuildVer = 0;
+		build_ver_t m_completedBuildVer = 0;
+		core::unordered_set<blas_smart_ptr_t> m_trackedBLASes;
 };
 
 }
diff --git a/include/nbl/video/IGPUCommandBuffer.h b/include/nbl/video/IGPUCommandBuffer.h
@@ -536,6 +536,9 @@ class NBL_API2 IGPUCommandBuffer : public IBackendObject
         //! Secondary CommandBuffer execute
         bool executeCommands(const uint32_t count, IGPUCommandBuffer* const* const cmdbufs);
 
+        // in case you want the commandbuffer to hold onto things as long as its not RESET
+        bool recordReferences(const std::span<const IReferenceCounted*> refs);
+
         virtual bool insertDebugMarker(const char* name, const core::vector4df_SIMD& color = core::vector4df_SIMD(1.0, 1.0, 1.0, 1.0)) = 0;
         virtual bool beginDebugMarker(const char* name, const core::vector4df_SIMD& color = core::vector4df_SIMD(1.0, 1.0, 1.0, 1.0)) = 0;
         virtual bool endDebugMarker() = 0;
@@ -708,6 +711,7 @@ class NBL_API2 IGPUCommandBuffer : public IBackendObject
             m_state = STATE::INITIAL;
 
             m_boundDescriptorSetsRecord.clear();
+            m_TLASToBLASReferenceSets.clear();
             m_boundGraphicsPipeline= nullptr;
             m_boundComputePipeline= nullptr;
             m_boundRayTracingPipeline= nullptr;
@@ -725,6 +729,7 @@ class NBL_API2 IGPUCommandBuffer : public IBackendObject
         {
             deleteCommandList();
             m_boundDescriptorSetsRecord.clear();
+            m_TLASToBLASReferenceSets.clear();
             m_boundGraphicsPipeline= nullptr;
             m_boundComputePipeline= nullptr;
             m_boundRayTracingPipeline= nullptr;
@@ -858,10 +863,17 @@ class NBL_API2 IGPUCommandBuffer : public IBackendObject
         template<typename IndirectCommand> requires nbl::is_any_of_v<IndirectCommand, hlsl::DrawArraysIndirectCommand_t, hlsl::DrawElementsIndirectCommand_t>
         bool invalidDrawIndirectCount(const asset::SBufferBinding<const IGPUBuffer>& indirectBinding, const asset::SBufferBinding<const IGPUBuffer>& countBinding, const uint32_t maxDrawCount, const uint32_t stride);
 
+
         // This bound descriptor set record doesn't include the descriptor sets whose layout has _any_ one of its bindings
         // created with IGPUDescriptorSetLayout::SBinding::E_CREATE_FLAGS::ECF_UPDATE_AFTER_BIND_BIT
         // or IGPUDescriptorSetLayout::SBinding::E_CREATE_FLAGS::ECF_UPDATE_UNUSED_WHILE_PENDING_BIT.
         core::unordered_map<const IGPUDescriptorSet*,uint64_t> m_boundDescriptorSetsRecord;
+
+        // If the user wants the builds to be tracking, and make the TLAS remember the BLASes that have been built into it.
+        // NOTE: We know that a TLAS may be rebuilt multiple times per frame on purpose and not only the final BLASes need to be kept alive till submission finishes.
+        // However, the Command Pool already tracks resources referenced in the Build Infos, so we only need pointers into those records.
+        core::unordered_map<IGPUTopLevelAccelerationStructure*,std::span<const IGPUTopLevelAccelerationStructure::blas_smart_ptr_t>> m_TLASToBLASReferenceSets;
+
         const IGPUGraphicsPipeline* m_boundGraphicsPipeline;
         const IGPUComputePipeline* m_boundComputePipeline;
         const IGPURayTracingPipeline* m_boundRayTracingPipeline;
diff --git a/include/nbl/video/IGPUCommandPool.h b/include/nbl/video/IGPUCommandPool.h
@@ -140,6 +140,7 @@ class IGPUCommandPool : public IBackendObject
         class CBlitImageCmd;
         class CCopyImageToBufferCmd;
         class CExecuteCommandsCmd;
+        class CCustomReferenceCmd;
         class CWaitEventsCmd;
         class CCopyImageCmd;
         class CResolveImageCmd;
@@ -686,6 +687,17 @@ class IGPUCommandPool::CExecuteCommandsCmd final : public IVariableSizeCommand<C
         }
 };
 
+class IGPUCommandPool::CCustomReferenceCmd final : public IVariableSizeCommand<CCustomReferenceCmd>
+{
+    public:
+        CCustomReferenceCmd(const uint32_t count) : IVariableSizeCommand<CCustomReferenceCmd>(count) {}
+
+        static uint32_t calc_resources(const uint32_t count)
+        {
+            return count;
+        }
+};
+
 class IGPUCommandPool::CWaitEventsCmd final : public IVariableSizeCommand<CWaitEventsCmd>
 {
     public:
diff --git a/include/nbl/video/ILogicalDevice.h b/include/nbl/video/ILogicalDevice.h
@@ -539,11 +539,41 @@ class NBL_API2 ILogicalDevice : public core::IReferenceCounted, public IDeviceMe
             // track things created
             if (result==DEFERRABLE_RESULT::DEFERRED)
             {
+                constexpr bool IsTLAS = std::is_same_v<AccelerationStructure,IGPUTopLevelAccelerationStructure>;
+                struct TLASCallback
+                {
+                    // upon completion set the BLASes tracked
+                    inline void operator()(IDeferredOperation*) const
+                    {
+                        for (const auto& set : m_TLASToBLASReferenceSets)
+                        {
+                            auto tlas = set.first;
+                            // we know the build is completed immediately after performing it, so we get our pending stamp then
+                            tlas->setTrackedBLASes(set.second.begin(),set.second.end(),tlas->registerNextBuildVer());
+                        }
+                    }
+
+                    // the rawpointers are already smartpointers in whatever else the `fillTracking` declared above writes
+                    core::unordered_map<IGPUTopLevelAccelerationStructure*,std::span<const IGPUTopLevelAccelerationStructure::blas_smart_ptr_t>> m_TLASToBLASReferenceSets;
+                } callback = {};
+
                 auto& tracking = deferredOperation->m_resourceTracking;
                 tracking.resize(trackingReservation);
                 auto oit = tracking.data();
                 for (const auto& info : infos)
+                {
                     oit = info.fillTracking(oit);
+                    if constexpr (IsTLAS)
+                    {
+                        const auto blasCount = info.trackedBLASes.size();
+                        if (blasCount)
+                            callback.m_TLASToBLASReferenceSets[info.dstAS] = {reinterpret_cast<const IGPUTopLevelAccelerationStructure::blas_smart_ptr_t*>(oit-blasCount),blasCount};
+                        else
+                            callback.m_TLASToBLASReferenceSets[info.dstAS] = {};
+                    }
+                }
+                if constexpr (IsTLAS)
+                    deferredOperation->m_callback = std::move(callback);
             }
             return result!=DEFERRABLE_RESULT::SOME_ERROR;
         }
diff --git a/include/nbl/video/IQueue.h b/include/nbl/video/IQueue.h
@@ -92,6 +92,12 @@ class IQueue : public core::Interface, public core::Unmovable
             std::span<const SSemaphoreInfo> waitSemaphores = {};
             std::span<const SCommandBufferInfo> commandBuffers = {};
             std::span<const SSemaphoreInfo> signalSemaphores = {};
+            // No guarantees are given about when it will execute, except that it will execute:
+            // 1) after the `signalSemaphore.back()` signals
+            // 2) in order w.r.t. all other submits on this queue
+            // 3) after all lifetime tracking has been performed (so transient resources will already be dead!)
+            // NOTE: This `std::function` WILL be copied! 
+            std::function<void()>* completionCallback = nullptr;
 
             inline bool valid() const
             {
@@ -116,45 +122,34 @@ class IQueue : public core::Interface, public core::Unmovable
         virtual const void* getNativeHandle() const = 0;
         
 		// only public because MultiTimelineEventHandlerST needs to know about it
-		class DeferredSubmitResourceDrop final
+		class DeferredSubmitCallback final
 		{
+                //
+                struct STLASBuildMetadata
+                {
+                    core::unordered_set<IGPUTopLevelAccelerationStructure::blas_smart_ptr_t> m_BLASes;
+                    uint32_t m_buildVer;
+                };
+                core::unordered_map<IGPUTopLevelAccelerationStructure*,STLASBuildMetadata> m_TLASToBLASReferenceSets;
+                //
                 using smart_ptr = core::smart_refctd_ptr<IBackendObject>;
-				core::smart_refctd_dynamic_array<smart_ptr> m_resources;
+                core::smart_refctd_dynamic_array<smart_ptr> m_resources;
+                //
+                std::function<void()> m_callback;
 
 			public:
-				inline DeferredSubmitResourceDrop(const SSubmitInfo& info)
-				{
-                    // We could actually not hold any signal semaphore because you're expected to use the signal result somewhere else.
-                    // However it's possible to you might only wait on one from the set and then drop the rest (UB) 
-                    m_resources = core::make_refctd_dynamic_array<decltype(m_resources)>(info.signalSemaphores.size()-1+info.commandBuffers.size()+info.waitSemaphores.size());
-                    auto outRes = m_resources->data();
-                    for (const auto& sema : info.waitSemaphores)
-                        *(outRes++) = smart_ptr(sema.semaphore);
-                    for (const auto& cb : info.commandBuffers)
-                        *(outRes++) = smart_ptr(cb.cmdbuf);
-                    // We don't hold the last signal semaphore, because the timeline does as an Event trigger.
-                    for (auto i=0u; i<info.signalSemaphores.size()-1; i++)
-                        *(outRes++) = smart_ptr(info.signalSemaphores[i].semaphore);
-				}
-                DeferredSubmitResourceDrop(const DeferredSubmitResourceDrop& other) = delete;
-				inline DeferredSubmitResourceDrop(DeferredSubmitResourceDrop&& other) : m_resources(nullptr)
+                DeferredSubmitCallback(const SSubmitInfo& info);
+                DeferredSubmitCallback(const DeferredSubmitCallback& other) = delete;
+				inline DeferredSubmitCallback(DeferredSubmitCallback&& other) : m_resources(nullptr)
 				{
 					this->operator=(std::move(other));
 				}
 
-                DeferredSubmitResourceDrop& operator=(const DeferredSubmitResourceDrop& other) = delete;
-				inline DeferredSubmitResourceDrop& operator=(DeferredSubmitResourceDrop&& other)
-				{
-                    m_resources = std::move(other.m_resources);
-                    other.m_resources = nullptr;
-					return *this;
-				}
+                DeferredSubmitCallback& operator=(const DeferredSubmitCallback& other) = delete;
+                DeferredSubmitCallback& operator=(DeferredSubmitCallback&& other);
 
                 // always exhaustive poll, because we need to get rid of resources ASAP
-                inline void operator()()
-                {
-                    m_resources = nullptr;
-                }
+                void operator()();
 		};
 
     protected:
@@ -170,7 +165,7 @@ class IQueue : public core::Interface, public core::Unmovable
         virtual RESULT waitIdle_impl() const = 0;
 
         // Refcounts all resources used by Pending Submits, gets occasionally cleared out
-        std::unique_ptr<MultiTimelineEventHandlerST<DeferredSubmitResourceDrop,false>> m_submittedResources;
+        std::unique_ptr<MultiTimelineEventHandlerST<DeferredSubmitCallback,false>> m_submittedResources;
         const ILogicalDevice* m_originDevice;
         const uint32_t m_familyIndex;
         const float m_priority;
diff --git a/src/nbl/video/IGPUCommandBuffer.cpp b/src/nbl/video/IGPUCommandBuffer.cpp
@@ -816,6 +816,7 @@ uint32_t IGPUCommandBuffer::buildAccelerationStructures_common(const std::span<c
 
     if (indirectBuffer)
     {
+        // TODO: maybe hoist the check
         if (!features.accelerationStructureIndirectBuild)
         {
             NBL_LOG_ERROR("'accelerationStructureIndirectBuild' feature not enabled!");
@@ -835,7 +836,18 @@ uint32_t IGPUCommandBuffer::buildAccelerationStructures_common(const std::span<c
     if (indirectBuffer)
         *(oit++) = core::smart_refctd_ptr<const IGPUBuffer>(indirectBuffer);
     for (const auto& info : infos)
+    {
         oit = info.fillTracking(oit);
+        // we still need to clear the BLAS tracking list if the TLAS has nothing to track
+        if constexpr (std::is_same_v<DeviceBuildInfo,IGPUTopLevelAccelerationStructure::DeviceBuildInfo>)
+        {
+            const auto blasCount = info.trackedBLASes.size();
+            if (blasCount)
+                m_TLASToBLASReferenceSets[info.dstAS] = {reinterpret_cast<const IGPUTopLevelAccelerationStructure::blas_smart_ptr_t*>(oit-blasCount),blasCount};
+            else
+                m_TLASToBLASReferenceSets[info.dstAS] = {};
+        }
+    }
 
     return totalGeometries;
 }
@@ -2066,4 +2078,22 @@ bool IGPUCommandBuffer::executeCommands(const uint32_t count, IGPUCommandBuffer*
     return executeCommands_impl(count,cmdbufs);
 }
 
+bool IGPUCommandBuffer::recordReferences(const std::span<const IReferenceCounted*> refs)
+{
+    if (!checkStateBeforeRecording(queue_flags_t::COMPUTE_BIT|queue_flags_t::GRAPHICS_BIT|queue_flags_t::TRANSFER_BIT|queue_flags_t::SPARSE_BINDING_BIT))
+        return false;
+    
+    auto cmd = m_cmdpool->m_commandListPool.emplace<IGPUCommandPool::CCustomReferenceCmd>(m_commandList,refs.size());
+    if (!cmd)
+    {
+        NBL_LOG_ERROR("out of host memory!");
+        return false;
+    }
+    auto oit = cmd->getVariableCountResources();
+    for (const auto& ref : refs)
+        *(oit++) = core::smart_refctd_ptr<const core::IReferenceCounted>(ref);
+
+    return true;
+}
+
 }
diff --git a/src/nbl/video/IQueue.cpp b/src/nbl/video/IQueue.cpp
