Add SBT support in the vulkan backend. Ray tracing should be feature complete for the Vulkan backend

src/anki/gr/AccelerationStructure.h

@@ -67,7 +67,7 @@ class AccelerationStructureInstance
 public:
 	AccelerationStructurePtr m_bottomLevel;
 	Mat3x4 m_transform = Mat3x4::getIdentity();
-	U32 m_stbRecordIndex = 0; ///< Points to the STB record.
+	U32 m_sbtRecordIndex = 0; ///< Points to the SBT record.
 };
 
 /// @memberof AccelerationStructureInitInfo

src/anki/gr/CommandBuffer.h

@@ -335,6 +335,16 @@ public:
 
 	void dispatchCompute(U32 groupCountX, U32 groupCountY, U32 groupCountZ);
 
+	/// Trace rays.
+	/// @param[in] sbtBuffer The SBT buffer.
+	/// @param sbtBufferOffset Offset inside the sbtBuffer where SBT records start.
+	/// @param hitGroupSbtRecordCount The number of SBT records that contain hit groups.
+	/// @param width Width.
+	/// @param height Height.
+	/// @param depth Depth.
+	void traceRays(BufferPtr sbtBuffer, PtrSize sbtBufferOffset, U32 hitGroupSbtRecordCount, U32 width, U32 height,
+				   U32 depth);
+
 	/// Generate mipmaps for non-3D textures. You have to transition all the mip levels of this face and layer to
 	/// TextureUsageBit::GENERATE_MIPMAPS before calling this method.
 	/// @param texView The texture view to generate mips. It should point to a subresource that contains the whole

src/anki/gr/Common.h

@@ -134,10 +134,10 @@ public:
 	/// Max push constant size.
 	PtrSize m_pushConstantsSize = 128;
 
-	/// The size and alignment of an STB record.
-	U32 m_stbRecordSize = 0;
+	/// The size and alignment of an SBT record.
+	U32 m_sbtRecordSize = 0;
 
-	/// The size of a shader group handle that will be placed inside an STB record.
+	/// The size of a shader group handle that will be placed inside an SBT record.
 	U32 m_shaderGroupHandleSize = 0;
 
 	/// GPU vendor.

src/anki/gr/Enums.h

@@ -613,6 +613,7 @@ enum class BufferUsageBit : U64
 	TRANSFER_DESTINATION = 1ull << 26ull,
 
 	ACCELERATION_STRUCTURE_BUILD = 1ull << 27ull, ///< Will be used as a position or index buffer in a BLAS build.
+	SBT = 1ull << 28ull, ///< Will be used as SBT in a traceRays() command.
 
 	// Derived
 	ALL_UNIFORM = UNIFORM_GEOMETRY | UNIFORM_FRAGMENT | UNIFORM_COMPUTE | UNIFORM_TRACE_RAYS,
@@ -631,13 +632,13 @@ enum class BufferUsageBit : U64
 	ALL_COMPUTE = UNIFORM_COMPUTE | STORAGE_COMPUTE_READ | STORAGE_COMPUTE_WRITE | TEXTURE_COMPUTE_READ
 				  | TEXTURE_COMPUTE_WRITE | INDIRECT_COMPUTE,
 	ALL_TRACE_RAYS = UNIFORM_TRACE_RAYS | STORAGE_TRACE_RAYS_READ | STORAGE_TRACE_RAYS_WRITE | TEXTURE_TRACE_RAYS_READ
-					 | TEXTURE_TRACE_RAYS_WRITE | INDIRECT_TRACE_RAYS,
+					 | TEXTURE_TRACE_RAYS_WRITE | INDIRECT_TRACE_RAYS | SBT,
 
 	ALL_RAY_TRACING = ALL_TRACE_RAYS | ACCELERATION_STRUCTURE_BUILD,
 	ALL_READ = ALL_UNIFORM | STORAGE_GEOMETRY_READ | STORAGE_FRAGMENT_READ | STORAGE_COMPUTE_READ
 			   | STORAGE_TRACE_RAYS_READ | TEXTURE_GEOMETRY_READ | TEXTURE_FRAGMENT_READ | TEXTURE_COMPUTE_READ
 			   | TEXTURE_TRACE_RAYS_READ | INDEX | VERTEX | INDIRECT_COMPUTE | INDIRECT_DRAW | INDIRECT_TRACE_RAYS
