[vulkan] Device API explicit semaphores

python/taichi/ui/staging_buffer.py

@@ -89,21 +89,21 @@ def copy_image_f32_to_u8(src: ti.template(), dst: ti.template(),
             c = src[i, j][k]
             c = max(0.0, min(1.0, c))
             c = c * 255
-            dst[i, j][k] = int(c)
+            dst[i, j][k] = ti.cast(c, ti.u8)
         if num_components < 4:
             # alpha channel
-            dst[i, j][3] = 255
+            dst[i, j][3] = ti.cast(255, ti.u8)
 
 
 @ti.kernel
 def copy_image_u8_to_u8(src: ti.template(), dst: ti.template(),
                         num_components: ti.template()):
     for i, j in src:
         for k in ti.static(range(num_components)):
-            dst[i, j][k] = src[i, j][k]
+            dst[i, j][k] = ti.cast(src[i, j][k], ti.u8)
         if num_components < 4:
             # alpha channel
-            dst[i, j][3] = 255
+            dst[i, j][3] = ti.cast(255, ti.u8)
 
 
 # ggui renderer always assumes the input image to be u8 RGBA

taichi/aot/module_loader.h

@@ -166,6 +166,9 @@ class TargetDevice : public Device {
   Stream *get_compute_stream() override {
     TI_NOT_IMPLEMENTED;
   }
+  void wait_idle() override {
+    TI_NOT_IMPLEMENTED;
+  }
 };
 
 }  // namespace aot

taichi/backends/cpu/cpu_device.h

@@ -69,8 +69,13 @@ class CpuStream : public Stream {
   ~CpuStream() override{};
 
   std::unique_ptr<CommandList> new_command_list() override{TI_NOT_IMPLEMENTED};
-  void submit(CommandList *cmdlist) override{TI_NOT_IMPLEMENTED};
-  void submit_synced(CommandList *cmdlist) override{TI_NOT_IMPLEMENTED};
+  StreamSemaphore submit(CommandList *cmdlist,
+                         const std::vector<StreamSemaphore> &wait_semaphores =
+                             {}) override{TI_NOT_IMPLEMENTED};
+  StreamSemaphore submit_synced(
+      CommandList *cmdlist,
+      const std::vector<StreamSemaphore> &wait_semaphores = {}) override{
+      TI_NOT_IMPLEMENTED};
 
   void command_sync() override{TI_NOT_IMPLEMENTED};
 };
@@ -111,6 +116,8 @@ class CpuDevice : public LlvmDevice {
 
   Stream *get_compute_stream() override{TI_NOT_IMPLEMENTED};
 
+  void wait_idle() override{TI_NOT_IMPLEMENTED};
+
  private:
   std::vector<AllocInfo> allocations_;
   std::unordered_map<int, std::unique_ptr<VirtualMemoryAllocator>>

taichi/backends/cuda/cuda_device.h

@@ -69,8 +69,13 @@ class CudaStream : public Stream {
   ~CudaStream() override{};
 
   std::unique_ptr<CommandList> new_command_list() override{TI_NOT_IMPLEMENTED};
-  void submit(CommandList *cmdlist) override{TI_NOT_IMPLEMENTED};
-  void submit_synced(CommandList *cmdlist) override{TI_NOT_IMPLEMENTED};
+  StreamSemaphore submit(CommandList *cmdlist,
+                         const std::vector<StreamSemaphore> &wait_semaphores =
+                             {}) override{TI_NOT_IMPLEMENTED};
+  StreamSemaphore submit_synced(
+      CommandList *cmdlist,
+      const std::vector<StreamSemaphore> &wait_semaphores = {}) override{
+      TI_NOT_IMPLEMENTED};
 
   void command_sync() override{TI_NOT_IMPLEMENTED};
 };
@@ -123,6 +128,8 @@ class CudaDevice : public LlvmDevice {
 
   Stream *get_compute_stream() override{TI_NOT_IMPLEMENTED};
 
+  void wait_idle() override{TI_NOT_IMPLEMENTED};
+
  private:
   std::vector<AllocInfo> allocations_;
   void validate_device_alloc(const DeviceAllocation alloc) {

taichi/backends/device.h

@@ -397,13 +397,26 @@ inline bool operator&(AllocUsage a, AllocUsage b) {
   return static_cast<int>(a) & static_cast<int>(b);
 }
 
+class StreamSemaphoreObject {
+ public:
+  virtual ~StreamSemaphoreObject() {
+  }
+};
+
+using StreamSemaphore = std::shared_ptr<StreamSemaphoreObject>;
+
 class Stream {
  public:
-  virtual ~Stream(){};
+  virtual ~Stream() {
+  }
 
   virtual std::unique_ptr<CommandList> new_command_list() = 0;
-  virtual void submit(CommandList *cmdlist) = 0;
-  virtual void submit_synced(CommandList *cmdlist) = 0;
+  virtual StreamSemaphore submit(
+      CommandList *cmdlist,
+      const std::vector<StreamSemaphore> &wait_semaphores = {}) = 0;
+  virtual StreamSemaphore submit_synced(
+      CommandList *cmdlist,
+      const std::vector<StreamSemaphore> &wait_semaphores = {}) = 0;
 
   virtual void command_sync() = 0;
 };
@@ -457,6 +470,9 @@ class Device {
   // Each thraed will acquire its own stream
   virtual Stream *get_compute_stream() = 0;
 
+  // Wait for all tasks to complete (task from all streams)
+  virtual void wait_idle() = 0;
+
   // Mapping can fail and will return nullptr
   virtual void *map_range(DevicePtr ptr, uint64_t size) = 0;
   virtual void *map(DeviceAllocation alloc) = 0;
@@ -498,8 +514,10 @@ class Surface {
   virtual ~Surface() {
   }
 
+  virtual StreamSemaphore acquire_next_image() = 0;
   virtual DeviceAllocation get_target_image() = 0;
-  virtual void present_image() = 0;
+  virtual void present_image(
+      const std::vector<StreamSemaphore> &wait_semaphores = {}) = 0;
   virtual std::pair<uint32_t, uint32_t> get_size() = 0;
   virtual int get_image_count() = 0;
   virtual BufferFormat image_format() = 0;