Fix memory allocation of acceleration structures

AnKi/Resource/AccelerationStructureScratchAllocator.h

@@ -0,0 +1,71 @@
+// Copyright (C) 2009-present, Panagiotis Christopoulos Charitos and contributors.
+// All rights reserved.
+// Code licensed under the BSD License.
+// http://www.anki3d.org/LICENSE
+
+#pragma once
+
+#include <AnKi/Resource/Common.h>
+#include <AnKi/Gr/Buffer.h>
+#include <AnKi/Gr/GrManager.h>
+
+namespace anki {
+
+/// @addtogroup resource
+/// @{
+
+/// Ring buffer used for AS scratch memory.
+class AccelerationStructureScratchAllocator
+{
+public:
+	static constexpr PtrSize kBufferSize = 64_MB;
+	static constexpr PtrSize kMaxBufferSize = 256_MB;
+
+	AccelerationStructureScratchAllocator()
+	{
+		BufferInitInfo buffInit("BLAS scratch");
+		buffInit.m_size = kBufferSize;
+		buffInit.m_usage = BufferUsageBit::kAccelerationStructureBuildScratch;
+		m_buffer = GrManager::getSingleton().newBuffer(buffInit);
+	}
+
+	BufferView allocate(PtrSize size, Bool& addBarrierBefore)
+	{
+		if(size > kMaxBufferSize)
+		{
+			ANKI_RESOURCE_LOGF("Asked for too much BLAS scratch memory: %llu", size);
+		}
+
+		if(size > m_buffer->getSize())
+		{
+			BufferInitInfo buffInit("BLAS scratch");
+			buffInit.m_size = min(kMaxBufferSize, m_buffer->getSize() * 2u);
+			buffInit.m_usage = BufferUsageBit::kAccelerationStructureBuildScratch;
+			m_buffer = GrManager::getSingleton().newBuffer(buffInit);
+
+			m_offset = 0;
+		}
+
+		if(m_offset + size > m_buffer->getSize())
+		{
+			m_offset = 0;
+			addBarrierBefore = true;
+		}
+		else
+		{
+			addBarrierBefore = false;
+		}
+
+		const BufferView view(m_buffer.get(), m_offset, size);
+		m_offset += size;
+
+		return view;
+	}
+
+private:
+	BufferPtr m_buffer;
+	PtrSize m_offset = 0;
+};
+/// @}
+
+} // end namespace anki

AnKi/Resource/MeshResource.cpp

@@ -7,6 +7,7 @@
 #include <AnKi/Resource/ResourceManager.h>
 #include <AnKi/Resource/MeshBinaryLoader.h>
 #include <AnKi/Resource/AsyncLoader.h>
+#include <AnKi/Resource/AccelerationStructureScratchAllocator.h>
 #include <AnKi/Util/Functions.h>
 #include <AnKi/Util/Filesystem.h>
 #include <AnKi/Core/App.h>
@@ -389,13 +390,20 @@ Error MeshResource::loadAsync(MeshBinaryLoader& loader) const
 			// Build BLASes
 			for(U32 lodIdx = 0; lodIdx < m_lods.getSize(); ++lodIdx)
 			{
-				// TODO find a temp buffer
-				BufferInitInfo buffInit("BLAS scratch");
-				buffInit.m_size = submesh.m_blas[lodIdx]->getBuildScratchBufferSize();
-				buffInit.m_usage = BufferUsageBit::kAccelerationStructureBuildScratch;
-				BufferPtr scratchBuff = GrManager::getSingleton().newBuffer(buffInit);
+				Bool addBarrier;
+				const BufferView scratchBuff = ResourceManager::getSingleton().getAccelerationStructureScratchAllocator().allocate(
+					submesh.m_blas[lodIdx]->getBuildScratchBufferSize(), addBarrier);
 
-				cmdb->buildAccelerationStructure(submesh.m_blas[lodIdx].get(), BufferView(scratchBuff.get()));
+				if(addBarrier)
+				{
+					BufferBarrierInfo barr;
+					barr.m_bufferView = scratchBuff;
+					barr.m_previousUsage = BufferUsageBit::kAccelerationStructureBuildScratch;
+					barr.m_previousUsage = BufferUsageBit::kAccelerationStructureBuildScratch;
+					cmdb->setPipelineBarrier({}, {&barr, 1}, {});
+				}
+
+				cmdb->buildAccelerationStructure(submesh.m_blas[lodIdx].get(), scratchBuff);
 			}
 
 			// Barriers again