-			   | TRANSFER_SOURCE | ACCELERATION_STRUCTURE_BUILD,
+			   | TRANSFER_SOURCE | ACCELERATION_STRUCTURE_BUILD | SBT,
 	ALL_WRITE = STORAGE_GEOMETRY_WRITE | STORAGE_FRAGMENT_WRITE | STORAGE_COMPUTE_WRITE | STORAGE_TRACE_RAYS_WRITE
 				| TEXTURE_GEOMETRY_WRITE | TEXTURE_FRAGMENT_WRITE | TEXTURE_COMPUTE_WRITE | TEXTURE_TRACE_RAYS_WRITE
 				| TRANSFER_DESTINATION,

src/anki/gr/ShaderProgram.h

@@ -151,6 +151,10 @@ class ShaderProgram : public GrObject
 public:
 	static const GrObjectType CLASS_TYPE = GrObjectType::SHADER_PROGRAM;
 
+	/// Get the shader group handles that will be used in the SBTs. The size of each handle is
+	/// GpuDeviceCapabilities::m_shaderGroupHandleSize.
+	ConstWeakArray<U8> getShaderGroupHandles() const;
+
 protected:
 	/// Construct.
 	ShaderProgram(GrManager* manager, CString name)

src/anki/gr/vulkan/AccelerationStructureImpl.cpp

@@ -190,10 +190,10 @@ void AccelerationStructureImpl::initBuildInfo()
 				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));
+				memcpy(&outInst.transform, &inInst.m_transform, sizeof(inInst.m_transform));
 				outInst.instanceCustomIndex = i & 0xFFFFFF;
 				outInst.mask = 0xFF;
-				outInst.instanceShaderBindingTableRecordOffset = inInst.m_stbRecordIndex & 0xFFFFFF;
+				outInst.instanceShaderBindingTableRecordOffset = inInst.m_sbtRecordIndex & 0xFFFFFF;
 				outInst.flags = VK_GEOMETRY_INSTANCE_TRIANGLE_FRONT_COUNTERCLOCKWISE_BIT_KHR;
 				outInst.accelerationStructureReference =
 					static_cast<const AccelerationStructureImpl&>(*inInst.m_bottomLevel).m_bottomLevelInfo.m_gpuAddress;

src/anki/gr/vulkan/CommandBuffer.cpp

@@ -280,6 +280,13 @@ void CommandBuffer::dispatchCompute(U32 groupCountX, U32 groupCountY, U32 groupC
 	self.dispatchCompute(groupCountX, groupCountY, groupCountZ);
 }
 
