Continue work on the acceleration structures

src/anki/gr/AccelerationStructure.h

@@ -5,7 +5,9 @@
 
 #pragma once
 
-#include <anki/gr/GrObject.h>
+#include <anki/gr/Buffer.h>
+#include <anki/Math.h>
+#include <anki/util/WeakArray.h>
 
 namespace anki
 {
@@ -13,27 +15,75 @@ namespace anki
 /// @addtogroup graphics
 /// @{
 
+/// @memberof AccelerationStructureInitInfo
+class BottomLevelAccelerationStructureInitInfo
+{
+public:
+	BufferPtr m_indexBuffer;
+	PtrSize m_indexBufferOffset = 0;
+	U32 m_indexCount = 0;
+	IndexType m_indexType = IndexType::COUNT;
+
+	BufferPtr m_positionBuffer;
+	PtrSize m_positionBufferOffset = 0;
+	U32 m_positionStride = 0;
+	Format m_positionsFormat = Format::NONE;
+	U32 m_positionCount = 0;
+
+	Bool isValid() const
+	{
+		if(m_indexBuffer.get() == nullptr || m_indexCount == 0 || m_indexType == IndexType::COUNT
+			|| m_positionBuffer.get() == nullptr || m_positionStride == 0 || m_positionsFormat == Format::NONE
+			|| m_positionCount == 0)
+		{
+			return false;
+		}
+
+		const PtrSize posRange = m_positionBufferOffset + m_positionStride * m_positionCount;
+		if(m_positionStride < getFormatBytes(m_positionsFormat)
+			|| (m_positionStride % getFormatBytes(m_positionsFormat)) != 0 || posRange > m_positionBuffer->getSize())
+		{
+			return false;
+		}
+
+		const PtrSize idxStride = (m_indexType == IndexType::U16) ? 2 : 4;
+		if(m_indexBufferOffset + idxStride * m_indexCount > m_indexBuffer->getSize())
+		{
+			return false;
+		}
+
+		return true;
+	}
+};
+
+/// @memberof AccelerationStructureInitInfo
+class AccelerationStructureInstance
+{
+public:
+	AccelerationStructurePtr m_bottomLevel;
+	Mat3x4 m_transform = Mat3x4::getIdentity();
+};
+
+/// @memberof AccelerationStructureInitInfo
+class TopLevelAccelerationStructureInitInfo
+{
+public:
+	ConstWeakArray<AccelerationStructureInstance> m_instances;
+
+	Bool isValid() const
+	{
+		return m_instances.getSize() > 0;
+	}
+};
+
 /// Acceleration struture init info.
 /// @memberof AccelerationStructure
 class AccelerationStructureInitInfo : public GrBaseInitInfo
 {
 public:
 	AccelerationStructureType m_type = AccelerationStructureType::COUNT;
-
-	class
-	{
-	public:
-		IndexType m_indexType = IndexType::COUNT;
-		Format m_positionsFormat = Format::NONE;
-		U32 m_indexCount = 0;
-		U32 m_vertexCount = 0;
-	} m_bottomLevel;
-
-	class
-	{
-	public:
-		U32 m_bottomLevelCount = 0;
-	} m_topLevel;
+	BottomLevelAccelerationStructureInitInfo m_bottomLevel;
+	TopLevelAccelerationStructureInitInfo m_topLevel;
 
 	AccelerationStructureInitInfo(CString name = {})
 		: GrBaseInitInfo(name)
@@ -47,23 +97,7 @@ public:
 			return false;
 		}
 
-		if(m_type == AccelerationStructureType::BOTTOM_LEVEL)
-		{
-			if(m_bottomLevel.m_indexType == IndexType::COUNT || m_bottomLevel.m_positionsFormat == Format::NONE
-				|| m_bottomLevel.m_indexCount == 0 || m_bottomLevel.m_vertexCount == 0)
-			{
-				return false;
-			}
-		}
-		else
-		{
-			if(m_topLevel.m_bottomLevelCount == 0)
-			{
-				return false;
-			}
-		}
-
-		return true;
+		return (m_type == AccelerationStructureType::BOTTOM_LEVEL) ? m_bottomLevel.isValid() : m_topLevel.isValid();
 	}
 };
 
@@ -75,6 +109,12 @@ class AccelerationStructure : public GrObject
 public:
 	static const GrObjectType CLASS_TYPE = GrObjectType::ACCELERATION_STRUCTURE;
 
+	AccelerationStructureType getType() const
+	{
+		ANKI_ASSERT(m_type != AccelerationStructureType::COUNT);
+		return m_type;
+	}
+
 protected:
 	AccelerationStructureType m_type = AccelerationStructureType::COUNT;
 

