[Vulkan] Device API Multi-streams, multi-queue, and initial multi-thread support

taichi/backends/device.cpp

@@ -25,25 +25,28 @@ void Device::memcpy(DevicePtr dst, DevicePtr src, uint64_t size) {
 void GraphicsDevice::image_transition(DeviceAllocation img,
                                       ImageLayout old_layout,
                                       ImageLayout new_layout) {
-  auto cmd_list = new_command_list({CommandListType::Graphics});
+  Stream *stream = get_graphics_stream();
+  auto cmd_list = stream->new_command_list();
   cmd_list->image_transition(img, old_layout, new_layout);
-  submit_synced(cmd_list.get());
+  stream->submit_synced(cmd_list.get());
 }
 void GraphicsDevice::buffer_to_image(DeviceAllocation dst_img,
                                      DevicePtr src_buf,
                                      ImageLayout img_layout,
                                      const BufferImageCopyParams &params) {
-  auto cmd_list = new_command_list({CommandListType::Graphics});
+  Stream *stream = get_graphics_stream();
+  auto cmd_list = stream->new_command_list();
   cmd_list->buffer_to_image(dst_img, src_buf, img_layout, params);
-  submit_synced(cmd_list.get());
+  stream->submit_synced(cmd_list.get());
 }
 void GraphicsDevice::image_to_buffer(DevicePtr dst_buf,
                                      DeviceAllocation src_img,
                                      ImageLayout img_layout,
                                      const BufferImageCopyParams &params) {
-  auto cmd_list = new_command_list({CommandListType::Graphics});
+  Stream *stream = get_graphics_stream();
+  auto cmd_list = stream->new_command_list();
   cmd_list->image_to_buffer(dst_buf, src_img, img_layout, params);
-  submit_synced(cmd_list.get());
+  stream->submit_synced(cmd_list.get());
 }
 
 }  // namespace lang

taichi/backends/device.h

@@ -203,12 +203,6 @@ class Pipeline {
   virtual ResourceBinder *resource_binder() = 0;
 };
 
