[rhi] Update Stream new_command_list API (#7073)

Issue: #6832

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

cpp_examples/rhi_examples/sample_2_triangle.cpp

@@ -84,7 +84,9 @@ class SampleApp : public App {
 
   std::vector<StreamSemaphore> render_loop(
       StreamSemaphore image_available_semaphore) override {
-    auto cmdlist = device->get_graphics_stream()->new_command_list();
+    auto [cmdlist, res] =
+        device->get_graphics_stream()->new_command_list_unique();
+    TI_ASSERT(res == RhiResult::success);
 
     // Set-up our frame buffer attachment
     DeviceAllocation surface_image = surface->get_target_image();
@@ -100,7 +102,9 @@ class SampleApp : public App {
 
     // Bind our triangle pipeline
     cmdlist->bind_pipeline(pipeline.get());
-    cmdlist->bind_raster_resources(raster_resources.get());
+    res = cmdlist->bind_raster_resources(raster_resources.get());
+    TI_ASSERT_INFO(res == RhiResult::success,
+                   "Raster res bind fault: RhiResult({})", res);
     // Render the triangle
     cmdlist->draw(3, 0);
     // End rendering

taichi/program/program.cpp

@@ -466,7 +466,8 @@ Ndarray *Program::create_ndarray(const DataType type,
       // Device api support for dx12 & metal backend are not complete yet
       Stream *stream =
           program_impl_->get_compute_device()->get_compute_stream();
-      auto cmdlist = stream->new_command_list();
+      auto [cmdlist, res] = stream->new_command_list_unique();
+      TI_ASSERT(res == RhiResult::success);
       cmdlist->buffer_fill(arr->ndarray_alloc_.get_ptr(0),
                            arr->get_element_size() * arr->get_nelement(),
                            /*data=*/0);

taichi/program/texture.cpp

@@ -209,7 +209,8 @@ void Texture::from_ndarray(Ndarray *ndarray) {
                            ImageLayout::transfer_dst);
 
   Stream *stream = device->get_compute_stream();
-  auto cmdlist = stream->new_command_list();
+  auto [cmdlist, res] = stream->new_command_list_unique();
+  TI_ASSERT(res == RhiResult::success);
 
   BufferImageCopyParams params;
   params.buffer_row_length = ndarray->shape[0];
@@ -250,7 +251,8 @@ void Texture::from_snode(SNode *snode) {
   DevicePtr devptr = get_device_ptr(prog_, snode);
 
   Stream *stream = device->get_compute_stream();
-  auto cmdlist = stream->new_command_list();
+  auto [cmdlist, res] = stream->new_command_list_unique();
+  TI_ASSERT(res == RhiResult::success);
 
   BufferImageCopyParams params;
   params.buffer_row_length = snode->shape_along_axis(0);

taichi/rhi/amdgpu/amdgpu_device.h

@@ -41,7 +41,8 @@ class AmdgpuStream : public Stream {
  public:
   ~AmdgpuStream() override{};
 
-  std::unique_ptr<CommandList> new_command_list() override{TI_NOT_IMPLEMENTED};
+  RhiResult new_command_list(CommandList **out_cmdlist) noexcept final{
+      TI_NOT_IMPLEMENTED};
   StreamSemaphore submit(CommandList *cmdlist,
                          const std::vector<StreamSemaphore> &wait_semaphores =
                              {}) override{TI_NOT_IMPLEMENTED};

taichi/rhi/cpu/cpu_device.h

@@ -46,7 +46,8 @@ class CpuStream : public Stream {
  public:
   ~CpuStream() override{};
 
-  std::unique_ptr<CommandList> new_command_list() override{TI_NOT_IMPLEMENTED};
+  RhiResult new_command_list(CommandList **out_cmdlist) noexcept override{
+      TI_NOT_IMPLEMENTED};
   StreamSemaphore submit(CommandList *cmdlist,
                          const std::vector<StreamSemaphore> &wait_semaphores =
                              {}) override{TI_NOT_IMPLEMENTED};

taichi/rhi/cuda/cuda_device.h

@@ -46,7 +46,8 @@ class CudaStream : public Stream {
  public:
   ~CudaStream() override{};
 
-  std::unique_ptr<CommandList> new_command_list() override{TI_NOT_IMPLEMENTED};
+  RhiResult new_command_list(CommandList **out_cmdlist) noexcept final{
+      TI_NOT_IMPLEMENTED};
   StreamSemaphore submit(CommandList *cmdlist,
                          const std::vector<StreamSemaphore> &wait_semaphores =
                              {}) override{TI_NOT_IMPLEMENTED};

taichi/rhi/device.cpp

@@ -111,7 +111,8 @@ void GraphicsDevice::image_transition(DeviceAllocation img,
                                       ImageLayout old_layout,
                                       ImageLayout new_layout) {
   Stream *stream = get_graphics_stream();
-  auto cmd_list = stream->new_command_list();
+  auto [cmd_list, res] = stream->new_command_list_unique();
+  TI_ASSERT(res == RhiResult::success);
   cmd_list->image_transition(img, old_layout, new_layout);
   stream->submit_synced(cmd_list.get());
 }
@@ -120,7 +121,8 @@ void GraphicsDevice::buffer_to_image(DeviceAllocation dst_img,
                                      ImageLayout img_layout,
                                      const BufferImageCopyParams &params) {
   Stream *stream = get_graphics_stream();
-  auto cmd_list = stream->new_command_list();
+  auto [cmd_list, res] = stream->new_command_list_unique();
+  TI_ASSERT(res == RhiResult::success);
   cmd_list->buffer_to_image(dst_img, src_buf, img_layout, params);
   stream->submit_synced(cmd_list.get());
 }
@@ -129,7 +131,8 @@ void GraphicsDevice::image_to_buffer(DevicePtr dst_buf,
                                      ImageLayout img_layout,
                                      const BufferImageCopyParams &params) {
   Stream *stream = get_graphics_stream();
-  auto cmd_list = stream->new_command_list();
+  auto [cmd_list, res] = stream->new_command_list_unique();
+  TI_ASSERT(res == RhiResult::success);
   cmd_list->image_to_buffer(dst_buf, src_img, img_layout, params);
   stream->submit_synced(cmd_list.get());
 }

taichi/rhi/device.h

@@ -520,7 +520,15 @@ class TI_DLL_EXPORT Stream {
   virtual ~Stream() {
   }
 
-  virtual std::unique_ptr<CommandList> new_command_list() = 0;
+  virtual RhiResult new_command_list(CommandList **out_cmdlist) noexcept = 0;
+
+  inline std::pair<std::unique_ptr<CommandList>, RhiResult>
+  new_command_list_unique() {
+    CommandList *cmdlist{nullptr};
+    RhiResult res = this->new_command_list(&cmdlist);
+    return std::make_pair(std::unique_ptr<CommandList>(cmdlist), res);
+  }
+
   virtual StreamSemaphore submit(
       CommandList *cmdlist,
       const std::vector<StreamSemaphore> &wait_semaphores = {}) = 0;

taichi/rhi/dx/dx_device.cpp

@@ -929,8 +929,9 @@ Dx11Stream::Dx11Stream(Dx11Device *device_) : device_(device_) {
 Dx11Stream::~Dx11Stream() {
 }
 
-std::unique_ptr<CommandList> Dx11Stream::new_command_list() {
-  return std::make_unique<Dx11CommandList>(device_);
+RhiResult Dx11Stream::new_command_list(CommandList **out_cmdlist) noexcept {
+  *out_cmdlist = new Dx11CommandList(device_);
+  return RhiResult::success;
 }
 
 StreamSemaphore Dx11Stream::submit(

taichi/rhi/dx/dx_device.h

@@ -94,7 +94,7 @@ class Dx11Stream : public Stream {
   Dx11Stream(Dx11Device *);
   ~Dx11Stream() override;
 
-  std::unique_ptr<CommandList> new_command_list() override;
+  RhiResult new_command_list(CommandList **out_cmdlist) noexcept final;
   StreamSemaphore submit(
       CommandList *cmdlist,
       const std::vector<StreamSemaphore> &wait_semaphores = {}) override;

taichi/rhi/interop/vulkan_cpu_interop.cpp

@@ -51,7 +51,8 @@ void memcpy_cpu_to_vulkan_via_staging(DevicePtr dst,
   vk_dev->unmap(staging);
 
   auto stream = vk_dev->get_compute_stream();
-  auto cmd_list = stream->new_command_list();
+  auto [cmd_list, res] = stream->new_command_list_unique();
+  TI_ASSERT(res == RhiResult::success);
   cmd_list->buffer_copy(dst, staging, size);
   stream->submit_synced(cmd_list.get());
 }

taichi/rhi/metal/device.cpp

@@ -229,11 +229,12 @@ class StreamImpl : public Stream {
       : command_queue_(command_queue), alloc_buf_mapper_(alloc_buf_mapper) {
   }
 
-  std::unique_ptr<CommandList> new_command_list() override {
+  RhiResult new_command_list(CommandList **out_cmdlist) noexcept final {
     auto cb = new_command_buffer(command_queue_);
     TI_ASSERT(cb != nullptr);
     set_label(cb.get(), fmt::format("command_buffer_{}", list_counter_++));
-    return std::make_unique<CommandListImpl>(std::move(cb), alloc_buf_mapper_);
+    *out_cmdlist = new CommandListImpl(std::move(cb), alloc_buf_mapper_);
+    return RhiResult::success;
   }
 
   StreamSemaphore submit(
@@ -361,7 +362,8 @@ class DeviceImpl : public Device, public AllocToMTLBufferMapper {
 
   void memcpy_internal(DevicePtr dst, DevicePtr src, uint64_t size) override {
     Stream *stream = get_compute_stream();
-    std::unique_ptr<CommandList> cmd = stream->new_command_list();
+    auto [cmd, res] = stream->new_command_list_unique();
+    assert(res == RhiResult::success);
     cmd->buffer_copy(dst, src, size);
     stream->submit_synced(cmd.get());
   }

taichi/rhi/opengl/opengl_device.cpp

@@ -464,8 +464,9 @@ void GLCommandList::run_commands() {
 GLStream::~GLStream() {
 }
 
-std::unique_ptr<CommandList> GLStream::new_command_list() {
-  return std::make_unique<GLCommandList>(device_);
+RhiResult GLStream::new_command_list(CommandList **out_cmdlist) noexcept {
+  *out_cmdlist = new GLCommandList(device_);
+  return RhiResult::success;
 }
 
 StreamSemaphore GLStream::submit(

taichi/rhi/opengl/opengl_device.h

@@ -204,7 +204,7 @@ class GLStream : public Stream {
   }
   ~GLStream() override;
 
-  std::unique_ptr<CommandList> new_command_list() override;
+  RhiResult new_command_list(CommandList **out_cmdlist) noexcept final;
   StreamSemaphore submit(
       CommandList *cmdlist,
       const std::vector<StreamSemaphore> &wait_semaphores = {}) override;

taichi/rhi/vulkan/vulkan_api.cpp

@@ -219,7 +219,9 @@ IVkCommandBuffer allocate_command_buffer(IVkCommandPool pool,
     info.commandBufferCount = 1;
 
     VkResult res = vkAllocateCommandBuffers(pool->device, &info, &cmdbuf);
-    BAIL_ON_VK_BAD_RESULT_NO_RETURN(res, "failed to allocate command buffer");
+    if (res != VK_SUCCESS) {
+      return nullptr;
+    }
   }
 
   IVkCommandBuffer obj = std::make_shared<DeviceObjVkCommandBuffer>();

taichi/rhi/vulkan/vulkan_device.cpp

@@ -935,12 +935,12 @@ RhiResult VulkanCommandList::bind_raster_resources(
     return RhiResult::invalid_usage;
   }
 
-  if (res->index_type >= VK_INDEX_TYPE_MAX_ENUM) {
-    return RhiResult::not_supported;
-  }
-
   if (res->index_binding.buffer != nullptr) {
     // We have a valid index buffer
+    if (res->index_type >= VK_INDEX_TYPE_MAX_ENUM) {
+      return RhiResult::not_supported;
+    }
+
     vkapi::IVkBuffer index_buffer = res->index_binding.buffer;
     vkCmdBindIndexBuffer(buffer_->buffer, index_buffer->buffer,
                          res->index_binding.offset, res->index_type);
@@ -1809,7 +1809,8 @@ void VulkanDevice::memcpy_internal(DevicePtr dst,
                                    uint64_t size) {
   // TODO: always create a queue specifically for transfer
   Stream *stream = get_compute_stream();
-  std::unique_ptr<CommandList> cmd = stream->new_command_list();
+  auto [cmd, res] = stream->new_command_list_unique();
+  TI_ASSERT(res == RhiResult::success);
   cmd->buffer_copy(dst, src, size);
   stream->submit_synced(cmd.get());
 }
@@ -1847,11 +1848,16 @@ void VulkanDevice::wait_idle() {
   }
 }
 
-std::unique_ptr<CommandList> VulkanStream::new_command_list() {
+RhiResult VulkanStream::new_command_list(CommandList **out_cmdlist) noexcept {
   vkapi::IVkCommandBuffer buffer =
       vkapi::allocate_command_buffer(command_pool_);
 
-  return std::make_unique<VulkanCommandList>(&device_, this, buffer);
+  if (buffer == nullptr) {
+    return RhiResult::out_of_memory;
+  }
+
+  *out_cmdlist = new VulkanCommandList(&device_, this, buffer);
+  return RhiResult::success;
 }
 
 StreamSemaphore VulkanStream::submit(
@@ -2788,13 +2794,12 @@ DeviceAllocation VulkanSurface::get_depth_data(DeviceAllocation &depth_alloc) {
     depth_buffer_ = device_->allocate_memory(params);
   }
 
-  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();
+  auto [cmd_list, res] = stream->new_command_list_unique();
+  assert(res == RhiResult::success && "Failed to allocate 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,
@@ -2832,8 +2837,6 @@ DeviceAllocation VulkanSurface::get_image_data() {
     screenshot_buffer_ = device_->allocate_memory(params);
   }
 
-  std::unique_ptr<CommandList> cmd_list{nullptr};
-
   /*
   if (config_.window_handle) {
     // TODO: check if blit is supported, and use copy_image if not
@@ -2851,7 +2854,8 @@ DeviceAllocation VulkanSurface::get_image_data() {
   copy_params.image_extent.x = w;
   copy_params.image_extent.y = h;
   copy_params.image_aspect_flag = VK_IMAGE_ASPECT_COLOR_BIT;
-  cmd_list = stream->new_command_list();
+  auto [cmd_list, res] = stream->new_command_list_unique();
+  assert(res == RhiResult::success && "Failed to allocate command list");
   cmd_list->image_transition(img_alloc, ImageLayout::present_src,
                              ImageLayout::transfer_src);
   // TODO: directly map the image to cpu memory

taichi/rhi/vulkan/vulkan_device.h

@@ -548,7 +548,7 @@ class VulkanStream : public Stream {
                uint32_t queue_family_index);
   ~VulkanStream() override;
 
-  std::unique_ptr<CommandList> new_command_list() override;
+  RhiResult new_command_list(CommandList **out_cmdlist) noexcept final;
   StreamSemaphore submit(
       CommandList *cmdlist,
       const std::vector<StreamSemaphore> &wait_semaphores = {}) override;

taichi/rhi/vulkan/vulkan_device_creator.cpp

@@ -69,7 +69,7 @@ vk_debug_callback(VkDebugUtilsMessageSeverityFlagBitsEXT message_severity,
   }
 
   if (message_severity > VK_DEBUG_UTILS_MESSAGE_SEVERITY_INFO_BIT_EXT) {
-    char msg_buf[512];
+    char msg_buf[4096];
     snprintf(msg_buf, sizeof(msg_buf), "Vulkan validation layer: %d, %s",
              message_type, p_callback_data->pMessage);
 

taichi/runtime/gfx/runtime.cpp

@@ -597,7 +597,9 @@ StreamSemaphore GfxRuntime::flush() {
     sema = device_->get_compute_stream()->submit(current_cmdlist_.get());
     current_cmdlist_ = nullptr;
   } else {
-    auto cmdlist = device_->get_compute_stream()->new_command_list();
+    auto [cmdlist, res] =
+        device_->get_compute_stream()->new_command_list_unique();
+    TI_ASSERT(res == RhiResult::success);
     cmdlist->memory_barrier();
     sema = device_->get_compute_stream()->submit(cmdlist.get());
   }
@@ -613,7 +615,10 @@ void GfxRuntime::ensure_current_cmdlist() {
   if (!current_cmdlist_) {
     ctx_buffers_.clear();
     current_cmdlist_pending_since_ = high_res_clock::now();
-    current_cmdlist_ = device_->get_compute_stream()->new_command_list();
+    auto [cmdlist, res] =
+        device_->get_compute_stream()->new_command_list_unique();
+    TI_ASSERT(res == RhiResult::success);
+    current_cmdlist_ = std::move(cmdlist);
   }
 }
 void GfxRuntime::submit_current_cmdlist_if_timeout() {
@@ -643,7 +648,9 @@ void GfxRuntime::init_nonroot_buffers() {
 
   // Need to zero fill the buffers, otherwise there could be NaN.
   Stream *stream = device_->get_compute_stream();
-  auto cmdlist = stream->new_command_list();
+  auto [cmdlist, res] =
+      device_->get_compute_stream()->new_command_list_unique();
+  TI_ASSERT(res == RhiResult::success);
 
   cmdlist->buffer_fill(global_tmps_buffer_->get_ptr(0), kBufferSizeEntireSize,
                        /*data=*/0);
@@ -663,7 +670,9 @@ void GfxRuntime::add_root_buffer(size_t root_buffer_size) {
            /*host_write=*/false, /*host_read=*/false,
            /*export_sharing=*/false, AllocUsage::Storage});
   Stream *stream = device_->get_compute_stream();
-  auto cmdlist = stream->new_command_list();
+  auto [cmdlist, res] =
+      device_->get_compute_stream()->new_command_list_unique();
+  TI_ASSERT(res == RhiResult::success);
   cmdlist->buffer_fill(new_buffer->get_ptr(0), kBufferSizeEntireSize,
                        /*data=*/0);
   stream->submit_synced(cmdlist.get());

taichi/ui/backends/vulkan/gui.cpp

@@ -59,7 +59,8 @@ void Gui::init_render_resources(VkRenderPass render_pass) {
   // Upload Fonts
   {
     auto stream = device.get_graphics_stream();
-    std::unique_ptr<CommandList> cmd_list = stream->new_command_list();
+    auto [cmd_list, res] = stream->new_command_list_unique();
+    assert(res == RhiResult::success && "Failed to allocate command list");
     VkCommandBuffer command_buffer =
         static_cast<VulkanCommandList *>(cmd_list.get())
             ->vk_command_buffer()

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

@@ -127,7 +127,8 @@ void SetImage::update_data(const SetImageInfo &info) {
     prog->enqueue_compute_op_lambda(copy_op, {});
   } else {
     auto stream = app_context_->device().get_graphics_stream();
-    auto cmd_list = stream->new_command_list();
+    auto [cmd_list, res] = stream->new_command_list_unique();
+    assert(res == RhiResult::success && "Failed to allocate command list");
     copy_op(&app_context_->device(), cmd_list.get());
     if (sema) {
       stream->submit(cmd_list.get(), {sema});
@@ -184,7 +185,8 @@ void SetImage::update_data(Texture *tex) {
                                     ImageLayout::transfer_src}});
   } else {
     auto stream = app_context_->device().get_graphics_stream();
-    auto cmd_list = stream->new_command_list();
+    auto [cmd_list, res] = stream->new_command_list_unique();
+    assert(res == RhiResult::success && "Failed to allocate command list");
     copy_op(&app_context_->device(), cmd_list.get());
     stream->submit(cmd_list.get());
   }