[AOT] Support importing external Vulkan buffers (#5020)

taichi-dev · May 21, 2022 · fba92cf · fba92cf
1 parent 25644e3
commit fba92cf
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/taichi/backends/vulkan/vulkan_device.cpp b/taichi/backends/vulkan/vulkan_device.cpp
@@ -1682,6 +1682,22 @@ vkapi::IVkFramebuffer VulkanDevice::get_framebuffer(
   return framebuffer;
 }
 
+DeviceAllocation VulkanDevice::import_vkbuffer(vkapi::IVkBuffer buffer) {
+  AllocationInternal alloc_int{};
+  alloc_int.external = true;
+  alloc_int.buffer = buffer;
+  alloc_int.mapped = nullptr;
+  alloc_int.addr = 0;
+
+  DeviceAllocation alloc;
+  alloc.device = this;
+  alloc.alloc_id = alloc_cnt_++;
+
+  allocations_[alloc.alloc_id] = alloc_int;
+
+  return alloc;
+}
+
 DeviceAllocation VulkanDevice::import_vk_image(vkapi::IVkImage image,
                                                vkapi::IVkImageView view,
                                                VkFormat format) {

diff --git a/taichi/backends/vulkan/vulkan_device.h b/taichi/backends/vulkan/vulkan_device.h
@@ -604,6 +604,9 @@ class TI_DLL_EXPORT VulkanDevice : public GraphicsDevice {
 
   std::tuple<vkapi::IVkImage, vkapi::IVkImageView, VkFormat> get_vk_image(
       const DeviceAllocation &alloc) const;
+
+  DeviceAllocation import_vkbuffer(vkapi::IVkBuffer buffer);
+
   DeviceAllocation import_vk_image(vkapi::IVkImage image,
                                    vkapi::IVkImageView view,
                                    VkFormat format);
@@ -642,6 +645,7 @@ class TI_DLL_EXPORT VulkanDevice : public GraphicsDevice {
 
   // Memory allocation
   struct AllocationInternal {
+    bool external{false};
     VmaAllocationInfo alloc_info;
     vkapi::IVkBuffer buffer;
     void *mapped{nullptr};