[gui] GGUI Image IO (well, it's actually just O...)

taichi/backends/device.h

@@ -219,7 +219,8 @@ enum class ImageLayout {
   depth_attachment,
   depth_attachment_read,
   transfer_dst,
-  transfer_src
+  transfer_src,
+  present_src
 };
 
 struct BufferImageCopyParams {
@@ -240,6 +241,12 @@ struct BufferImageCopyParams {
   uint32_t image_layer_count{1};
 };
 
+struct ImageCopyParams {
+  uint32_t width{1};
+  uint32_t height{1};
+  uint32_t depth{1};
+};
+
 class CommandList {
  public:
   virtual ~CommandList() {
@@ -303,6 +310,20 @@ class CommandList {
                                const BufferImageCopyParams &params) {
     TI_NOT_IMPLEMENTED
   }
+  virtual void copy_image(DeviceAllocation dst_img,
+                          DeviceAllocation src_img,
+                          ImageLayout dst_img_layout,
+                          ImageLayout src_img_layout,
+                          const ImageCopyParams &params) {
+    TI_NOT_IMPLEMENTED
+  }
+  virtual void blit_image(DeviceAllocation dst_img,
+                          DeviceAllocation src_img,
+                          ImageLayout dst_img_layout,
+                          ImageLayout src_img_layout,
+                          const ImageCopyParams &params) {
+    TI_NOT_IMPLEMENTED
+  }
 };
 
 struct PipelineSourceDesc {
@@ -421,6 +442,9 @@ class Surface {
   virtual std::pair<uint32_t, uint32_t> get_size() = 0;
   virtual BufferFormat image_format() = 0;
   virtual void resize(uint32_t width, uint32_t height) = 0;
+  virtual DeviceAllocation get_image_data() {
+    TI_NOT_IMPLEMENTED
+  }
 };
 
 struct VertexInputBinding {

taichi/backends/vulkan/vulkan_device.cpp

@@ -95,7 +95,8 @@ const std::unordered_map<ImageLayout, VkImageLayout> image_layout_ti_2_vk = {
     {ImageLayout::depth_attachment_read,
      VK_IMAGE_LAYOUT_DEPTH_READ_ONLY_OPTIMAL},
     {ImageLayout::transfer_dst, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL},
-    {ImageLayout::transfer_src, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL}};
+    {ImageLayout::transfer_src, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL},
+    {ImageLayout::present_src, VK_IMAGE_LAYOUT_PRESENT_SRC_KHR}};
 
 VkImageLayout image_layout_ti_to_vk(ImageLayout layout) {
   if (image_layout_ti_2_vk.find(layout) == image_layout_ti_2_vk.end()) {
@@ -956,17 +957,21 @@ void VulkanCommandList::image_transition(DeviceAllocation img,
   static std::unordered_map<VkImageLayout, VkPipelineStageFlagBits> stages;
   stages[VK_IMAGE_LAYOUT_UNDEFINED] = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
   stages[VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL] = VK_PIPELINE_STAGE_TRANSFER_BIT;
+  stages[VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL] = VK_PIPELINE_STAGE_TRANSFER_BIT;
   stages[VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL] =
       VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
   stages[VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL] =
       VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
+  stages[VK_IMAGE_LAYOUT_PRESENT_SRC_KHR] = VK_PIPELINE_STAGE_TRANSFER_BIT;
 
   static std::unordered_map<VkImageLayout, VkAccessFlagBits> access;
   access[VK_IMAGE_LAYOUT_UNDEFINED] = (VkAccessFlagBits)0;
   access[VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL] = VK_ACCESS_TRANSFER_WRITE_BIT;
+  access[VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL] = VK_ACCESS_TRANSFER_READ_BIT;
   access[VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL] = VK_ACCESS_SHADER_READ_BIT;
   access[VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL] =
       VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
+  access[VK_IMAGE_LAYOUT_PRESENT_SRC_KHR] = VK_ACCESS_MEMORY_READ_BIT;
 
   if (stages.find(old_layout) == stages.end() ||
       stages.find(new_layout) == stages.end()) {
@@ -1040,6 +1045,60 @@ void VulkanCommandList::image_to_buffer(DevicePtr dst_buf,
   buffer_->refs.push_back(buffer);
 }
 
+void VulkanCommandList::copy_image(DeviceAllocation dst_img,
+                                   DeviceAllocation src_img,
+                                   ImageLayout dst_img_layout,
+                                   ImageLayout src_img_layout,
+                                   const ImageCopyParams &params) {
+  VkImageCopy copy{};
+  copy.srcSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
+  copy.srcSubresource.layerCount = 1;
+  copy.dstSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
+  copy.dstSubresource.layerCount = 1;
+  copy.extent.width = params.width;
+  copy.extent.height = params.height;
+  copy.extent.depth = params.depth;
+
+  auto [dst_vk_image, dst_view, dst_format] = ti_device_->get_vk_image(dst_img);
+  auto [src_vk_image, src_view, src_format] = ti_device_->get_vk_image(src_img);
+
+  vkCmdCopyImage(buffer_->buffer, src_vk_image->image,
+                 image_layout_ti_to_vk(src_img_layout), dst_vk_image->image,
+                 image_layout_ti_to_vk(dst_img_layout), 1, &copy);
+
+  buffer_->refs.push_back(dst_vk_image);
+  buffer_->refs.push_back(src_vk_image);
+}
+
+void VulkanCommandList::blit_image(DeviceAllocation dst_img,
+                                   DeviceAllocation src_img,
+                                   ImageLayout dst_img_layout,
+                                   ImageLayout src_img_layout,
+                                   const ImageCopyParams &params) {
+  VkOffset3D blit_size;
+  blit_size.x = params.width;
+  blit_size.y = params.height;
+  blit_size.z = params.depth;
+  VkImageBlit blit{};
+  blit.srcSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
+  blit.srcSubresource.layerCount = 1;
+  blit.srcOffsets[1] = blit_size;
+  blit.dstSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
+  blit.dstSubresource.layerCount = 1;
+  blit.dstOffsets[1] = blit_size;
+
+  auto [dst_vk_image, dst_view, dst_format] = ti_device_->get_vk_image(dst_img);
+  auto [src_vk_image, src_view, src_format] = ti_device_->get_vk_image(src_img);
+
+  vkCmdBlitImage(buffer_->buffer, src_vk_image->image,
+                 image_layout_ti_to_vk(src_img_layout), dst_vk_image->image,
+                 image_layout_ti_to_vk(dst_img_layout), 1, &blit,
+                 VK_FILTER_NEAREST);
+
+  buffer_->refs.push_back(dst_vk_image);
+  buffer_->refs.push_back(src_vk_image);
+}
+
 void VulkanCommandList::set_line_width(float width) {
   vkCmdSetLineWidth(buffer_->buffer, width);
 }
@@ -1902,7 +1961,8 @@ void VulkanSurface::create_swap_chain() {
   createInfo.imageColorSpace = surface_format.colorSpace;
   createInfo.imageExtent = extent;
   createInfo.imageArrayLayers = 1;
-  createInfo.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;
+  createInfo.imageUsage =
+      VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT;
   createInfo.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;
   createInfo.queueFamilyIndexCount = 0;
   createInfo.pQueueFamilyIndices = nullptr;
@@ -1966,6 +2026,12 @@ VulkanSurface::~VulkanSurface() {
   destroy_swap_chain();
   vkDestroySemaphore(device_->vk_device(), image_available_, nullptr);
   vkDestroySurfaceKHR(device_->vk_instance(), surface_, nullptr);
+  if (screenshot_buffer_ != kDeviceNullAllocation) {
+    device_->dealloc_memory(screenshot_buffer_);
+  }
+  if (screenshot_image_ != kDeviceNullAllocation) {
+    device_->destroy_image(screenshot_image_);
+  }
 }
 
 void VulkanSurface::resize(uint32_t width, uint32_t height) {
@@ -2006,6 +2072,55 @@ void VulkanSurface::present_image() {
   vkQueuePresentKHR(device_->graphics_queue(), &presentInfo);
 }
 
+DeviceAllocation VulkanSurface::get_image_data() {
+  auto *stream = device_->get_graphics_stream();
+  DeviceAllocation &img_alloc = swapchain_images_[image_index_];
+  auto [w, h] = get_size();
+  size_t size_bytes = w * h * 4;
+  if (screenshot_image_ == kDeviceNullAllocation) {
+    ImageParams params = {ImageDimension::d2D,
+                          BufferFormat::rgba8,
+                          ImageLayout::transfer_dst,
+                          w,
+                          h,
+                          1,
+                          false};
+    screenshot_image_ = device_->create_image(params);
+  }
+  if (screenshot_buffer_ == kDeviceNullAllocation) {
+    Device::AllocParams params{size_bytes, /*host_wrtie*/ false,
+                               /*host_read*/ true, /*export_sharing*/ false,
+                               AllocUsage::Uniform};
+    screenshot_buffer_ = device_->allocate_memory(params);
+  }
+
+  device_->image_transition(img_alloc, ImageLayout::present_src,
+                            ImageLayout::transfer_src);
+
+  auto cmd_list = stream->new_command_list();
+  // TODO: check if blit is suppoted, and use copy_image if not
+  cmd_list->blit_image(screenshot_image_, img_alloc, ImageLayout::transfer_dst,
+                       ImageLayout::transfer_src, {w, h, 1});
+  cmd_list->image_transition(screenshot_image_, ImageLayout::transfer_dst,
+                             ImageLayout::transfer_src);
+  stream->submit_synced(cmd_list.get());
+
+  BufferImageCopyParams copy_params;
+  copy_params.image_extent.x = w;
+  copy_params.image_extent.y = h;
+  cmd_list = stream->new_command_list();
+  // TODO: directly map the image to cpu memory
+  cmd_list->image_to_buffer(screenshot_buffer_.get_ptr(), screenshot_image_,
+                            ImageLayout::transfer_src, copy_params);
+  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_;
+}
+
 VulkanStream::VulkanStream(VulkanDevice &device,
                            VkQueue queue,
                            uint32_t queue_family_index)

taichi/backends/vulkan/vulkan_device.h

@@ -314,6 +314,18 @@ class VulkanCommandList : public CommandList {
                        ImageLayout img_layout,
                        const BufferImageCopyParams &params) override;
 
+  void copy_image(DeviceAllocation dst_img,
+                  DeviceAllocation src_img,
+                  ImageLayout dst_img_layout,
+                  ImageLayout src_img_layout,
+                  const ImageCopyParams &params) override;
+
+  void blit_image(DeviceAllocation dst_img,
+                  DeviceAllocation src_img,
+                  ImageLayout dst_img_layout,
+                  ImageLayout src_img_layout,
+                  const ImageCopyParams &params) override;
+
   vkapi::IVkRenderPass current_renderpass();
 
   // Vulkan specific functions
@@ -348,6 +360,8 @@ class VulkanSurface : public Surface {
   BufferFormat image_format() override;
   virtual void resize(uint32_t width, uint32_t height);
 
+  DeviceAllocation get_image_data() override;
+
  private:
   void create_swap_chain();
   void destroy_swap_chain();
@@ -364,6 +378,9 @@ class VulkanSurface : public Surface {
   uint32_t image_index_{0};
 
   std::vector<DeviceAllocation> swapchain_images_;
+
+  DeviceAllocation screenshot_image_{kDeviceNullAllocation};
+  DeviceAllocation screenshot_buffer_{kDeviceNullAllocation};
 };
 
 struct DescPool {

taichi/python/export_ggui.cpp

@@ -244,6 +244,10 @@ struct PyWindow {
     window = new vulkan::Window(config);
   }
 
+  void write_image(const std::string &filename) {
+    window->write_image(filename);
+  }
+
   void show() {
     window->show();
   }
@@ -305,6 +309,7 @@ void export_ggui(py::module &m) {
       .def(py::init<std::string, py::tuple, bool, std::string, Arch, bool>())
       .def("get_canvas", &PyWindow::get_canvas)
       .def("show", &PyWindow::show)
+      .def("write_image", &PyWindow::write_image)
       .def("is_pressed", &PyWindow::is_pressed)
       .def("get_cursor_pos", &PyWindow::py_get_cursor_pos)
       .def("is_running", &PyWindow::is_running)

taichi/ui/backends/vulkan/swap_chain.cpp

@@ -1,6 +1,7 @@
 #include "taichi/ui/utils/utils.h"
 #include "taichi/ui/backends/vulkan/app_context.h"
 #include "taichi/ui/backends/vulkan/swap_chain.h"
+#include "taichi/util/image_io.h"
 
 TI_UI_NAMESPACE_BEGIN
 
@@ -62,6 +63,14 @@ taichi::lang::Surface &SwapChain::surface() {
   return *(surface_.get());
 }
 
+void SwapChain::write_image(const std::string &filename) {
+  auto [w, h] = surface_->get_size();
+  DeviceAllocation img_buffer = surface_->get_image_data();
+  unsigned char *ptr = (unsigned char *)app_context_->device().map(img_buffer);
+  imwrite(filename, (size_t)ptr, w, h, 4);
+  app_context_->device().unmap(img_buffer);
+}
+
 }  // namespace vulkan
 
 TI_UI_NAMESPACE_END

taichi/ui/backends/vulkan/swap_chain.h

@@ -16,6 +16,8 @@ class SwapChain {
 
   void resize(uint32_t width, uint32_t height);
 
+  void write_image(const std::string &filename);
+
   void cleanup();
 
  private:

taichi/ui/backends/vulkan/window.cpp

@@ -20,7 +20,9 @@ void Window::init(const AppConfig &config) {
 }
 
 void Window::show() {
-  draw_frame();
+  if (!drawn_frame_) {
+    draw_frame();
+  }
   present_frame();
   WindowBase::show();
   prepare_for_next_frame();
@@ -29,6 +31,7 @@ void Window::show() {
 void Window::prepare_for_next_frame() {
   renderer_->prepare_for_next_frame();
   gui_->prepare_for_next_frame();
+  drawn_frame_ = false;
 }
 
 CanvasBase *Window::get_canvas() {
@@ -67,6 +70,7 @@ void Window::resize() {
 
 void Window::draw_frame() {
   renderer_->draw_frame(gui_.get());
+  drawn_frame_ = true;
 }
 
 void Window::present_frame() {
@@ -79,6 +83,13 @@ Window::~Window() {
   glfwTerminate();
 }
 
+void Window::write_image(const std::string &filename) {
+  if (!drawn_frame_) {
+    draw_frame();
+  }
+  renderer_->swap_chain().write_image(filename);
+}
+
 }  // namespace vulkan
 
 TI_UI_NAMESPACE_END

taichi/ui/backends/vulkan/window.h

@@ -34,12 +34,15 @@ class Window final : public WindowBase {
   virtual CanvasBase *get_canvas() override;
   virtual GuiBase *GUI() override;
 
+  void write_image(const std::string &filename) override;
+
   ~Window();
 
  private:
   std::unique_ptr<Canvas> canvas_;
   std::unique_ptr<Gui> gui_;
   std::unique_ptr<Renderer> renderer_;
+  bool drawn_frame_{false};
 
  private:
   void init(const AppConfig &config);

taichi/ui/common/window_base.h

@@ -39,6 +39,8 @@ class WindowBase {
 
   virtual void show();
 
+  virtual void write_image(const std::string &filename) = 0;
+
   virtual GuiBase *GUI();
 
   virtual ~WindowBase();