src/anki/gr/CommandBuffer.h

@@ -386,19 +386,8 @@ public:
 	/// @param range Size to copy.
 	void copyBufferToBuffer(BufferPtr src, PtrSize srcOffset, BufferPtr dst, PtrSize dstOffset, PtrSize range);
 
-	/// Build the bottom level acceleration structure.
-	/// @param[in,out] as The AS to build.
-	/// @param[in] positions The position buffer.
-	/// @param positionsOffset Offset to the positions buffer.
-	/// @param positionsStride Stride of the positions.
-	/// @param[in] indices The indices buffer.
-	/// @param indicesOffset The offset of the 1st index.
-	void buildBottomLevelAccelerationStructure(AccelerationStructurePtr as,
-		BufferPtr positions,
-		PtrSize positionsOffset,
-		PtrSize positionsStride,
-		BufferPtr indices,
-		PtrSize indicesOffset);
+	/// Build the acceleration structure.
+	void buildAccelerationStructure(AccelerationStructurePtr as);
 	/// @}
 
 	/// @name Sync

src/anki/gr/vulkan/AccelerationStructureImpl.cpp

@@ -29,14 +29,15 @@ Error AccelerationStructureImpl::init(const AccelerationStructureInitInfo& inf)
 
 	if(m_type == AccelerationStructureType::BOTTOM_LEVEL)
 	{
-		m_bottomLevelInfo.m_indexType = inf.m_bottomLevel.m_indexType;
-		m_bottomLevelInfo.m_positionsFormat = inf.m_bottomLevel.m_positionsFormat;
-		m_bottomLevelInfo.m_indexCount = inf.m_bottomLevel.m_indexCount;
-		m_bottomLevelInfo.m_vertexCount = inf.m_bottomLevel.m_vertexCount;
+		static_cast<BottomLevelAccelerationStructureInitInfo&>(m_bottomLevelInfo) = inf.m_bottomLevel;
 	}
 	else
 	{
-		m_topLevelInfo.m_bottomLevelCount = inf.m_topLevel.m_bottomLevelCount;
+		m_topLevelInfo.m_instances.create(getAllocator(), inf.m_topLevel.m_instances.getSize());
+		for(U32 i = 0; i < inf.m_topLevel.m_instances.getSize(); ++i)
+		{
+			m_topLevelInfo.m_instances[i] = inf.m_topLevel.m_instances[i];
+		}
 	}
 
 	// Create the handle
@@ -48,19 +49,19 @@ Error AccelerationStructureImpl::init(const AccelerationStructureInitInfo& inf)
 		if(m_type == AccelerationStructureType::BOTTOM_LEVEL)
 		{
 			geom.geometryType = VK_GEOMETRY_TYPE_TRIANGLES_KHR;
-			geom.indexType = convertIndexType(inf.m_bottomLevel.m_indexType);
-			geom.vertexFormat = convertFormat(inf.m_bottomLevel.m_positionsFormat);
-			geom.maxPrimitiveCount = inf.m_bottomLevel.m_indexCount / 3;
-			geom.maxVertexCount = inf.m_bottomLevel.m_vertexCount;
-			geom.allowsTransforms = false;
+			geom.indexType = convertIndexType(m_bottomLevelInfo.m_indexType);
+			geom.vertexFormat = convertFormat(m_bottomLevelInfo.m_positionsFormat);
+			geom.maxPrimitiveCount = m_bottomLevelInfo.m_indexCount / 3;
+			geom.maxVertexCount = m_bottomLevelInfo.m_positionCount;
 		}
 		else
 		{
 			geom.geometryType = VK_GEOMETRY_TYPE_INSTANCES_KHR;
-			geom.maxPrimitiveCount = inf.m_topLevel.m_bottomLevelCount;
-			geom.allowsTransforms = true;
+			geom.maxPrimitiveCount = m_topLevelInfo.m_instances.getSize();
 		}
 
+		geom.allowsTransforms = false; // Only for triangle types and at the moment not used
+
 		VkAccelerationStructureCreateInfoKHR ci{};
 		ci.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_CREATE_INFO_KHR;
 
@@ -116,7 +117,114 @@ Error AccelerationStructureImpl::init(const AccelerationStructureInitInfo& inf)
 		ANKI_VK_CHECK(vkBindAccelerationStructureMemoryKHR(getDevice(), 1, &bindInfo));
 	}
 