+void CommandBuffer::traceRays(BufferPtr sbtBuffer, PtrSize sbtBufferOffset, U32 hitGroupSbtRecordCount, U32 width,
+							  U32 height, U32 depth)
+{
+	ANKI_VK_SELF(CommandBufferImpl);
+	self.traceRaysInternal(sbtBuffer, sbtBufferOffset, hitGroupSbtRecordCount, width, height, depth);
+}
+
 void CommandBuffer::generateMipmaps2d(TextureViewPtr texView)
 {
 	ANKI_VK_SELF(CommandBufferImpl);

src/anki/gr/vulkan/CommandBufferImpl.h

@@ -306,6 +306,9 @@ public:
 
 	void dispatchCompute(U32 groupCountX, U32 groupCountY, U32 groupCountZ);
 
+	void traceRaysInternal(BufferPtr& sbtBuffer, PtrSize sbtBufferOffset, U32 hitGroupSbtRecordCount, U32 width,
+						   U32 height, U32 depth);
+
 	void resetOcclusionQuery(OcclusionQueryPtr query);
 
 	void beginOcclusionQuery(OcclusionQueryPtr query);
@@ -407,6 +410,8 @@ private:
 
 	ShaderProgramImpl* m_computeProg ANKI_DEBUG_CODE(= nullptr);
 
+	ShaderProgramImpl* m_rtProg ANKI_DEBUG_CODE(= nullptr);
+
 	VkSubpassContents m_subpassContents = VK_SUBPASS_CONTENTS_MAX_ENUM;
 
 	CommandBufferCommandType m_lastCmdType = CommandBufferCommandType::ANY_OTHER_COMMAND;

src/anki/gr/vulkan/CommandBufferImpl.inl.h

@@ -378,6 +378,83 @@ inline void CommandBufferImpl::dispatchCompute(U32 groupCountX, U32 groupCountY,
 	ANKI_CMD(vkCmdDispatch(m_handle, groupCountX, groupCountY, groupCountZ), ANY_OTHER_COMMAND);
 }
 
+inline void CommandBufferImpl::traceRaysInternal(BufferPtr& sbtBuffer, PtrSize sbtBufferOffset,
+												 U32 hitGroupSbtRecordCount, U32 width, U32 height, U32 depth)
+{
+	ANKI_ASSERT(hitGroupSbtRecordCount > 0);
+	ANKI_ASSERT(width > 0 && height > 0 && depth > 0);
+	ANKI_ASSERT(m_rtProg);
+	const ShaderProgramImpl& sprog = static_cast<const ShaderProgramImpl&>(*m_rtProg);
+	ANKI_ASSERT(sprog.getReflectionInfo().m_pushConstantsSize == m_setPushConstantsSize
+				&& "Forgot to set pushConstants");
+
+	const U32 sbtRecordCount = 1 + sprog.getMissShaderCount() + hitGroupSbtRecordCount;
+	const U32 shaderGroupBaseAlignment =
+		getGrManagerImpl().getPhysicalDeviceRayTracingProperties().shaderGroupBaseAlignment;
+	const PtrSize sbtBufferSize = sbtRecordCount * shaderGroupBaseAlignment;
+	ANKI_ASSERT(sbtBufferSize + sbtBufferOffset <= sbtBuffer->getSize());
+	ANKI_ASSERT(isAligned(shaderGroupBaseAlignment, sbtBufferOffset));
+
+	commandCommon();
+
+	// Bind descriptors
+	for(U32 i = 0; i < MAX_DESCRIPTOR_SETS; ++i)
+	{
+		if(sprog.getReflectionInfo().m_descriptorSetMask.get(i))
+		{
+			DescriptorSet dset;
+			Bool dirty;
+			Array<PtrSize, MAX_BINDINGS_PER_DESCRIPTOR_SET> dynamicOffsetsPtrSize;
+			U32 dynamicOffsetCount;
+			if(getGrManagerImpl().getDescriptorSetFactory().newDescriptorSet(
+				   m_tid, m_alloc, m_dsetState[i], dset, dirty, dynamicOffsetsPtrSize, dynamicOffsetCount))
+			{
+				ANKI_VK_LOGF("Cannot recover");
+			}
+
+			if(dirty)
+			{
+				// Vulkan should have had the dynamic offsets as VkDeviceSize and not U32. Workaround that.
+				Array<U32, MAX_BINDINGS_PER_DESCRIPTOR_SET> dynamicOffsets;
+				for(U32 i = 0; i < dynamicOffsetCount; ++i)
+				{
+					dynamicOffsets[i] = U32(dynamicOffsetsPtrSize[i]);
+				}
+
+				VkDescriptorSet dsHandle = dset.getHandle();
+
+				ANKI_CMD(vkCmdBindDescriptorSets(m_handle, VK_PIPELINE_BIND_POINT_RAY_TRACING_KHR,
+												 sprog.getPipelineLayout().getHandle(), i, 1, &dsHandle,
+												 dynamicOffsetCount, &dynamicOffsets[0]),
+						 ANY_OTHER_COMMAND);
+			}
+		}
+	}
+
+	Array<VkStridedBufferRegionKHR, 3> regions;
+
+	// Rgen
+	regions[0].buffer = static_cast<const BufferImpl&>(*sbtBuffer).getHandle();
+	regions[0].offset = sbtBufferOffset;
+	regions[0].stride = shaderGroupBaseAlignment;
+	regions[0].size = shaderGroupBaseAlignment;
+
+	// Miss
+	regions[1].buffer = regions[0].buffer;
+	regions[1].offset = regions[0].offset + regions[0].size;
+	regions[1].stride = shaderGroupBaseAlignment;
+	regions[1].size = shaderGroupBaseAlignment * sprog.getMissShaderCount();
+
+	// Hit
+	regions[2].buffer = regions[1].buffer;
+	regions[2].offset = regions[1].offset + regions[1].size;
+	regions[2].stride = shaderGroupBaseAlignment;
+	regions[2].size = shaderGroupBaseAlignment * hitGroupSbtRecordCount;
+
+	ANKI_CMD(vkCmdTraceRaysKHR(m_handle, &regions[0], &regions[1], &regions[2], nullptr, width, height, depth),
+			 ANY_OTHER_COMMAND);
+}
+
 inline void CommandBufferImpl::resetOcclusionQuery(OcclusionQueryPtr query)
 {
 	commandCommon();
@@ -750,17 +827,31 @@ inline void CommandBufferImpl::bindShaderProgram(ShaderProgramPtr& prog)
 	{
 		m_graphicsProg = &impl;
 		m_computeProg = nullptr; // Unbind the compute prog. Doesn't work like vulkan
+		m_rtProg = nullptr; // See above
 		m_state.bindShaderProgram(prog);
 	}
-	else
+	else if(!!(impl.getStages() & ShaderTypeBit::COMPUTE))
 	{
 		m_computeProg = &impl;
 		m_graphicsProg = nullptr; // See comment in the if()
+		m_rtProg = nullptr; // See above
 
 		// Bind the pipeline now
 		ANKI_CMD(vkCmdBindPipeline(m_handle, VK_PIPELINE_BIND_POINT_COMPUTE, impl.getComputePipelineHandle()),
 				 ANY_OTHER_COMMAND);
 	}
+	else
+	{
+		ANKI_ASSERT(!!(impl.getStages() & ShaderTypeBit::ALL_RAY_TRACING));
+		m_computeProg = nullptr;
+		m_graphicsProg = nullptr;
+		m_rtProg = &impl;
+
+		// Bind now
+		ANKI_CMD(
+			vkCmdBindPipeline(m_handle, VK_PIPELINE_BIND_POINT_RAY_TRACING_KHR, impl.getRayTracingPipelineHandle()),
+			ANY_OTHER_COMMAND);
+	}
 
 	for(U32 i = 0; i < MAX_DESCRIPTOR_SETS; ++i)
 	{

src/anki/gr/vulkan/GrManagerImpl.cpp

@@ -430,7 +430,7 @@ Error GrManagerImpl::initInstance(const GrManagerInitInfo& init)
 	m_capabilities.m_minorApiVersion = vulkanMinor;
 
 	m_capabilities.m_shaderGroupHandleSize = m_rtProps.shaderGroupHandleSize;
-	m_capabilities.m_stbRecordSize = m_rtProps.shaderGroupBaseAlignment;
+	m_capabilities.m_sbtRecordSize = m_rtProps.shaderGroupBaseAlignment;
 
 	return Error::NONE;
 }

src/anki/gr/vulkan/ShaderProgram.cpp

@@ -23,4 +23,9 @@ ShaderProgram* ShaderProgram::newInstance(GrManager* manager, const ShaderProgra
 	return impl;
 }
 
+ConstWeakArray<U8> ShaderProgram::getShaderGroupHandles() const
+{
+	return static_cast<const ShaderProgramImpl&>(*this).getShaderGroupHandles();
+}
+
 } // end namespace anki

src/anki/gr/vulkan/ShaderProgramImpl.cpp

@@ -25,6 +25,11 @@ ShaderProgramImpl::~ShaderProgramImpl()
 		vkDestroyPipeline(getDevice(), m_compute.m_ppline, nullptr);
 	}
 
