Fix implementation problems in HLSL and write initial transferProperties function

deprilula28 · deprilula28 · commit 4be1a3ce9fe6 · 2024-01-22T14:43:38.000-03:00
diff --git a/include/nbl/builtin/hlsl/property_pool/copy.comp.hlsl b/include/nbl/builtin/hlsl/property_pool/copy.comp.hlsl
@@ -14,16 +14,22 @@ namespace property_pools
 template<bool Fill, bool SrcIndexIota, bool DstIndexIota, uint64_t SrcIndexSizeLog2, uint64_t DstIndexSizeLog2>
 struct TransferLoop
 {
-    void iteration(uint propertyId, uint64_t propertySize, uint64_t srcAddr, uint64_t dstAddr, uint invocationIndex)
+    void iteration(uint propertyId, TransferRequest transferRequest, uint invocationIndex)
     {
-        const uint64_t srcOffset = uint64_t(invocationIndex) * (uint64_t(1) << SrcIndexSizeLog2) * propertySize;
-        const uint64_t dstOffset = uint64_t(invocationIndex) * (uint64_t(1) << DstIndexSizeLog2) * propertySize;
+        const uint64_t srcIndexSize = uint64_t(1) << SrcIndexSizeLog2;
+        const uint64_t dstIndexSize = uint64_t(1) << DstIndexSizeLog2;
+
+        const uint64_t srcOffset = uint64_t(invocationIndex) * srcIndexSize * transferRequest.propertySize;
+        const uint64_t dstOffset = uint64_t(invocationIndex) * dstIndexSize * transferRequest.propertySize;
         
-        const uint64_t srcIndexAddress = Fill ? srcAddr + srcOffset : srcAddr;
-        const uint64_t dstIndexAddress = Fill ? dstAddr + dstOffset : dstAddr;
+        const uint64_t srcIndexAddress = Fill ? transferRequest.srcIndexAddr + srcOffset : transferRequest.srcIndexAddr;
+        const uint64_t dstIndexAddress = Fill ? transferRequest.dstIndexAddr + dstOffset : transferRequest.dstIndexAddr;
+
+        const uint64_t srcAddressBufferOffset = SrcIndexIota ? srcIndexAddress : vk::RawBufferLoad<uint32_t>(srcIndexAddress);
+        const uint64_t dstAddressBufferOffset = DstIndexIota ? dstIndexAddress : vk::RawBufferLoad<uint32_t>(dstIndexAddress);
 
-        const uint64_t srcAddressMapped = SrcIndexIota ? srcIndexAddress : vk::RawBufferLoad<uint64_t>(srcIndexAddress); 
-        const uint64_t dstAddressMapped = DstIndexIota ? dstIndexAddress : vk::RawBufferLoad<uint64_t>(dstIndexAddress); 
+        const uint64_t srcAddressMapped = transferRequest.srcAddr + srcAddressBufferOffset * srcIndexSize; 
+        const uint64_t dstAddressMapped = transferRequest.dstAddr + dstAddressBufferOffset * dstIndexSize; 
 
         if (SrcIndexSizeLog2 == 0) {} // we can't write individual bytes
         else if (SrcIndexSizeLog2 == 1) vk::RawBufferStore<uint16_t>(dstAddressMapped, vk::RawBufferLoad<uint16_t>(srcAddressMapped));
@@ -35,17 +41,21 @@ struct TransferLoop
     {
         uint64_t elementCount = uint64_t(transferRequest.elementCount32)
             | uint64_t(transferRequest.elementCountExtra) << 32;
-        uint lastInvocation = min(elementCount, globals.endOffset);
-        for (uint invocationIndex = globals.beginOffset + baseInvocationIndex; invocationIndex < lastInvocation; invocationIndex += dispatchSize)
+        uint64_t lastInvocation = min(elementCount, globals.endOffset);
+        for (uint64_t invocationIndex = globals.beginOffset + baseInvocationIndex; invocationIndex < lastInvocation; invocationIndex += dispatchSize)
         {
-            iteration(propertyId, transferRequest.propertySize, transferRequest.srcAddr, transferRequest.dstAddr, invocationIndex);
+            iteration(propertyId, transferRequest, invocationIndex);
         }
     }
 };
 
 // For creating permutations of the functions based on parameters that are constant over the transfer request
 // These branches should all be scalar, and because of how templates are compiled statically, the loops shouldn't have any
 // branching within them
+// 
+// Permutations:
+// 2 (fill or not) * 2 (src index iota or not) * 2 (dst index iota or not) * 4 (src index size) * 4 (dst index size)
+// Total amount of permutations: 128
 
 template<bool Fill, bool SrcIndexIota, bool DstIndexIota, uint64_t SrcIndexSizeLog2>
 struct TransferLoopPermutationSrcIndexSizeLog
@@ -76,7 +86,7 @@ struct TransferLoopPermutationSrcIota
 {
     void copyLoop(uint baseInvocationIndex, uint propertyId, TransferRequest transferRequest, uint dispatchSize)
     {
-        bool dstIota = transferRequest.dstAddr == 0;
+        bool dstIota = transferRequest.dstIndexAddr == 0;
         if (dstIota) { TransferLoopPermutationDstIota<Fill, SrcIndexIota, true> loop; loop.copyLoop(baseInvocationIndex, propertyId, transferRequest, dispatchSize); }
         else { TransferLoopPermutationDstIota<Fill, SrcIndexIota, false> loop; loop.copyLoop(baseInvocationIndex, propertyId, transferRequest, dispatchSize); }
     }
@@ -87,7 +97,7 @@ struct TransferLoopPermutationFill
 {
     void copyLoop(uint baseInvocationIndex, uint propertyId, TransferRequest transferRequest, uint dispatchSize)
     {
-        bool srcIota = transferRequest.srcAddr == 0;
+        bool srcIota = transferRequest.srcIndexAddr == 0;
         if (srcIota) { TransferLoopPermutationSrcIota<Fill, true> loop; loop.copyLoop(baseInvocationIndex, propertyId, transferRequest, dispatchSize); }
         else { TransferLoopPermutationSrcIota<Fill, false> loop; loop.copyLoop(baseInvocationIndex, propertyId, transferRequest, dispatchSize); }
     }
diff --git a/include/nbl/builtin/hlsl/property_pool/transfer.hlsl b/include/nbl/builtin/hlsl/property_pool/transfer.hlsl
@@ -1,6 +1,8 @@
 #ifndef _NBL_BUILTIN_HLSL_GLSL_PROPERTY_POOLS_TRANSFER_
 #define _NBL_BUILTIN_HLSL_GLSL_PROPERTY_POOLS_TRANSFER_
 
+#include "nbl/builtin/hlsl/cpp_compat.hlsl"
+
 namespace nbl
 {
 namespace hlsl
@@ -40,6 +42,9 @@ struct GlobalPushContants
     // BDA address (GPU pointer) into the transfer commands buffer
     uint64_t transferCommandsAddress;
 };
+
+NBL_CONSTEXPR uint32_t MaxPropertiesPerDispatch = 128;
+
 }
 }
 }