+	// Get scratch buffer size
+	{
+		VkMemoryRequirements2 req{};
+		req.sType = VK_STRUCTURE_TYPE_MEMORY_REQUIREMENTS_2;
+
+		VkAccelerationStructureMemoryRequirementsInfoKHR inf{};
+		inf.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_MEMORY_REQUIREMENTS_INFO_KHR;
+		inf.type = VK_ACCELERATION_STRUCTURE_MEMORY_REQUIREMENTS_TYPE_BUILD_SCRATCH_KHR;
+		inf.buildType = VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR;
+		inf.accelerationStructure = m_handle;
+		vkGetAccelerationStructureMemoryRequirementsKHR(getDevice(), &inf, &req);
+
+		m_scratchBufferSize = U32(req.memoryRequirements.size);
+	}
+
+	// Get GPU address
+	if(m_type == AccelerationStructureType::BOTTOM_LEVEL)
+	{
+		VkAccelerationStructureDeviceAddressInfoKHR inf{};
+		inf.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_DEVICE_ADDRESS_INFO_KHR;
+		inf.accelerationStructure = m_handle;
+
+		m_bottomLevelInfo.m_gpuAddress = vkGetAccelerationStructureDeviceAddressKHR(getDevice(), &inf);
+		ANKI_ASSERT(m_bottomLevelInfo.m_gpuAddress);
+	}
+
 	return Error::NONE;
 }
 
+void AccelerationStructureImpl::initBuildInfo()
+{
+	if(m_type == AccelerationStructureType::BOTTOM_LEVEL)
+	{
+		m_geometry.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR;
+		m_geometry.geometryType = VK_GEOMETRY_TYPE_TRIANGLES_KHR;
+		m_geometry.flags = VK_GEOMETRY_OPAQUE_BIT_KHR; // TODO
+
+		VkAccelerationStructureGeometryTrianglesDataKHR& triangles = m_geometry.geometry.triangles;
+		triangles.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR;
+		triangles.vertexFormat = convertFormat(m_bottomLevelInfo.m_positionsFormat);
+		triangles.vertexData.deviceAddress =
+			m_bottomLevelInfo.m_positionBuffer->getGpuAddress() + m_bottomLevelInfo.m_positionBufferOffset;
+		triangles.vertexStride = m_bottomLevelInfo.m_positionStride;
+		triangles.indexType = convertIndexType(m_bottomLevelInfo.m_indexType);
+		triangles.indexData.deviceAddress =
+			m_bottomLevelInfo.m_indexBuffer->getGpuAddress() + m_bottomLevelInfo.m_indexBufferOffset;
+	}
+	else
+	{
+		const U32 instanceCount = m_topLevelInfo.m_instances.getSize();
+
+		// Create the instances buffer
+		{
+			BufferInitInfo buffInit("RT_instances");
+			buffInit.m_size = sizeof(VkAccelerationStructureInstanceKHR) * instanceCount;
+			buffInit.m_usage = BufferImpl::ACCELERATION_STRUCTURE_BUILD_SCRATCH_USAGE;
+			buffInit.m_exposeGpuAddress = true;
+			buffInit.m_access = BufferMapAccessBit::WRITE;
+			m_topLevelInfo.m_instancesBuff = getManager().newBuffer(buffInit);
+		}
+
+		// Populate the instances buffer
+		{
+			VkAccelerationStructureInstanceKHR* instances = static_cast<VkAccelerationStructureInstanceKHR*>(
+				m_topLevelInfo.m_instancesBuff->map(0, MAX_PTR_SIZE, BufferMapAccessBit::WRITE));
+			for(U32 i = 0; i < instanceCount; ++i)
+			{
+				VkAccelerationStructureInstanceKHR& outInst = instances[i];
+				const AccelerationStructureInstance& inInst = m_topLevelInfo.m_instances[i];
+				static_assert(sizeof(outInst.transform) == sizeof(inInst.m_transform), "See file");
+				memcpy(&outInst.transform.matrix[0][0], &inInst.m_transform, sizeof(inInst.m_transform));
+				outInst.instanceCustomIndex = i;
+				outInst.mask = 0xFF;
+				outInst.instanceShaderBindingTableRecordOffset = 0;
+				outInst.flags = VK_GEOMETRY_INSTANCE_TRIANGLE_FRONT_COUNTERCLOCKWISE_BIT_KHR;
+				outInst.accelerationStructureReference =
+					static_cast<const AccelerationStructureImpl&>(*inInst.m_bottomLevel).m_bottomLevelInfo.m_gpuAddress;
+				ANKI_ASSERT(outInst.accelerationStructureReference != 0);
+			}
+
+			m_topLevelInfo.m_instancesBuff->unmap();
+		}
+
+		// TODO geometry
+	}
+
+	m_buildInfo.sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR;
+	m_buildInfo.type = convertAccelerationStructureType(m_type);
+	m_buildInfo.flags = VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR;
+	m_buildInfo.update = false;
+	m_buildInfo.dstAccelerationStructure = m_handle;
+	m_buildInfo.geometryArrayOfPointers = false;
+	m_buildInfo.geometryCount = 1;
+	m_buildInfo.ppGeometries = &m_geometryPtr;
+	m_buildInfo.scratchData.deviceAddress = MAX_U64;
+
+	// Ofset info
+	m_offsetInfo.firstVertex = 0;
+	if(m_type == AccelerationStructureType::BOTTOM_LEVEL)
+	{
+		m_offsetInfo.primitiveCount = m_bottomLevelInfo.m_indexCount / 3;
+	}
+	else
+	{
+		m_offsetInfo.primitiveCount = m_topLevelInfo.m_instances.getSize();
+	}
+	m_offsetInfo.primitiveOffset = 0;
+	m_offsetInfo.transformOffset = 0;
+}
+
 } // end namespace anki