+	if(m_rt.m_ppline)
+	{
+		vkDestroyPipeline(getDevice(), m_rt.m_ppline, nullptr);
+	}
+
 	m_shaders.destroy(getAllocator());
 	m_rt.m_allHandles.destroy(getAllocator());
 }
@@ -63,6 +68,8 @@ Error ShaderProgramImpl::init(const ShaderProgramInitInfo& inf)
 			m_shaders.emplaceBack(getAllocator(), s);
 		}
 
+		m_rt.m_missShaderCount = inf.m_rayTracingShaders.m_missShaders.getSize();
+
 		for(const RayTracingHitGroup& group : inf.m_rayTracingShaders.m_hitGroups)
 		{
 			if(group.m_anyHitShader)
@@ -315,14 +322,14 @@ Error ShaderProgramImpl::init(const ShaderProgramInitInfo& inf)
 		{
 			ANKI_TRACE_SCOPED_EVENT(VK_PIPELINE_CREATE);
 			ANKI_VK_CHECK(vkCreateRayTracingPipelinesKHR(getDevice(), getGrManagerImpl().getPipelineCache(), 1, &ci,
-														 nullptr, &m_rt.m_rtPpline));
+														 nullptr, &m_rt.m_ppline));
 		}
 
 		// Get RT handles
 		const U32 handleArraySize =
 			getGrManagerImpl().getPhysicalDeviceRayTracingProperties().shaderGroupHandleSize * groupCount;
 		m_rt.m_allHandles.create(getAllocator(), handleArraySize, 0);
