[vulkan] Improve Vulkan RHI impl with lower overhead internal impleme…

…ntations (taichi-dev#6912)

Also added partial RHI support for Vulkan Dynamic Rendering, but there
are some interactions and bugs to figure out between this and ImGUI, so
the device creator will not actively enable this feature for now.

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

taichi/rhi/device.h

@@ -17,6 +17,7 @@ enum class RhiResult {
   error = -1,
   invalid_usage = -2,
   not_supported = -3,
+  out_of_memory = -4,
 };
 
 constexpr size_t kBufferSizeEntireSize = size_t(-1);
@@ -54,7 +55,7 @@ struct LLVMRuntime;
 
 // TODO: Figure out how to support images. Temporary solutions is to have all
 // opque types such as images work as an allocation
-using DeviceAllocationId = uint32_t;
+using DeviceAllocationId = uint64_t;
 
 struct TI_DLL_EXPORT DeviceAllocation {
   Device *device{nullptr};

taichi/rhi/impl_support.h

@@ -2,6 +2,9 @@
 
 #include "taichi/rhi/device.h"
 #include <assert.h>
+#include <forward_list>
+#include <unordered_set>
+#include <mutex>
 
 namespace taichi::lang {
 
@@ -89,5 +92,61 @@ struct BidirMap {
   }
 };
 
+// A synchronized list of objects that is pointer stable & reuse objects
+template <class T>
+class SyncedPtrStableObjectList {
+  using storage_block = std::array<uint8_t, sizeof(T)>;
+
+ public:
+  template <typename... Params>
+  T &acquire(Params &&...args) {
+    std::lock_guard<std::mutex> _(lock_);
+
+    void *storage = nullptr;
+    if (free_nodes_.empty()) {
+      storage = objects_.emplace_front().data();
+    } else {
+      storage = free_nodes_.back();
+      free_nodes_.pop_back();
+    }
+    return *new (storage) T(std::forward<Params>(args)...);
+  }
+
+  void release(T *ptr) {
+    std::lock_guard<std::mutex> _(lock_);
+
+    ptr->~T();
+    free_nodes_.push_back(ptr);
+  }
+
+  void clear() {
+    std::lock_guard<std::mutex> _(lock_);
+
+    // Transfer to quick look-up
+    std::unordered_set<void *> free_nodes_set(free_nodes_.begin(),
+                                              free_nodes_.end());
+    free_nodes_.clear();
+    // Destroy live objects
+    for (auto &storage : objects_) {
+      T *obj = reinterpret_cast<T *>(storage.data());
+      // Call destructor if object is not in the free list (thus live)
+      if (free_nodes_set.find(obj) == free_nodes_set.end()) {
+        obj->~T();
+      }
+    }
+    // Clear the storage
+    objects_.clear();
+  }
+
+  ~SyncedPtrStableObjectList() {
+    clear();
+  }
+
+ private:
+  std::mutex lock_;
+  std::forward_list<storage_block> objects_;
+  std::vector<void *> free_nodes_;
+};
+
 }  // namespace rhi_impl
 }  // namespace taichi::lang

taichi/rhi/opengl/opengl_device.cpp

@@ -562,7 +562,8 @@ DeviceAllocation GLDevice::allocate_memory(const AllocParams &params) {
 }
 
 void GLDevice::dealloc_memory(DeviceAllocation handle) {
-  glDeleteBuffers(1, &handle.alloc_id);
+  GLuint buffer = GLuint(handle.alloc_id);
+  glDeleteBuffers(1, &buffer);
   check_opengl_error("glDeleteBuffers");
 }
 
@@ -714,7 +715,8 @@ DeviceAllocation GLDevice::create_image(const ImageParams &params) {
 }
 
 void GLDevice::destroy_image(DeviceAllocation handle) {
-  glDeleteTextures(1, &handle.alloc_id);
+  GLuint texture = GLuint(handle.alloc_id);
+  glDeleteTextures(1, &texture);
   check_opengl_error("glDeleteTextures");
   image_to_dims_.erase(handle.alloc_id);
   image_to_int_format_.erase(handle.alloc_id);

taichi/rhi/vulkan/vulkan_api.cpp

@@ -369,6 +369,38 @@ IVkPipeline create_graphics_pipeline(VkDevice device,
   return obj;
 }
 
+IVkPipeline create_graphics_pipeline_dynamic(
+    VkDevice device,
+    VkGraphicsPipelineCreateInfo *create_info,
+    VkPipelineRenderingCreateInfoKHR *rendering_info,
+    IVkPipelineLayout layout,
+    IVkPipelineCache cache,
+    IVkPipeline base_pipeline) {
+  IVkPipeline obj = std::make_shared<DeviceObjVkPipeline>();
+  obj->device = device;
+  obj->ref_layout = layout;
+  obj->ref_cache = cache;
+  obj->ref_renderpass = nullptr;
+
+  create_info->pNext = rendering_info;
+  create_info->layout = layout->layout;
+
+  if (base_pipeline) {
+    create_info->basePipelineHandle = base_pipeline->pipeline;
+    create_info->basePipelineIndex = -1;
+  } else {
+    create_info->basePipelineHandle = VK_NULL_HANDLE;
+    create_info->basePipelineIndex = 0;
+  }
+
+  VkResult res =
+      vkCreateGraphicsPipelines(device, cache ? cache->cache : VK_NULL_HANDLE,
+                                1, create_info, nullptr, &obj->pipeline);
+  BAIL_ON_VK_BAD_RESULT_NO_RETURN(res, "failed to create graphics pipeline");
+
+  return obj;
+}
+
 IVkPipeline create_raytracing_pipeline(
     VkDevice device,
     VkRayTracingPipelineCreateInfoKHR *create_info,

taichi/rhi/vulkan/vulkan_api.h

@@ -167,6 +167,13 @@ IVkPipeline create_graphics_pipeline(VkDevice device,
                                      IVkPipelineLayout layout,
                                      IVkPipelineCache cache = nullptr,
                                      IVkPipeline base_pipeline = nullptr);
+IVkPipeline create_graphics_pipeline_dynamic(
+    VkDevice device,
+    VkGraphicsPipelineCreateInfo *create_info,
+    VkPipelineRenderingCreateInfoKHR *rendering_info,
+    IVkPipelineLayout layout,
+    IVkPipelineCache cache = nullptr,
+    IVkPipeline base_pipeline = nullptr);
 IVkPipeline create_raytracing_pipeline(
     VkDevice device,
     VkRayTracingPipelineCreateInfoKHR *create_info,