src/anki/gr/vulkan/AccelerationStructureImpl.h

@@ -15,23 +15,6 @@ namespace anki
 /// @addtogroup vulkan
 /// @{
 
-/// @memberof AccelerationStructureImpl
-class ASBottomLevelInfo
-{
-public:
-	IndexType m_indexType = IndexType::COUNT;
-	Format m_positionsFormat = Format::NONE;
-	U32 m_indexCount = 0;
-	U32 m_vertexCount = 0;
-};
-
-/// @memberof AccelerationStructureImpl
-class ASTopLevelInfo
-{
-public:
-	U32 m_bottomLevelCount = 0;
-};
-
 /// AccelerationStructure implementation.
 class AccelerationStructureImpl final : public AccelerationStructure,
 										public VulkanObject<AccelerationStructure, AccelerationStructureImpl>
@@ -52,27 +35,49 @@ public:
 		return m_handle;
 	}
 
-	const ASBottomLevelInfo& getBottomLevelInfo() const
+	U32 getBuildScratchBufferSize() const
 	{
-		ANKI_ASSERT(m_type == AccelerationStructureType::BOTTOM_LEVEL);
-		return m_bottomLevelInfo;
+		ANKI_ASSERT(m_scratchBufferSize > 0);
+		return m_scratchBufferSize;
 	}
 
-	const ASTopLevelInfo& getTopLevelInfo() const
+	void generateBuildInfo(U64 scratchBufferAddress,
+		VkAccelerationStructureBuildGeometryInfoKHR& info,
+		VkAccelerationStructureBuildOffsetInfoKHR& offsetInfo) const
 	{
-		ANKI_ASSERT(m_type == AccelerationStructureType::TOP_LEVEL);
-		return m_topLevelInfo;
+		info = m_buildInfo;
+		info.scratchData.deviceAddress = scratchBufferAddress;
+		offsetInfo = m_offsetInfo;
 	}
 
 private:
-	VkAccelerationStructureKHR m_handle = VK_NULL_HANDLE;
-	GpuMemoryHandle m_memHandle;
+	class ASBottomLevelInfo : public BottomLevelAccelerationStructureInitInfo
+	{
+	public:
+		U64 m_gpuAddress = 0;
+	};
 
-	union
+	class ASTopLevelInfo
 	{
-		ASBottomLevelInfo m_bottomLevelInfo;
-		ASTopLevelInfo m_topLevelInfo;
+	public:
+		DynamicArray<AccelerationStructureInstance> m_instances;
+		BufferPtr m_instancesBuff;
 	};
+
+	VkAccelerationStructureKHR m_handle = VK_NULL_HANDLE;
+	GpuMemoryHandle m_memHandle;
+
+	ASBottomLevelInfo m_bottomLevelInfo;
+	ASTopLevelInfo m_topLevelInfo;
+
+	U32 m_scratchBufferSize = 0;
+
+	VkAccelerationStructureBuildGeometryInfoKHR m_buildInfo{};
+	VkAccelerationStructureGeometryKHR m_geometry{};
+	VkAccelerationStructureGeometryKHR* m_geometryPtr = &m_geometry;
+	VkAccelerationStructureBuildOffsetInfoKHR m_offsetInfo{};
+
+	void initBuildInfo();
 };
 /// @}
 

src/anki/gr/vulkan/BufferImpl.h

