[aot] Fixed buffer device address import (#6326)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
taichi-dev · Oct 14, 2022 · 8936fa6 · 8936fa6
1 parent 306dd88
commit 8936fa6
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/taichi/rhi/vulkan/vulkan_device.cpp b/taichi/rhi/vulkan/vulkan_device.cpp
@@ -1838,7 +1838,13 @@ DeviceAllocation VulkanDevice::import_vkbuffer(vkapi::IVkBuffer buffer) {
   alloc_int.external = true;
   alloc_int.buffer = buffer;
   alloc_int.mapped = nullptr;
-  alloc_int.addr = 0;
+  if (get_cap(DeviceCapability::spirv_has_physical_storage_buffer)) {
+    VkBufferDeviceAddressInfoKHR info{};
+    info.sType = VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO;
+    info.buffer = buffer->buffer;
+    info.pNext = nullptr;
+    alloc_int.addr = vkGetBufferDeviceAddress(device_, &info);
+  }
 
   DeviceAllocation alloc;
   alloc.device = this;