make the Host Acceleration Structure builds set

Author: devsh

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;
 };
 

include/nbl/video/IGPUCommandBuffer.h

@@ -872,7 +872,7 @@ class NBL_API2 IGPUCommandBuffer : public IBackendObject
         // 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<const IGPUTopLevelAccelerationStructure*,const IGPUTopLevelAccelerationStructure::blas_smart_ptr_t*> m_TLASToBLASReferenceSets;
+        core::unordered_map<IGPUTopLevelAccelerationStructure*,std::span<const IGPUTopLevelAccelerationStructure::blas_smart_ptr_t>> m_TLASToBLASReferenceSets;
 
         const IGPUGraphicsPipeline* m_boundGraphicsPipeline;
         const IGPUComputePipeline* m_boundComputePipeline;

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;
         }

src/nbl/video/IGPUCommandBuffer.cpp

@@ -840,7 +840,13 @@ uint32_t IGPUCommandBuffer::buildAccelerationStructures_common(const std::span<c
         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>)
-            m_TLASToBLASReferenceSets[info.dstAS] = info.trackedBLASes.empty() ? nullptr:reinterpret_cast<const IGPUTopLevelAccelerationStructure::blas_smart_ptr_t*>(oit-info.trackedBLASes.size());
+        {
+            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;