diff --git a/include/nbl/video/utilities/CPropertyPoolHandler.h b/include/nbl/video/utilities/CPropertyPoolHandler.h
@@ -74,12 +74,12 @@ class NBL_API2 CPropertyPoolHandler final : public core::IReferenceCounted, publ
 			asset::SBufferRange<IGPUBuffer> memblock = {};
 			E_FLAG flags = EF_NONE;
 			uint16_t elementSize = 0u;
-			uint32_t elementCount = 0u;
+			uint64_t elementCount = 0u;
 			// the source or destination buffer depending on the transfer type
 			asset::SBufferBinding<video::IGPUBuffer> buffer = {};
 			// can be invalid, if invalid, treated like an implicit {0,1,2,3,...} iota view
-			uint32_t srcAddressesOffset = IPropertyPool::invalid;
-			uint32_t dstAddressesOffset = IPropertyPool::invalid;
+			uint64_t srcAddressesOffset = IPropertyPool::invalid;
+			uint64_t dstAddressesOffset = IPropertyPool::invalid;
 		};
 		// Fence must be not pending yet, `cmdbuf` must be already in recording state.
 		[[nodiscard]] bool transferProperties(
@@ -206,7 +206,7 @@ class NBL_API2 CPropertyPoolHandler final : public core::IReferenceCounted, publ
     protected:
 		~CPropertyPoolHandler() {}
 
-		static inline constexpr auto MaxPropertiesPerDispatch = 0u; // TODO
+		static inline constexpr auto MaxPropertiesPerDispatch = nbl::hlsl::property_pools::MaxPropertiesPerDispatch;
 		static inline constexpr auto DescriptorCacheSize = 128u;
 
 
diff --git a/src/nbl/video/utilities/CPropertyPoolHandler.cpp b/src/nbl/video/utilities/CPropertyPoolHandler.cpp
@@ -31,6 +31,9 @@ CPropertyPoolHandler::CPropertyPoolHandler(core::smart_refctd_ptr<ILogicalDevice
 		return gpuShader;
 	};
 	auto shader = loadShader("../../../include/nbl/builtin/hlsl/property_pool/copy.comp.hlsl");
+	const asset::SPushConstantRange baseDWORD = {asset::IShader::ESS_COMPUTE,0u,sizeof(nbl::hlsl::property_pools::GlobalPushContants)};
+	auto layout = m_device->createPipelineLayout(&baseDWORD,&baseDWORD+1u);
+	m_pipeline = m_device->createComputePipeline(nullptr,std::move(layout),std::move(shader));
 
 #if 0
 	const auto& deviceLimits = m_device->getPhysicalDevice()->getLimits();
@@ -89,6 +92,73 @@ bool CPropertyPoolHandler::transferProperties(
 	system::logger_opt_ptr logger, const uint32_t baseDWORD, const uint32_t endDWORD
 )
 {
+	if (requestsBegin==requestsEnd)
+		return true;
+	if (!scratch.buffer || !scratch.buffer->getCreationParams().usage.hasFlags(IGPUBuffer::EUF_INLINE_UPDATE_VIA_CMDBUF))
+	{
+		logger.log("CPropertyPoolHandler: Need a valid scratch buffer which can have updates staged from the commandbuffer!",system::ILogger::ELL_ERROR);
+		return false;
+	}
+	// TODO: validate usage flags
+	uint32_t maxScratchSize = MaxPropertiesPerDispatch * sizeof(nbl::hlsl::property_pools::TransferRequest);
+	if (scratch.offset + maxScratchSize > scratch.buffer->getSize())
+	    logger.log("CPropertyPoolHandler: The scratch buffer binding provided might not be big enough in the worst case! (Scratch buffer size: %i Max scratch size: %i)",
+			system::ILogger::ELL_WARNING,
+			scratch.buffer->getSize() - scratch.offset,
+			maxScratchSize);
+
+	const auto totalProps = std::distance(requestsBegin,requestsEnd);
+	bool success = true;
+	
+	uint32_t numberOfPasses = totalProps / MaxPropertiesPerDispatch;
+	nbl::hlsl::property_pools::TransferRequest transferRequestsData[MaxPropertiesPerDispatch];
+	uint64_t scratchBufferDeviceAddr = m_device->getBufferDeviceAddress(scratch.buffer.get()) + scratch.offset;
+	uint64_t addressBufferDeviceAddr = m_device->getBufferDeviceAddress(addresses.buffer.get()) + addresses.offset;
+
+	for (uint32_t transferPassRequestsIndex = 0; transferPassRequestsIndex < totalProps; transferPassRequestsIndex += MaxPropertiesPerDispatch)
+	{
+		const TransferRequest* transferPassRequests = requestsBegin + transferPassRequestsIndex;
+		uint32_t requestsThisPass = core::min<uint32_t>(std::distance(transferPassRequests, requestsEnd), MaxPropertiesPerDispatch);
+		uint64_t maxElements = 0;
+		for (uint32_t i = 0; i < requestsThisPass; i ++)
+		{
+			auto& transferRequest = transferRequestsData[i];
+			auto srcRequest = transferPassRequests + i;
+			transferRequest.srcAddr = m_device->getBufferDeviceAddress(srcRequest->memblock.buffer.get()) + srcRequest->memblock.offset;
+			transferRequest.dstAddr = m_device->getBufferDeviceAddress(srcRequest->buffer.buffer.get()) + srcRequest->buffer.offset;
+			transferRequest.srcIndexAddr = srcRequest->srcAddressesOffset ? addressBufferDeviceAddr + srcRequest->srcAddressesOffset : 0;
+			transferRequest.dstIndexAddr = srcRequest->dstAddressesOffset ? addressBufferDeviceAddr + srcRequest->dstAddressesOffset : 0;
+			transferRequest.elementCount32 = uint32_t(srcRequest->elementCount & (uint64_t(1) << 32) - 1);
+			transferRequest.elementCountExtra = uint32_t(srcRequest->elementCount >> 32);
+			transferRequest.propertySize = srcRequest->elementSize;
+			transferRequest.fill = 0; // TODO
+			transferRequest.srcIndexSizeLog2 = 1u; // TODO
+			transferRequest.dstIndexSizeLog2 = 1u; // TODO
+
+			maxElements = core::max<uint64_t>(maxElements, srcRequest->elementCount);
+		}
+		cmdbuf->updateBuffer(scratch.buffer.get(),scratch.offset,sizeof(TransferRequest)*requestsThisPass,transferRequestsData);
+		// TODO: pipeline barrier
+		cmdbuf->bindComputePipeline(m_pipeline.get());
+		
+		nbl::hlsl::property_pools::GlobalPushContants pushConstants;
+		{
+			pushConstants.beginOffset = baseDWORD;
+			pushConstants.endOffset = endDWORD;
+			pushConstants.transferCommandsAddress = scratchBufferDeviceAddr;
+		}
+		cmdbuf->pushConstants(m_pipeline->getLayout(), asset::IShader::ESS_COMPUTE, 0u, sizeof(nbl::hlsl::property_pools::GlobalPushContants), &pushConstants);
+
+		// dispatch
+		{
+			const auto& limits = m_device->getPhysicalDevice()->getLimits();
+			const auto invocationCoarseness = limits.maxOptimallyResidentWorkgroupInvocations * requestsThisPass;
+			cmdbuf->dispatch(limits.computeOptimalPersistentWorkgroupDispatchSize(maxElements,invocationCoarseness), requestsThisPass, 1u);
+		}
+		// TODO: pipeline barrier
+	}
+
+	return success;
 #if 0
 	if (requestsBegin==requestsEnd)
 		return true;

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,8 @@`
`1`	`1`	`#ifndef _NBL_BUILTIN_HLSL_GLSL_PROPERTY_POOLS_TRANSFER_`
`2`	`2`	`#define _NBL_BUILTIN_HLSL_GLSL_PROPERTY_POOLS_TRANSFER_`
`3`	`3`
	`4`	`+#include "nbl/builtin/hlsl/cpp_compat.hlsl"`
	`5`	`+`
`4`	`6`	`namespace nbl`
`5`	`7`	`{`
`6`	`8`	`namespace hlsl`
`@@ -40,6 +42,9 @@ struct GlobalPushContants`
`40`	`42`	`// BDA address (GPU pointer) into the transfer commands buffer`
`41`	`43`	`uint64_t transferCommandsAddress;`
`42`	`44`	`};`
	`45`	`+`
	`46`	`+NBL_CONSTEXPR uint32_t MaxPropertiesPerDispatch = 128;`
	`47`	`+`
`43`	`48`	`}`
`44`	`49`	`}`
`45`	`50`	`}`