BLASes referenced during a TLAS build (if provided) are recorded into the buildAccelerationStructures command in the pool.

Now just time for callbacks.

@Crisspl I also added that feature you wanted `bool IGPUCommandBuffer::recordReferences(const std::span<const IReferenceCounted*> refs);`

Author: devsh

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,7 +475,7 @@ 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
+				// [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>;
@@ -550,7 +555,7 @@ class IGPUTopLevelAccelerationStructure : public asset::ITopLevelAccelerationStr
 		//
 		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 nextBuildVer()
+		inline build_ver_t registerNextBuildVer()
 		{
 			return m_pendingBuildVer++;
 		}
@@ -595,6 +600,11 @@ class IGPUTopLevelAccelerationStructure : public asset::ITopLevelAccelerationStr
 			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)

include/nbl/video/IGPUCommandBuffer.h

@@ -267,9 +267,7 @@ class NBL_API2 IGPUCommandBuffer : public IBackendObject
         inline bool buildAccelerationStructures(const std::span<const IGPUTopLevelAccelerationStructure::DeviceBuildInfo> infos, const IGPUTopLevelAccelerationStructure::DirectBuildRangeRangeInfos buildRangeInfos)
         {
             if (buildAccelerationStructures_common(infos,buildRangeInfos))
-            {
                 return buildAccelerationStructures_impl(infos,buildRangeInfos);
-            }
             return false;
         }
         // We don't allow different indirect command addresses due to https://registry.khronos.org/vulkan/specs/1.3-extensions/html/vkspec.html#VUID-vkCmdBuildAccelerationStructuresIndirectKHR-pIndirectDeviceAddresses-03646
@@ -305,9 +303,7 @@ class NBL_API2 IGPUCommandBuffer : public IBackendObject
                 if constexpr(std::is_same_v<AccelerationStructure,IGPUBottomLevelAccelerationStructure>)
                     return buildAccelerationStructuresIndirect_impl(indirectRangeBuffer,infos,pIndirectOffsets,pIndirectStrides,maxPrimitiveOrInstanceCounts,totalGeometryCount);
                 else
-                {
                     return buildAccelerationStructuresIndirect_impl(indirectRangeBuffer,infos,pIndirectOffsets,pIndirectStrides,maxPrimitiveOrInstanceCounts);
-                }
             }
             return false;
         }
@@ -540,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;
@@ -712,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;
@@ -729,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;
@@ -862,13 +863,16 @@ 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
-        core::vector<core::unordered_set<core::smart_refctd_ptr<const IGPUBottomLevelAccelerationStructure>>> m_TLASToBLASReferenceSets;
+        // 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;
 
         const IGPUGraphicsPipeline* m_boundGraphicsPipeline;
         const IGPUComputePipeline* m_boundComputePipeline;

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:

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,12 @@ 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>)
+            m_TLASToBLASReferenceSets[info.dstAS] = info.trackedBLASes.empty() ? nullptr:reinterpret_cast<const IGPUTopLevelAccelerationStructure::blas_smart_ptr_t*>(oit-info.trackedBLASes.size());
+    }
 
     return totalGeometries;
 }
@@ -2066,4 +2072,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;
+}
+
 }