[Vulkan] Track image layout internally

taichi/program/program.h

@@ -304,6 +304,9 @@ class TI_DLL_EXPORT Program {
                                            uint64 *result_buffer) {
     return program_impl_->allocate_memory_ndarray(alloc_size, result_buffer);
   }
+  DeviceAllocation allocate_texture(const ImageParams &params) {
+    return program_impl_->allocate_texture(params);
+  }
 
   Ndarray *create_ndarray(
       const DataType type,

taichi/program/program_impl.h

@@ -102,6 +102,11 @@ class ProgramImpl {
                                                    uint64 *result_buffer) {
     return kDeviceNullAllocation;
   }
+
+  virtual DeviceAllocation allocate_texture(const ImageParams &params) {
+    return kDeviceNullAllocation;
+  }
+
   virtual ~ProgramImpl() {
   }
 

taichi/program/texture.cpp

@@ -35,7 +35,7 @@ Texture::Texture(Program *prog,
   img_params.y = height;
   img_params.z = depth;
   img_params.initial_layout = ImageLayout::undefined;
-  texture_alloc_ = device->create_image(img_params);
+  texture_alloc_ = prog_->allocate_texture(img_params);
 
   format_ = img_params.format;
 
@@ -116,6 +116,10 @@ void Texture::from_ndarray(Ndarray *ndarray) {
 
   GraphicsDevice *device =
       static_cast<GraphicsDevice *>(prog_->get_graphics_device());
+
+  device->image_transition(texture_alloc_, ImageLayout::undefined,
+                           ImageLayout::transfer_dst);
+
   Stream *stream = device->get_compute_stream();
   auto cmdlist = stream->new_command_list();
 
@@ -128,8 +132,6 @@ void Texture::from_ndarray(Ndarray *ndarray) {
   params.image_extent.z = depth_;
 
   cmdlist->buffer_barrier(ndarray->ndarray_alloc_);
-  cmdlist->image_transition(texture_alloc_, ImageLayout::undefined,
-                            ImageLayout::transfer_dst);
   cmdlist->buffer_to_image(texture_alloc_, ndarray->ndarray_alloc_.get_ptr(0),
                            ImageLayout::transfer_dst, params);
 
@@ -154,6 +156,9 @@ void Texture::from_snode(SNode *snode) {
   GraphicsDevice *device =
       static_cast<GraphicsDevice *>(prog_->get_graphics_device());
 
+  device->image_transition(texture_alloc_, ImageLayout::undefined,
+                           ImageLayout::transfer_dst);
+
   DevicePtr devptr = get_device_ptr(prog_, snode);
 
   Stream *stream = device->get_compute_stream();
@@ -168,8 +173,6 @@ void Texture::from_snode(SNode *snode) {
   params.image_extent.z = depth_;
 
   cmdlist->buffer_barrier(devptr);
-  cmdlist->image_transition(texture_alloc_, ImageLayout::undefined,
-                            ImageLayout::transfer_dst);
   cmdlist->buffer_to_image(texture_alloc_, devptr, ImageLayout::transfer_dst,
                            params);
 

taichi/rhi/device.h

@@ -600,7 +600,7 @@ struct SurfaceConfig {
 struct TI_DLL_EXPORT ImageParams {
   ImageDimension dimension;
   BufferFormat format;
-  ImageLayout initial_layout;
+  ImageLayout initial_layout{ImageLayout::undefined};
   uint32_t x{1};
   uint32_t y{1};
   uint32_t z{1};
@@ -642,6 +642,10 @@ class TI_DLL_EXPORT GraphicsDevice : public Device {
 
   virtual std::unique_ptr<Surface> create_surface(
       const SurfaceConfig &config) = 0;
+  // You are not expected to call this directly. If you want to use this image
+  // in a taichi kernel, you usually want to create the image via
+  // `GfxRuntime::create_image`. `GfxRuntime` is available in `ProgramImpl`
+  // of GPU backends.
   virtual DeviceAllocation create_image(const ImageParams &params) = 0;
   virtual void destroy_image(DeviceAllocation handle) = 0;
 

taichi/rhi/vulkan/vulkan_device.cpp

@@ -2505,6 +2505,11 @@ void VulkanSurface::present_image(
     const std::vector<StreamSemaphore> &wait_semaphores) {
   std::vector<VkSemaphore> vk_wait_semaphores;
 
+  // Already transitioned to `present_src` at the end of the render pass.
+  // device_->image_transition(get_target_image(),
+  // ImageLayout::color_attachment,
+  //                          ImageLayout::present_src);
+
   for (const StreamSemaphore &sema_ : wait_semaphores) {
     auto sema = std::static_pointer_cast<VulkanStreamSemaphoreObject>(sema_);
     vk_wait_semaphores.push_back(sema->vkapi_ref->semaphore);
@@ -2537,20 +2542,19 @@ DeviceAllocation VulkanSurface::get_depth_data(DeviceAllocation &depth_alloc) {
     depth_buffer_ = device_->allocate_memory(params);
   }
 
-  device_->image_transition(depth_alloc, ImageLayout::present_src,
-                            ImageLayout::transfer_src);
-
   std::unique_ptr<CommandList> cmd_list{nullptr};
 
   BufferImageCopyParams copy_params;
   copy_params.image_extent.x = w;
   copy_params.image_extent.y = h;
   copy_params.image_aspect_flag = VK_IMAGE_ASPECT_DEPTH_BIT;
   cmd_list = stream->new_command_list();
+  cmd_list->image_transition(depth_alloc, ImageLayout::depth_attachment,
+                             ImageLayout::transfer_src);
   cmd_list->image_to_buffer(depth_buffer_.get_ptr(), depth_alloc,
                             ImageLayout::transfer_src, copy_params);
   cmd_list->image_transition(depth_alloc, ImageLayout::transfer_src,
-                             ImageLayout::present_src);
+                             ImageLayout::depth_attachment);
   stream->submit_synced(cmd_list.get());
 
   return depth_buffer_;
@@ -2582,9 +2586,6 @@ DeviceAllocation VulkanSurface::get_image_data() {
     screenshot_buffer_ = device_->allocate_memory(params);
   }
 
-  device_->image_transition(img_alloc, ImageLayout::present_src,
-                            ImageLayout::transfer_src);
-
   std::unique_ptr<CommandList> cmd_list{nullptr};
 
   /*
@@ -2605,17 +2606,19 @@ DeviceAllocation VulkanSurface::get_image_data() {
   copy_params.image_extent.y = h;
   copy_params.image_aspect_flag = VK_IMAGE_ASPECT_COLOR_BIT;
   cmd_list = stream->new_command_list();
+  cmd_list->image_transition(img_alloc, ImageLayout::present_src,
+                             ImageLayout::transfer_src);
   // TODO: directly map the image to cpu memory
   cmd_list->image_to_buffer(screenshot_buffer_.get_ptr(), img_alloc,
                             ImageLayout::transfer_src, copy_params);
+  cmd_list->image_transition(img_alloc, ImageLayout::transfer_src,
+                             ImageLayout::present_src);
   /*
   if (config_.window_handle) {
     cmd_list->image_transition(screenshot_image_, ImageLayout::transfer_src,
                                ImageLayout::transfer_dst);
   }
   */
-  cmd_list->image_transition(img_alloc, ImageLayout::transfer_src,
-                             ImageLayout::present_src);
   stream->submit_synced(cmd_list.get());
 
   return screenshot_buffer_;

taichi/runtime/gfx/runtime.cpp

@@ -498,12 +498,10 @@ void GfxRuntime::launch_kernel(KernelHandle handle, RuntimeContext *host_ctx) {
     for (auto &bind : attribs.texture_binds) {
       DeviceAllocation texture = textures.at(bind.arg_id);
       if (bind.is_storage) {
-        current_cmdlist_->image_transition(texture, ImageLayout::undefined,
-                                           ImageLayout::shader_read_write);
+        transition_image(texture, ImageLayout::shader_read_write);
         binder->rw_image(0, bind.binding, texture, 0);
       } else {
-        current_cmdlist_->image_transition(texture, ImageLayout::undefined,
-                                           ImageLayout::shader_read);
+        transition_image(texture, ImageLayout::shader_read);
         binder->image(0, bind.binding, texture, {});
       }
     }
@@ -561,6 +559,22 @@ void GfxRuntime::buffer_copy(DevicePtr dst, DevicePtr src, size_t size) {
   submit_current_cmdlist_if_timeout();
 }
 
+DeviceAllocation GfxRuntime::create_image(const ImageParams &params) {
+  GraphicsDevice *gfx_device = dynamic_cast<GraphicsDevice *>(device_);
+  TI_ERROR_IF(gfx_device == nullptr,
+              "Image can only be created on a graphics device");
+  DeviceAllocation image = gfx_device->create_image(params);
+  last_image_layouts_[image.alloc_id] = params.initial_layout;
+  return image;
+}
+
+void GfxRuntime::transition_image(DeviceAllocation image, ImageLayout layout) {
+  ImageLayout &last_layout = last_image_layouts_.at(image.alloc_id);
+  ensure_current_cmdlist();
+  current_cmdlist_->image_transition(image, last_layout, layout);
+  last_layout = layout;
+}
+
 void GfxRuntime::signal_event(DeviceEvent *event) {
   ensure_current_cmdlist();
   current_cmdlist_->signal_event(event);

taichi/runtime/gfx/runtime.h

@@ -102,6 +102,9 @@ class TI_DLL_EXPORT GfxRuntime {
 
   void buffer_copy(DevicePtr dst, DevicePtr src, size_t size);
 
+  DeviceAllocation create_image(const ImageParams &params);
+  void transition_image(DeviceAllocation image, ImageLayout layout);
+
   void signal_event(DeviceEvent *event);
   void reset_event(DeviceEvent *event);
   void wait_event(DeviceEvent *event);
@@ -142,6 +145,7 @@ class TI_DLL_EXPORT GfxRuntime {
   std::vector<std::unique_ptr<CompiledTaichiKernel>> ti_kernels_;
 
   std::unordered_map<DeviceAllocation *, size_t> root_buffers_size_map_;
+  std::unordered_map<DeviceAllocationId, ImageLayout> last_image_layouts_;
 };
 
 GfxRuntime::RegisterParams run_codegen(

taichi/runtime/program_impls/opengl/opengl_program.cpp

@@ -80,6 +80,10 @@ DeviceAllocation OpenglProgramImpl::allocate_memory_ndarray(
       {alloc_size, /*host_write=*/false, /*host_read=*/false,
        /*export_sharing=*/false});
 }
+DeviceAllocation OpenglProgramImpl::allocate_texture(
+    const ImageParams &params) {
+  return runtime_->create_image(params);
+}
 
 std::unique_ptr<aot::Kernel> OpenglProgramImpl::make_aot_kernel(
     Kernel &kernel) {

taichi/runtime/program_impls/opengl/opengl_program.h

@@ -43,6 +43,7 @@ class OpenglProgramImpl : public ProgramImpl {
 
   DeviceAllocation allocate_memory_ndarray(std::size_t alloc_size,
                                            uint64 *result_buffer) override;
+  DeviceAllocation allocate_texture(const ImageParams &params) override;
 
   Device *get_compute_device() override {
     return device_.get();

taichi/runtime/program_impls/vulkan/vulkan_program.cpp

@@ -193,6 +193,11 @@ DeviceAllocation VulkanProgramImpl::allocate_memory_ndarray(
        /*export_sharing=*/false});
 }
 
+DeviceAllocation VulkanProgramImpl::allocate_texture(
+    const ImageParams &params) {
+  return vulkan_runtime_->create_image(params);
+}
+
 std::unique_ptr<aot::Kernel> VulkanProgramImpl::make_aot_kernel(
     Kernel &kernel) {
   spirv::lower(&kernel);

taichi/runtime/program_impls/vulkan/vulkan_program.h

@@ -65,6 +65,7 @@ class VulkanProgramImpl : public ProgramImpl {
 
   DeviceAllocation allocate_memory_ndarray(std::size_t alloc_size,
                                            uint64 *result_buffer) override;
+  DeviceAllocation allocate_texture(const ImageParams &params) override;
 
   Device *get_compute_device() override {
     if (embedded_device_) {

taichi/ui/backends/vulkan/renderables/set_image.cpp

@@ -67,7 +67,7 @@ void SetImage::update_data(const SetImageInfo &info) {
 
   int pixels = width * height;
 
-  app_context_->device().image_transition(texture_, ImageLayout::shader_read,
+  app_context_->device().image_transition(texture_, ImageLayout::undefined,
                                           ImageLayout::transfer_dst);
 
   uint64_t img_size = pixels * data_type_size(texture_dtype_) * 4;
@@ -98,6 +98,8 @@ void SetImage::update_data(const SetImageInfo &info) {
 
   auto stream = app_context_->device().get_graphics_stream();
   auto cmd_list = stream->new_command_list();
+  cmd_list->image_transition(texture_, ImageLayout::undefined,
+                             ImageLayout::transfer_dst);
   cmd_list->buffer_to_image(texture_, gpu_staging_buffer_.get_ptr(0),
                             ImageLayout::transfer_dst, copy_params);
 

taichi/ui/backends/vulkan/renderer.cpp

@@ -154,6 +154,8 @@ void Renderer::draw_frame(Gui *gui) {
                                      background_color_[2], 1};
   auto semaphore = swap_chain_.surface().acquire_next_image();
   auto image = swap_chain_.surface().get_target_image();
+  cmd_list->image_transition(image, ImageLayout::undefined,
+                             ImageLayout::color_attachment);
   auto depth_image = swap_chain_.depth_allocation();
   cmd_list->begin_renderpass(
       /*xmin=*/0, /*ymin=*/0, /*xmax=*/swap_chain_.width(),

taichi/ui/backends/vulkan/swap_chain.cpp

@@ -64,14 +64,13 @@ bool SwapChain::copy_depth_buffer_to_ndarray(
 
     auto depth_staging_buffer = device.allocate_memory(params);
 
-    device.image_transition(depth_allocation_, ImageLayout::depth_attachment,
-                            ImageLayout::transfer_src);
-
     BufferImageCopyParams copy_params;
     copy_params.image_extent.x = w;
     copy_params.image_extent.y = h;
     copy_params.image_aspect_flag = VK_IMAGE_ASPECT_DEPTH_BIT;
     cmd_list = stream->new_command_list();
+    cmd_list->image_transition(depth_allocation_, ImageLayout::depth_attachment,
+                               ImageLayout::transfer_src);
     cmd_list->image_to_buffer(depth_staging_buffer.get_ptr(), depth_allocation_,
                               ImageLayout::transfer_src, copy_params);
     cmd_list->image_transition(depth_allocation_, ImageLayout::transfer_src,