-		ANKI_VK_CHECK(vkGetRayTracingShaderGroupHandlesKHR(getDevice(), m_rt.m_rtPpline, 0, groupCount, handleArraySize,
+		ANKI_VK_CHECK(vkGetRayTracingShaderGroupHandlesKHR(getDevice(), m_rt.m_ppline, 0, groupCount, handleArraySize,
 														   &m_rt.m_allHandles[0]));
 	}
 

src/anki/gr/vulkan/ShaderProgramImpl.h

@@ -82,12 +82,30 @@ public:
 		return m_compute.m_ppline;
 	}
 
+	VkPipeline getRayTracingPipelineHandle() const
+	{
+		ANKI_ASSERT(m_rt.m_ppline);
+		return m_rt.m_ppline;
+	}
+
 	ShaderTypeBit getStages() const
 	{
 		ANKI_ASSERT(!!m_stages);
 		return m_stages;
 	}
 
+	U32 getMissShaderCount() const
+	{
+		ANKI_ASSERT(m_rt.m_missShaderCount > 0);
+		return m_rt.m_missShaderCount;
+	}
+
+	ConstWeakArray<U8> getShaderGroupHandles() const
+	{
+		ANKI_ASSERT(m_rt.m_allHandles.getSize() > 0);
+		return m_rt.m_allHandles;
+	}
+
 private:
 	DynamicArray<ShaderPtr> m_shaders;
 	ShaderTypeBit m_stages = ShaderTypeBit::NONE;
@@ -114,9 +132,9 @@ private:
 	class
 	{
 	public:
-		VkPipeline m_rtPpline = VK_NULL_HANDLE;
-		BufferPtr m_stb;
+		VkPipeline m_ppline = VK_NULL_HANDLE;
 		DynamicArray<U8> m_allHandles;
+		U32 m_missShaderCount = 0;
 	} m_rt;
 };
 /// @}