AnKi/Resource/ResourceManager.cpp

@@ -7,6 +7,7 @@
 #include <AnKi/Resource/AsyncLoader.h>
 #include <AnKi/Resource/ShaderProgramResourceSystem.h>
 #include <AnKi/Resource/AnimationResource.h>
+#include <AnKi/Resource/AccelerationStructureScratchAllocator.h>
 #include <AnKi/Util/Logger.h>
 #include <AnKi/Util/CVarSet.h>
 
@@ -43,6 +44,8 @@ ResourceManager::~ResourceManager()
 	static_cast<TypeData<className>&>(m_allTypes).m_map.destroy();
 #include <AnKi/Resource/Resources.def.h>
 
+	deleteInstance(ResourceMemoryPool::getSingleton(), m_asScratchAlloc);
+
 	ResourceMemoryPool::freeSingleton();
 }
 
@@ -65,6 +68,11 @@ Error ResourceManager::init(AllocAlignedCallback allocCallback, void* allocCallb
 	m_shaderProgramSystem = newInstance<ShaderProgramResourceSystem>(ResourceMemoryPool::getSingleton());
 	ANKI_CHECK(m_shaderProgramSystem->init());
 
+	if(GrManager::getSingleton().getDeviceCapabilities().m_rayTracingEnabled)
+	{
+		m_asScratchAlloc = newInstance<AccelerationStructureScratchAllocator>(ResourceMemoryPool::getSingleton());
+	}
+
 	return Error::kNone;
 }
 

AnKi/Resource/ResourceManager.h

@@ -23,6 +23,7 @@ class AsyncLoader;
 class ResourceManagerModel;
 class ShaderCompilerCache;
 class ShaderProgramResourceSystem;
+class AccelerationStructureScratchAllocator;
 
 /// @addtogroup resource
 /// @{
@@ -75,6 +76,12 @@ public:
 		return *m_fs;
 	}
 
+	ANKI_INTERNAL AccelerationStructureScratchAllocator& getAccelerationStructureScratchAllocator()
+	{
+		ANKI_ASSERT(m_asScratchAlloc);
+		return *m_asScratchAlloc;
+	}
+
 private:
 	template<typename Type>
 	class TypeData
@@ -113,6 +120,7 @@ public \
 	AsyncLoader* m_asyncLoader = nullptr; ///< Async loading thread
 	ShaderProgramResourceSystem* m_shaderProgramSystem = nullptr;
 	TransferGpuAllocator* m_transferGpuAlloc = nullptr;
+	AccelerationStructureScratchAllocator* m_asScratchAlloc = nullptr;
 
 	AllTypeData m_allTypes;
 

AnKi/Shaders/GBufferGeneric.ankiprog

@@ -685,7 +685,7 @@ GBufferPixelOut main(
 // ===========================================================================
 #if ANKI_TECHNIQUE_RtMaterialFetch
 
-#	define USE_POSITION_FETCH_NORMALS ANKI_GR_BACKEND_VULKAN && ANKI_PLATFORM_MOBILE
+#	define USE_POSITION_FETCH_NORMALS ANKI_GR_BACKEND_VULKAN&& ANKI_PLATFORM_MOBILE
 
 #	if ANKI_CLOSEST_HIT_SHADER
 [[vk::shader_record_ext]] ConstantBuffer<GpuSceneRenderableInstance> g_gpuSceneRenderable : register(b0); // TODO that won't work on D3D