-enum class CommandListType { Graphics, Compute };
-
-struct CommandListConfig {
-  CommandListType type;
-};
-
 enum class ImageDimension { d1D, d2D, d3D };
 
 enum class ImageLayout {
@@ -327,6 +321,18 @@ inline bool operator&(AllocUsage a, AllocUsage b) {
   return static_cast<int>(a) & static_cast<int>(b);
 }
 
+class Stream {
+ public:
+  virtual ~Stream(){};
+
+  virtual std::unique_ptr<CommandList> new_command_list() = 0;
+  virtual void dealloc_command_list(CommandList *cmdlist) = 0;
+  virtual void submit(CommandList *cmdlist) = 0;
+  virtual void submit_synced(CommandList *cmdlist) = 0;
+
+  virtual void command_sync() = 0;
+};
+
 class Device {
  public:
   virtual ~Device(){};
@@ -378,14 +384,9 @@ class Device {
   // Copy memory inter or intra devices (synced)
   static void memcpy(DevicePtr dst, DevicePtr src, uint64_t size);
 
-  // TODO: Add a flag to select graphics / compute pool
-  virtual std::unique_ptr<CommandList> new_command_list(
-      CommandListConfig config) = 0;
-  virtual void dealloc_command_list(CommandList *cmdlist) = 0;
-  virtual void submit(CommandList *cmdlist) = 0;
-  virtual void submit_synced(CommandList *cmdlist) = 0;
-
-  virtual void command_sync() = 0;
+  // Each thraed will acquire its own stream
+  virtual Stream *get_compute_stream() = 0;
+  virtual Stream *get_graphics_stream() = 0;
 
  private:
   std::unordered_map<DeviceCapability, uint32_t> caps_;

taichi/backends/vulkan/runtime.cpp

@@ -133,6 +133,22 @@ class HostDeviceContextBlitter {
       return;
     }
 
+    bool require_sync = ctx_attribs_->rets().size() > 0;
+    if (!require_sync) {
+      for (int i = 0; i < ctx_attribs_->args().size(); ++i) {
+        const auto &arg = ctx_attribs_->args()[i];
+        if (arg.is_array) {
+          require_sync = true;
+        }
+      }
+    }
+
+    if (require_sync) {
+      device_->get_compute_stream()->command_sync();
+    } else {
+      return;
+    }
+
     char *const device_base =
         reinterpret_cast<char *>(device_->map(*host_shadow_buffer_));
 
@@ -231,8 +247,9 @@ class CompiledTaichiKernel {
   };
 
   CompiledTaichiKernel(const Params &ti_params)
-      : ti_kernel_attribs_(*ti_params.ti_kernel_attribs) {
-    InputBuffersMap input_buffers = {
+      : ti_kernel_attribs_(*ti_params.ti_kernel_attribs),
+        device_(ti_params.device) {
+    input_buffers_ = {
         {BufferEnum::Root, ti_params.root_buffer},
         {BufferEnum::GlobalTmps, ti_params.global_tmps_buffer},
     };
@@ -241,44 +258,27 @@ class CompiledTaichiKernel {
       Device::AllocParams params;
       ctx_buffer_ = ti_params.device->allocate_memory_unique(
           {size_t(ctx_sz),
-           /*host_write*/ true, /*host_read*/ false});
+           /*host_write=*/true, /*host_read=*/false,
+           /*export_sharing=*/false, AllocUsage::Storage});
       ctx_buffer_host_ = ti_params.device->allocate_memory_unique(
           {size_t(ctx_sz),
-           /*host_write*/ false, /*host_read*/ true});
-      input_buffers[BufferEnum::Context] = ctx_buffer_.get();
+           /*host_write=*/false, /*host_read=*/true,
+           /*export_sharing=*/false, AllocUsage::Storage});
+      input_buffers_[BufferEnum::Context] = ctx_buffer_.get();
     }
 
     const auto &task_attribs = ti_kernel_attribs_.tasks_attribs;
     const auto &spirv_bins = ti_params.spirv_bins;
     TI_ASSERT(task_attribs.size() == spirv_bins.size());
 
-    cmdlist_ = ti_params.device->new_command_list({CommandListType::Compute});
     for (int i = 0; i < task_attribs.size(); ++i) {
-      const auto &attribs = task_attribs[i];
       PipelineSourceDesc source_desc{PipelineSourceType::spirv_binary,
                                      (void *)spirv_bins[i].data(),
                                      spirv_bins[i].size() * sizeof(uint32_t)};
       auto vp = ti_params.device->create_pipeline(source_desc,
                                                   ti_kernel_attribs_.name);
-      const int group_x = (attribs.advisory_total_num_threads +
-                           attribs.advisory_num_threads_per_group - 1) /
-                          attribs.advisory_num_threads_per_group;
-      ResourceBinder *binder = vp->resource_binder();
-      for (auto &pair : input_buffers) {
-        binder->rw_buffer(0, uint32_t(pair.first), *pair.second);
-      }
-      cmdlist_->bind_pipeline(vp.get());
-      cmdlist_->bind_resources(binder);
-      cmdlist_->dispatch(group_x);
-      cmdlist_->memory_barrier();
       pipelines_.push_back(std::move(vp));
     }
-
-    if (!ti_kernel_attribs_.ctx_attribs.empty()) {
-      cmdlist_->buffer_copy(ctx_buffer_host_->get_ptr(0),
-                            ctx_buffer_->get_ptr(0), ctx_sz);
-      cmdlist_->buffer_barrier(*ctx_buffer_host_);
-    }
   }
 
   const TaichiKernelAttributes &ti_kernel_attribs() const {
@@ -297,12 +297,40 @@ class CompiledTaichiKernel {
     return ctx_buffer_host_.get();
   }
 
-  CommandList *command_list() const {
-    return cmdlist_.get();
+  void command_list(CommandList *cmdlist) const {
+    const auto &task_attribs = ti_kernel_attribs_.tasks_attribs;
+
+    for (int i = 0; i < task_attribs.size(); ++i) {
+      const auto &attribs = task_attribs[i];
+      auto vp = pipelines_[i].get();
+      const int group_x = (attribs.advisory_total_num_threads +
+                           attribs.advisory_num_threads_per_group - 1) /
+                          attribs.advisory_num_threads_per_group;
+      ResourceBinder *binder = vp->resource_binder();
+      for (auto &pair : input_buffers_) {
+        binder->rw_buffer(0, uint32_t(pair.first), *pair.second);
+      }
+      cmdlist->bind_pipeline(vp);
+      cmdlist->bind_resources(binder);
+      cmdlist->dispatch(group_x);
+      cmdlist->memory_barrier();
+    }
+
+    const auto ctx_sz = ti_kernel_attribs_.ctx_attribs.total_bytes();
+    if (!ti_kernel_attribs_.ctx_attribs.empty()) {
+      cmdlist->buffer_copy(ctx_buffer_host_->get_ptr(0),
+                           ctx_buffer_->get_ptr(0), ctx_sz);
+      cmdlist->buffer_barrier(*ctx_buffer_host_);
+    }
   }
 
  private:
   TaichiKernelAttributes ti_kernel_attribs_;
+  std::vector<TaskAttributes> tasks_attribs_;
+
+  Device *device_;
+
+  InputBuffersMap input_buffers_;
 
   // Right now |ctx_buffer_| is allocated from a HOST_VISIBLE|COHERENT
   // memory, because we do not do computation on this buffer anyway, and it may
@@ -312,8 +340,6 @@ class CompiledTaichiKernel {
   std::unique_ptr<DeviceAllocationGuard> ctx_buffer_{nullptr};
   std::unique_ptr<DeviceAllocationGuard> ctx_buffer_host_{nullptr};
   std::vector<std::unique_ptr<Pipeline>> pipelines_;
-
-  std::unique_ptr<CommandList> cmdlist_;
 };
 
 }  // namespace
@@ -376,15 +402,22 @@ class VkRuntime ::Impl {
       ctx_blitter->host_to_device();
     }
 
-    device_->submit(ti_kernel->command_list());
+    if (!current_cmdlist_) {
+      current_cmdlist_ = device_->get_compute_stream()->new_command_list();
+    }
+
+    ti_kernel->command_list(current_cmdlist_.get());
+
     if (ctx_blitter) {
-      synchronize();
+      device_->get_compute_stream()->submit(current_cmdlist_.get());
       ctx_blitter->device_to_host();
+
+      current_cmdlist_ = nullptr;
     }
   }
 
   void synchronize() {
-    device_->command_sync();
+    device_->get_compute_stream()->command_sync();
   }
 
   Device *get_ti_device() const {
@@ -397,16 +430,23 @@ class VkRuntime ::Impl {
     size_t root_buffer_size = 64 * 1024 * 1024;
     size_t gtmp_buffer_size = 1024 * 1024;
 
-    root_buffer_ = device_->allocate_memory_unique({root_buffer_size});
-    global_tmps_buffer_ = device_->allocate_memory_unique({gtmp_buffer_size});
+    root_buffer_ = device_->allocate_memory_unique(
+        {root_buffer_size,
+         /*host_write=*/false, /*host_read=*/false,
+         /*export_sharing=*/false, AllocUsage::Storage});
+    global_tmps_buffer_ = device_->allocate_memory_unique(
+        {gtmp_buffer_size,
+         /*host_write=*/false, /*host_read=*/false,
+         /*export_sharing=*/false, AllocUsage::Storage});
 
     // Need to zero fill the buffers, otherwise there could be NaN.
-    auto cmdlist = device_->new_command_list({CommandListType::Compute});
+    Stream *stream = device_->get_compute_stream();
+    auto cmdlist = stream->new_command_list();
     cmdlist->buffer_fill(root_buffer_->get_ptr(0), root_buffer_size,
                          /*data=*/0);
     cmdlist->buffer_fill(global_tmps_buffer_->get_ptr(0), gtmp_buffer_size,
                          /*data=*/0);
-    device_->submit_synced(cmdlist.get());
+    stream->submit_synced(cmdlist.get());
   }
 
   const SNodeDescriptorsMap *const snode_descriptors_;
@@ -419,6 +459,8 @@ class VkRuntime ::Impl {
 
   Device *device_;
 
+  std::unique_ptr<CommandList> current_cmdlist_{nullptr};
+
   std::vector<std::unique_ptr<CompiledTaichiKernel>> ti_kernels_;
 };
 

taichi/backends/vulkan/spirv_ir_builder.cpp

@@ -314,6 +314,14 @@ SType IRBuilder::get_struct_array_type(const SType &value_type,
     TI_ERROR("buffer type must be primitive or snode struct");
   }
 
+  if (nbytes == 0) {
+    if (value_type.flag == TypeKind::kPrimitive) {
+      TI_WARN("Invalid primitive bit size");
+    } else {
+      TI_WARN("Invalid container stride");
+    }
+  }
+
   // decorate the array type
   this->decorate(spv::OpDecorate, arr_type, spv::DecorationArrayStride, nbytes);
   // declare struct of array