[gui] Use Device API memcpy in GGUI (taichi-dev#3163)

* wip

* wip

* remove legacy vulkan_cuda_interop

python/taichi/ui/utils.py

@@ -10,29 +10,6 @@
 from taichi.type.primitive_types import u64
 
 
-@kernel
-def get_field_addr_0D(x: template()) -> u64:
-    return get_addr(x, [None])
-
-
-@kernel
-def get_field_addr_ND(x: template()) -> u64:
-    return get_addr(x, [0 for _ in x.shape])
-
-
-field_addr_cache = {}
-
-
-def get_field_addr(x):
-    if x not in field_addr_cache:
-        if len(x.shape) == 0:
-            addr = get_field_addr_0D(x)
-        else:
-            addr = get_field_addr_ND(x)
-        field_addr_cache[x] = addr
-    return field_addr_cache[x]
-
-
 def get_field_info(field):
     info = _ti_core.FieldInfo()
     if field is None:
@@ -48,7 +25,7 @@ def get_field_info(field):
     info.shape = [n for n in field.shape]
 
     info.dtype = field.dtype
-    info.data = get_field_addr(field)
+    info.snode = field.snode.ptr
 
     if hasattr(field, 'n'):
         info.field_type = _ti_core.FieldType.Matrix

taichi/python/export_ggui.cpp

@@ -368,7 +368,7 @@ void export_ggui(py::module &m) {
       .def_property("dtype", &FieldInfo::get_dtype, &FieldInfo::set_dtype)
       .def_property("field_source", &FieldInfo::get_field_source,
                     &FieldInfo::set_field_source)
-      .def_property("data", &FieldInfo::get_data, &FieldInfo::set_data)
+      .def_property("snode", &FieldInfo::get_snode, &FieldInfo::set_snode)
       .def_property("shape", &FieldInfo::get_shape, &FieldInfo::set_shape)
       .def_property("valid", &FieldInfo::get_valid, &FieldInfo::set_valid);
 

taichi/ui/backends/vulkan/canvas.cpp

@@ -1,9 +1,6 @@
 #include "canvas.h"
 #include "taichi/ui/utils/utils.h"
 
-#include "taichi/ui/backends/vulkan/vulkan_cuda_interop.h"
-#include "taichi/ui/backends/vulkan/vulkan_cuda_interop.h"
-
 TI_UI_NAMESPACE_BEGIN
 
 namespace vulkan {

taichi/ui/backends/vulkan/renderable.cpp

@@ -1,8 +1,6 @@
 #include "taichi/ui/backends/vulkan/renderable.h"
 #include "taichi/ui/utils/utils.h"
 
-#include "taichi/ui/backends/vulkan/vulkan_cuda_interop.h"
-
 TI_UI_NAMESPACE_BEGIN
 
 namespace vulkan {
@@ -37,23 +35,6 @@ void Renderable::init_buffers() {
   create_storage_buffers();
 
   create_bindings();
-
-  if (app_context_->config.ti_arch == Arch::cuda) {
-    auto [vb_mem, vb_offset, vb_size] =
-        app_context_->device().get_vkmemory_offset_size(vertex_buffer_);
-
-    auto [ib_mem, ib_offset, ib_size] =
-        app_context_->device().get_vkmemory_offset_size(index_buffer_);
-
-    auto block_size = VulkanDevice::kMemoryBlockSize;
-
-    vertex_buffer_device_ptr_ =
-        (float *)get_memory_pointer(vb_mem, block_size, vb_offset, vb_size,
-                                    app_context_->device().vk_device());
-    index_buffer_device_ptr_ =
-        (int *)get_memory_pointer(ib_mem, block_size, ib_offset, ib_size,
-                                  app_context_->device().vk_device());
-  }
 }
 
 void Renderable::update_data(const RenderableInfo &info) {
@@ -76,42 +57,36 @@ void Renderable::update_data(const RenderableInfo &info) {
     init_buffers();
   }
 
-  if (info.vbo.field_source == FieldSource::TaichiCuda) {
-    cuda_memcpy(vertex_buffer_device_ptr_, (void *)info.vbo.data,
-                sizeof(Vertex) * num_vertices);
-
-    if (info.indices.valid) {
-      indexed_ = true;
-      cuda_memcpy(index_buffer_device_ptr_, (int *)info.indices.data,
-                  num_indices * sizeof(int));
+  Program &program = get_current_program();
+  DevicePtr vbo_dev_ptr = get_device_ptr(&program, info.vbo.snode);
+  uint64_t vbo_size = sizeof(Vertex) * num_vertices;
+
+  Device::MemcpyCapability memcpy_cap = Device::check_memcpy_capability(
+      vertex_buffer_.get_ptr(), vbo_dev_ptr, vbo_size);
+  if (memcpy_cap == Device::MemcpyCapability::Direct) {
+    Device::memcpy_direct(vertex_buffer_.get_ptr(), vbo_dev_ptr.get_ptr(),
+                          vbo_size);
+  } else if (memcpy_cap == Device::MemcpyCapability::RequiresStagingBuffer) {
+    Device::memcpy_via_staging(vertex_buffer_.get_ptr(),
+                               staging_vertex_buffer_.get_ptr(), vbo_dev_ptr,
+                               vbo_size);
+  } else {
+    TI_NOT_IMPLEMENTED;
+  }
 
+  if (info.indices.valid) {
+    indexed_ = true;
+    DevicePtr ibo_dev_ptr = get_device_ptr(&program, info.indices.snode);
+    uint64_t ibo_size = num_indices * sizeof(int);
+    if (memcpy_cap == Device::MemcpyCapability::Direct) {
+      Device::memcpy_direct(index_buffer_.get_ptr(), ibo_dev_ptr, ibo_size);
+    } else if (memcpy_cap == Device::MemcpyCapability::RequiresStagingBuffer) {
+      Device::memcpy_via_staging(index_buffer_.get_ptr(),
+                                 staging_index_buffer_.get_ptr(), ibo_dev_ptr,
+                                 ibo_size);
     } else {
-      indexed_ = false;
-    }
-  } else if (info.vbo.field_source == FieldSource::TaichiX64) {
-    float *mapped_vbo =
-        (float *)app_context_->device().map(staging_vertex_buffer_);
-    memcpy(mapped_vbo, (void *)info.vbo.data, sizeof(Vertex) * num_vertices);
-    app_context_->device().unmap(staging_vertex_buffer_);
-
-    int *mapped_ibo = (int *)app_context_->device().map(staging_index_buffer_);
-    if (info.indices.valid) {
-      indexed_ = true;
-      memcpy(mapped_ibo, (int *)info.indices.data, num_indices * sizeof(int));
-    } else {
-      indexed_ = false;
+      TI_NOT_IMPLEMENTED;
     }
-    app_context_->device().unmap(staging_index_buffer_);
-
-    auto stream = app_context_->device().get_graphics_stream();
-    auto cmd_list = stream->new_command_list();
-    cmd_list->buffer_copy(vertex_buffer_.get_ptr(0),
-                          staging_vertex_buffer_.get_ptr(0),
-                          config_.vertices_count * sizeof(Vertex));
-    cmd_list->buffer_copy(index_buffer_.get_ptr(0),
-                          staging_index_buffer_.get_ptr(0),
-                          config_.indices_count * sizeof(int));
-    stream->submit_synced(cmd_list.get());
   }
 }
 

taichi/ui/backends/vulkan/renderable.h

@@ -16,7 +16,7 @@
 #include "taichi/ui/backends/vulkan/vertex.h"
 #include "taichi/ui/backends/vulkan/app_context.h"
 #include "taichi/ui/backends/vulkan/swap_chain.h"
-#include "taichi/ui/backends/vulkan/vulkan_cuda_interop.h"
+
 #include "taichi/ui/common/renderable_info.h"
 #include "taichi/backends/vulkan/vulkan_device.h"
 
@@ -65,9 +65,6 @@ class Renderable {
   taichi::lang::DeviceAllocation uniform_buffer_;
   taichi::lang::DeviceAllocation storage_buffer_;
 
-  float *vertex_buffer_device_ptr_;
-  int *index_buffer_device_ptr_;
-
   bool indexed_{false};
 
  protected:

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

@@ -1,5 +1,4 @@
 #include "circles.h"
-#include "taichi/ui/backends/vulkan/vulkan_cuda_interop.h"
 
 #include "taichi/ui/utils/utils.h"
 

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

@@ -1,5 +1,4 @@
 #include "lines.h"
-#include "taichi/ui/backends/vulkan/vulkan_cuda_interop.h"
 
 #include "taichi/ui/utils/utils.h"
 

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

@@ -1,5 +1,4 @@
 #include "mesh.h"
-#include "taichi/ui/backends/vulkan/vulkan_cuda_interop.h"
 
 #include "taichi/ui/utils/utils.h"
 #include "taichi/backends/vulkan/vulkan_device.h"

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

@@ -1,5 +1,4 @@
 #include "particles.h"
-#include "taichi/ui/backends/vulkan/vulkan_cuda_interop.h"
 
 #include "taichi/ui/utils/utils.h"
 

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

@@ -1,6 +1,5 @@
 #include "set_image.h"
-#include "taichi/ui/backends/vulkan/vulkan_cuda_interop.h"
-#include "taichi/ui/backends/vulkan/vulkan_cuda_interop.h"
+
 #include "taichi/ui/utils/utils.h"
 
 TI_UI_NAMESPACE_BEGIN
@@ -44,45 +43,35 @@ void SetImage::update_data(const SetImageInfo &info) {
   app_context_->device().image_transition(texture_, ImageLayout::shader_read,
                                           ImageLayout::transfer_dst);
 
+  Program &program = get_current_program();
+  DevicePtr img_dev_ptr = get_device_ptr(&program, img.snode);
+  uint64_t img_size = pixels * 4;
+
+  Device::MemcpyCapability memcpy_cap = Device::check_memcpy_capability(
+      gpu_staging_buffer_.get_ptr(), img_dev_ptr, img_size);
+  if (memcpy_cap == Device::MemcpyCapability::Direct) {
+    Device::memcpy_direct(gpu_staging_buffer_.get_ptr(), img_dev_ptr, img_size);
+  } else if (memcpy_cap == Device::MemcpyCapability::RequiresStagingBuffer) {
+    Device::memcpy_via_staging(gpu_staging_buffer_.get_ptr(),
+                               cpu_staging_buffer_.get_ptr(), img_dev_ptr,
+                               img_size);
+  } else {
+    TI_NOT_IMPLEMENTED;
+  }
+
   BufferImageCopyParams copy_params;
   // these are flipped because taichi is y-major and vulkan is x-major
   copy_params.image_extent.x = height;
   copy_params.image_extent.y = width;
 
-  if (img.field_source == FieldSource::TaichiCuda) {
-    unsigned char *mapped = device_ptr_;
-
-    cuda_memcpy(mapped, (unsigned char *)img.data, pixels * 4);
-
-    auto stream = app_context_->device().get_graphics_stream();
-    auto cmd_list = stream->new_command_list();
-    cmd_list->buffer_to_image(texture_, gpu_staging_buffer_.get_ptr(0),
-                              ImageLayout::transfer_dst, copy_params);
-
-    cmd_list->image_transition(texture_, ImageLayout::transfer_dst,
-                               ImageLayout::shader_read);
-    stream->submit_synced(cmd_list.get());
-
-  } else if (img.field_source == FieldSource::TaichiX64) {
-    unsigned char *mapped =
-        (unsigned char *)app_context_->device().map(cpu_staging_buffer_);
-
-    memcpy(mapped, (unsigned char *)img.data, pixels * 4);
+  auto stream = app_context_->device().get_graphics_stream();
+  auto cmd_list = stream->new_command_list();
+  cmd_list->buffer_to_image(texture_, gpu_staging_buffer_.get_ptr(0),
+                            ImageLayout::transfer_dst, copy_params);
 
-    app_context_->device().unmap(cpu_staging_buffer_);
-
-    auto stream = app_context_->device().get_graphics_stream();
-    auto cmd_list = stream->new_command_list();
-    cmd_list->buffer_to_image(texture_, cpu_staging_buffer_.get_ptr(0),
-                              ImageLayout::transfer_dst, copy_params);
-
-    cmd_list->image_transition(texture_, ImageLayout::transfer_dst,
-                               ImageLayout::shader_read);
-    stream->submit_synced(cmd_list.get());
-
-  } else {
-    throw std::runtime_error("unsupported field source");
-  }
+  cmd_list->image_transition(texture_, ImageLayout::transfer_dst,
+                             ImageLayout::shader_read);
+  stream->submit_synced(cmd_list.get());
 }
 
 SetImage::SetImage(AppContext *app_context) {
@@ -139,16 +128,6 @@ void SetImage::create_texture() {
                                                 AllocUsage::Uniform};
   gpu_staging_buffer_ =
       app_context_->device().allocate_memory(gpu_staging_buffer_params);
-
-  if (app_context_->config.ti_arch == Arch::cuda) {
-    auto [mem, offset, size] =
-        app_context_->device().get_vkmemory_offset_size(gpu_staging_buffer_);
-
-    auto block_size = VulkanDevice::kMemoryBlockSize;
-
-    device_ptr_ = (unsigned char *)get_memory_pointer(
-        mem, block_size, offset, size, app_context_->device().vk_device());
-  }
 }
 
 void SetImage::destroy_texture() {

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

@@ -49,8 +49,6 @@ class SetImage final : public Renderable {
 
   taichi::lang::DeviceAllocation texture_;
 
-  unsigned char *device_ptr_{nullptr};
-
  private:
   void init_set_image(AppContext *app_context, int img_width, int img_height);
 

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

@@ -1,5 +1,4 @@
 #include "triangles.h"
-#include "taichi/ui/backends/vulkan/vulkan_cuda_interop.h"
 
 #include "taichi/ui/utils/utils.h"
 

taichi/ui/backends/vulkan/renderer.cpp

@@ -1,9 +1,6 @@
 #include "renderer.h"
 #include "taichi/ui/utils/utils.h"
 
-#include "taichi/ui/backends/vulkan/vulkan_cuda_interop.h"
-#include "taichi/ui/backends/vulkan/vulkan_cuda_interop.h"
-
 TI_UI_NAMESPACE_BEGIN
 
 namespace vulkan {
@@ -110,10 +107,6 @@ void Renderer::prepare_for_next_frame() {
 void Renderer::draw_frame(Gui *gui) {
   uint32_t image_index = 0;
 
-  if (app_context_.config.ti_arch == Arch::cuda) {
-    CUDADriver::get_instance().stream_synchronize(nullptr);
-  }
-
   auto stream = app_context_.device().get_graphics_stream();
   auto cmd_list = stream->new_command_list();
   bool color_clear = true;