Using std::unordered_map to store the Buffer Device Addresses

This commit finished off the build acceleration structure command. This is because in MVKDevice, we are now using a std::unordered_map instead of a custom map implementation.
KhronosGroup · Jul 10, 2023 · 43a987c · 43a987c
1 parent eab6f6f
commit 43a987c
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 12 deletions.
diff --git a/MoltenVK/MoltenVK/Commands/MVKCmdAccelerationStructure.mm b/MoltenVK/MoltenVK/Commands/MVKCmdAccelerationStructure.mm
@@ -37,21 +37,21 @@
 void MVKCmdBuildAccelerationStructure::encode(MVKCommandEncoder* cmdEncoder) {
     id<MTLAccelerationStructureCommandEncoder> accStructEncoder = cmdEncoder->getMTLAccelerationStructureEncoder(kMVKCommandUseNone);
     id<MTLAccelerationStructure> srcAccelerationStructure = (id<MTLAccelerationStructure>)_geometryInfos.srcAccelerationStructure;
-    id<MTLAccelerationStructure> dstAccelerationStructure = (id<MTLAccelerationStructure>)_geometryInfos.dstAccelerationStructure; // Target acceleration Structure
+    id<MTLAccelerationStructure> dstAccelerationStructure = (id<MTLAccelerationStructure>)_geometryInfos.dstAccelerationStructure;
 
     MTLAccelerationStructureDescriptor* accStructDescriptor = [MTLAccelerationStructureDescriptor new];
     accStructDescriptor.usage = MTLAccelerationStructureUsageNone;
-    /*
-      * The NVIDIA extension seemed to use to provide the scratch buffer offset, but not the KHR version
-      * However the KHR extension does not seem to have anything similar, for now I'll leave it 0, but
-      * it should be changed.
-    */
+
+    MVKDevice* mvkDvc = cmdEncoder->getDevice();
+    MVKBuffer* mvkBuffer = mvkDvc->getBufferAtAddress(_geometryInfos.scratchData.deviceAddress);
+
+    id<MTLBuffer> scratchBuffer = mvkBuffer->getMTLBuffer();
     int scratchBufferOffset = 0;
 
-//    [accStructEncoder buildAccelerationStructure:dstAccelerationStructure
-//                                    descriptor:accStructDescriptor
-//                                    scratchBuffer:nil
-//                                    scratchBufferOffset:scratchBufferOffset];
+    [accStructEncoder buildAccelerationStructure:dstAccelerationStructure
+                                    descriptor:accStructDescriptor
+                                    scratchBuffer:scratchBuffer
+                                    scratchBufferOffset:scratchBufferOffset];
 }
 
 VkResult MVKCmdCopyAccelerationStructure::setContent(MVKCommandBuffer* cmdBuff,
@@ -62,7 +62,7 @@
     MVKAccelerationStructure* mvkDstAccStruct = (MVKAccelerationStructure*)dstAccelerationStructure;
 
     _srcAccelerationStructure = mvkSrcAccStruct->getMTLAccelerationStructure();
-    _dstAccelerationStructure = mvkSrcAccStruct->getMTLAccelerationStructure();
+    _dstAccelerationStructure = mvkDstAccStruct->getMTLAccelerationStructure();
     return VK_SUCCESS;
 }
 

diff --git a/MoltenVK/MoltenVK/GPUObjects/MVKDevice.h b/MoltenVK/MoltenVK/GPUObjects/MVKDevice.h
@@ -524,6 +524,9 @@ class MVKDevice : public MVKDispatchableVulkanAPIObject {
 								 const VkCalibratedTimestampInfoEXT* pTimestampInfos,
 								 uint64_t* pTimestamps,
 								 uint64_t* pMaxDeviation);
+
+    /** Returns a pointer to the buffer at the provided address*/
+    MVKBuffer* getBufferAtAddress(uint64_t address);
 
 #pragma mark Object lifecycle
 
@@ -913,6 +916,7 @@ class MVKDevice : public MVKDispatchableVulkanAPIObject {
 	MVKSmallVector<MVKSmallVector<MVKQueue*, kMVKQueueCountPerQueueFamily>, kMVKQueueFamilyCount> _queuesByQueueFamilyIndex;
 	MVKSmallVector<MVKResource*, 256> _resources;
 	MVKSmallVector<MVKBuffer*, 8> _gpuAddressableBuffers;
+    std::unordered_map<uint64_t, MVKBuffer*> _gpuBufferAddressMap;
 	MVKSmallVector<MVKPrivateDataSlot*> _privateDataSlots;
 	MVKSmallVector<bool> _privateDataSlotsAvailability;
 	MVKSmallVector<MVKSemaphoreImpl*> _awaitingSemaphores;

diff --git a/MoltenVK/MoltenVK/GPUObjects/MVKDevice.mm b/MoltenVK/MoltenVK/GPUObjects/MVKDevice.mm
@@ -3640,6 +3640,20 @@ static uint32_t mvkGetEntryProperty(io_registry_entry_t entry, CFStringRef prope
 	*pMaxDeviation = cpuEnd - cpuStart;
 }
 
+MVKBuffer* MVKDevice::getBufferAtAddress(uint64_t address)
+{
+    auto mvkBufIt = _gpuBufferAddressMap.find(address);
+
+    if(mvkBufIt != _gpuBufferAddressMap.end())
+    {
+        return mvkBufIt->second;
+    }
+    else
+    {
+        return nullptr;
+    }
+}
+
 #pragma mark Object lifecycle
 
 uint32_t MVKDevice::getVulkanMemoryTypeIndex(MTLStorageMode mtlStorageMode) {
@@ -4132,7 +4146,9 @@ static uint32_t mvkGetEntryProperty(io_registry_entry_t entry, CFStringRef prope
 	_resources.push_back(mvkBuff);
 	if (mvkIsAnyFlagEnabled(mvkBuff->getUsage(), VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT)) {
 		_gpuAddressableBuffers.push_back(mvkBuff);
-	}
+        std::pair<uint64_t, MVKBuffer*> _bufferAddressPair = std::make_pair(mvkBuff->getMTLBufferGPUAddress(), mvkBuff);
+        _gpuBufferAddressMap.insert(_bufferAddressPair);
+    }
 	return mvkBuff;
 }
 
@@ -4143,6 +4159,8 @@ static uint32_t mvkGetEntryProperty(io_registry_entry_t entry, CFStringRef prope
 	mvkRemoveFirstOccurance(_resources, mvkBuff);
 	if (mvkIsAnyFlagEnabled(mvkBuff->getUsage(), VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT)) {
 		mvkRemoveFirstOccurance(_gpuAddressableBuffers, mvkBuff);
+        auto bufferAddressIt = _gpuBufferAddressMap.find(mvkBuff->getMTLBufferGPUAddress());
+        _gpuBufferAddressMap.erase(bufferAddressIt);
 	}
 	return mvkBuff;
 }