From 1e324c6834c2020013c3285f0f773d2b5039239b Mon Sep 17 00:00:00 2001
From: Bob Cao <bobcaocheng@gmail.com>
Date: Wed, 25 Jan 2023 18:32:38 -0800
Subject: [PATCH] [rhi] Add create_image_unique stub & misc RHI bug fixes
 (#7232)

Part of #7218
---
 taichi/rhi/device.cpp               |  4 ++++
 taichi/rhi/device.h                 | 14 ++++++++++++++
 taichi/rhi/opengl/opengl_device.cpp |  3 ++-
 taichi/rhi/vulkan/vulkan_device.cpp | 24 +++++++++++++++++++-----
 4 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/taichi/rhi/device.cpp b/taichi/rhi/device.cpp
index 81b97e7feafbc..2dfc3881837f7 100644
--- a/taichi/rhi/device.cpp
+++ b/taichi/rhi/device.cpp
@@ -18,6 +18,10 @@ DeviceAllocationGuard::~DeviceAllocationGuard() {
   device->dealloc_memory(*this);
 }
 
+DeviceImageGuard::~DeviceImageGuard() {
+  dynamic_cast<GraphicsDevice *>(device)->destroy_image(*this);
+}
+
 DevicePtr DeviceAllocation::get_ptr(uint64_t offset) const {
   return DevicePtr{{device, alloc_id}, offset};
 }
diff --git a/taichi/rhi/device.h b/taichi/rhi/device.h
index 821bdb5b8931e..47fb2331cf07b 100644
--- a/taichi/rhi/device.h
+++ b/taichi/rhi/device.h
@@ -78,6 +78,17 @@ struct TI_DLL_EXPORT DeviceAllocationGuard : public DeviceAllocation {
   ~DeviceAllocationGuard();
 };
 
+using DeviceAllocationUnique = std::unique_ptr<DeviceAllocationGuard>;
+
+struct TI_DLL_EXPORT DeviceImageGuard : public DeviceAllocation {
+  explicit DeviceImageGuard(DeviceAllocation alloc) : DeviceAllocation(alloc) {
+  }
+  DeviceImageGuard(const DeviceAllocationGuard &) = delete;
+  ~DeviceImageGuard();
+};
+
+using DeviceImageUnique = std::unique_ptr<DeviceImageGuard>;
+
 struct TI_DLL_EXPORT DevicePtr : public DeviceAllocation {
   uint64_t offset{0};
 
@@ -894,6 +905,9 @@ class TI_DLL_EXPORT GraphicsDevice : public Device {
   // `GfxRuntime::create_image`. `GfxRuntime` is available in `ProgramImpl`
   // of GPU backends.
   virtual DeviceAllocation create_image(const ImageParams &params) = 0;
+  inline DeviceImageUnique create_image_unique(const ImageParams &params) {
+    return std::make_unique<DeviceImageGuard>(this->create_image(params));
+  }
   virtual void destroy_image(DeviceAllocation handle) = 0;
 
   virtual void image_transition(DeviceAllocation img,
diff --git a/taichi/rhi/opengl/opengl_device.cpp b/taichi/rhi/opengl/opengl_device.cpp
index 2f28d0e7043cb..7485fef6ff1f9 100644
--- a/taichi/rhi/opengl/opengl_device.cpp
+++ b/taichi/rhi/opengl/opengl_device.cpp
@@ -847,7 +847,8 @@ void GLCommandList::CmdBindResources::execute() {
 }
 
 void GLCommandList::CmdBufferBarrier::execute() {
-  glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);
+  glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT | GL_BUFFER_UPDATE_BARRIER_BIT |
+                  GL_UNIFORM_BARRIER_BIT);
   check_opengl_error("glMemoryBarrier");
 }
 
diff --git a/taichi/rhi/vulkan/vulkan_device.cpp b/taichi/rhi/vulkan/vulkan_device.cpp
index c43c34b85a374..83fc6f230a55b 100644
--- a/taichi/rhi/vulkan/vulkan_device.cpp
+++ b/taichi/rhi/vulkan/vulkan_device.cpp
@@ -552,10 +552,16 @@ void VulkanPipeline::create_graphics_pipeline(
   color_blending.blendConstants[3] = 0.0f;
 
   if (raster_params.blending.size()) {
-    RHI_ASSERT(raster_params.blending.size() ==
-                   graphics_pipeline_template_->blend_attachments.size() &&
-               "RasterParams::blending (size={}) must either be zero sized "
-               "or match the number of fragment shader outputs (size={}).");
+    if (raster_params.blending.size() != color_blending.attachmentCount) {
+      std::array<char, 256> buf;
+      snprintf(buf.data(), buf.size(),
+               "RasterParams::blending (size=%u) must either be zero sized "
+               "or match the number of fragment shader outputs (size=%u).",
+               uint32_t(raster_params.blending.size()),
+               uint32_t(color_blending.attachmentCount));
+      RHI_LOG_ERROR(buf.data());
+      RHI_ASSERT(false);
+    }
 
     for (int i = 0; i < raster_params.blending.size(); i++) {
       auto &state = graphics_pipeline_template_->blend_attachments[i];
@@ -1694,6 +1700,8 @@ RhiResult VulkanDevice::create_pipeline(Pipeline **out_pipeline,
 DeviceAllocation VulkanDevice::allocate_memory(const AllocParams &params) {
   AllocationInternal &alloc = allocations_.acquire();
 
+  RHI_ASSERT(params.size > 0);
+
   VkBufferCreateInfo buffer_info{};
   buffer_info.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
   buffer_info.pNext = nullptr;
@@ -1770,7 +1778,13 @@ DeviceAllocation VulkanDevice::allocate_memory(const AllocParams &params) {
     alloc_info.requiredFlags = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT;
   }
 
-  if (get_caps().get(DeviceCapability::spirv_has_physical_storage_buffer)) {
+  if (get_caps().get(DeviceCapability::spirv_has_physical_storage_buffer) &&
+      ((alloc_info.usage & VK_BUFFER_USAGE_STORAGE_BUFFER_BIT) ||
+       (alloc_info.usage &
+        VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR) ||
+       (alloc_info.usage &
+        VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_STORAGE_BIT_KHR) ||
+       (alloc_info.usage & VK_BUFFER_USAGE_SHADER_BINDING_TABLE_BIT_KHR))) {
     buffer_info.usage |= VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT_KHR;
   }