@@ -19,6 +19,8 @@ namespace anki
 class BufferImpl final : public Buffer, public VulkanObject<Buffer, BufferImpl>
 {
 public:
+	static constexpr BufferUsageBit ACCELERATION_STRUCTURE_BUILD_SCRATCH_USAGE = BufferUsageBit(1ull << 63ull);
+
 	BufferImpl(GrManager* manager, CString name)
 		: Buffer(manager, name)
 	{

src/anki/gr/vulkan/CommandBuffer.cpp

@@ -6,6 +6,7 @@
 #include <anki/gr/CommandBuffer.h>
 #include <anki/gr/vulkan/CommandBufferImpl.h>
 #include <anki/gr/vulkan/GrManagerImpl.h>
+#include <anki/gr/AccelerationStructure.h>
 
 namespace anki
 {
@@ -325,16 +326,10 @@ void CommandBuffer::copyBufferToBuffer(
 	self.copyBufferToBuffer(src, srcOffset, dst, dstOffset, range);
 }
 
-void CommandBuffer::buildBottomLevelAccelerationStructure(AccelerationStructurePtr as,
-	BufferPtr positions,
-	PtrSize positionsOffset,
-	PtrSize positionsStride,
-	BufferPtr indices,
-	PtrSize indicesOffset)
+void CommandBuffer::buildAccelerationStructure(AccelerationStructurePtr as)
 {
 	ANKI_VK_SELF(CommandBufferImpl);
-	self.buildBottomLevelAccelerationStructureInternal(
-		as, positions, positionsOffset, positionsStride, indices, indicesOffset);
+	self.buildAccelerationStructureInternal(as);
 }
 
 void CommandBuffer::setTextureBarrier(

src/anki/gr/vulkan/CommandBufferImpl.cpp

@@ -813,26 +813,30 @@ void CommandBufferImpl::rebindDynamicState()
 	}
 }
 
-void CommandBufferImpl::buildBottomLevelAccelerationStructureInternal(AccelerationStructurePtr& as,
-	BufferPtr& positions,
-	PtrSize positionsOffset,
-	PtrSize positionsStride,
-	BufferPtr& indices,
-	PtrSize indicesOffset)
+void CommandBufferImpl::buildAccelerationStructureInternal(AccelerationStructurePtr& as)
 {
 	// Get objects
-	AccelerationStructureImpl& asImpl = static_cast<AccelerationStructureImpl&>(*as);
-	const BufferImpl& positionsImpl = static_cast<const BufferImpl&>(*positions);
-	const BufferImpl& indicesImpl = static_cast<const BufferImpl&>(*indices);
-
-	// Do checks
-#if ANKI_ENABLE_ASSERTS
-	{
-		const ASBottomLevelInfo& bottomInfo = asImpl.getBottomLevelInfo();
-	}
-#endif
-
-	ANKI_ASSERT(!"TODO");
+	const AccelerationStructureImpl& asImpl = static_cast<AccelerationStructureImpl&>(*as);
+
+	// Create the scrach buffer
+	BufferInitInfo bufferInit;
+	bufferInit.m_exposeGpuAddress = true;
+	bufferInit.m_usage = BufferImpl::ACCELERATION_STRUCTURE_BUILD_SCRATCH_USAGE;
+	bufferInit.m_size = asImpl.getBuildScratchBufferSize();
+	BufferPtr scratchBuff = getManager().newBuffer(bufferInit);
+
+	// Create the build info
+	VkAccelerationStructureBuildGeometryInfoKHR buildInfo;
+	VkAccelerationStructureBuildOffsetInfoKHR offsetInfo;
+	asImpl.generateBuildInfo(scratchBuff->getGpuAddress(), buildInfo, offsetInfo);
+
+	// Do the command
+	VkAccelerationStructureBuildOffsetInfoKHR* offsetInfoPtr = &offsetInfo;
+	ANKI_CMD(vkCmdBuildAccelerationStructureKHR(m_handle, 1, &buildInfo, &offsetInfoPtr), ANY_OTHER_COMMAND);
+
+	// Push refs
+	m_microCmdb->pushObjectRef(as);
+	m_microCmdb->pushObjectRef(scratchBuff);
 }
 
 } // end namespace anki

src/anki/gr/vulkan/CommandBufferImpl.h

@@ -372,12 +372,7 @@ public:
 
 	void copyBufferToBuffer(BufferPtr& src, PtrSize srcOffset, BufferPtr& dst, PtrSize dstOffset, PtrSize range);
 
-	void buildBottomLevelAccelerationStructureInternal(AccelerationStructurePtr& as,
-		BufferPtr& positions,
-		PtrSize positionsOffset,
-		PtrSize positionsStride,
-		BufferPtr& indices,
-		PtrSize indicesOffset);
+	void buildAccelerationStructureInternal(AccelerationStructurePtr& as);
 
 	void setPushConstants(const void* data, U32 dataSize);