From 3dce9b7c5de177aaba39d05b4effa01871c233a4 Mon Sep 17 00:00:00 2001
From: "Zhekun(Josh) Zhang" <32320144+zhekunz2@users.noreply.github.com>
Date: Fri, 29 Mar 2024 01:22:37 -0700
Subject: [PATCH] [Runtime] Support host stream (#166)

Add host stream support using host callbacks `cudaLaunchHostFunc` rather
than immediate execution.
---
 .../brt/backends/cpu/device/cpu_work_queue.h  |  18 ++-
 .../backends/cuda/device/cuda_work_queue.h    |  60 +++++-----
 runtime/include/brt/core/context/work_queue.h |  20 ++--
 .../brt/core/framework/op_kernel_impl_base.h  |  33 +++++-
 .../brt/core/framework/op_kernel_info.h       |  26 ++---
 .../lib/backends/cpu/device/cpu_work_queue.cc |  37 +++---
 .../providers/default/custom_call/tf_equal.cc |   2 +-
 .../default/custom_call/tf_select.cc          |   2 +-
 .../custom_call/tf_string_to_number.cc        |  17 ++-
 .../providers/default/custom_call/tf_where.cc |   9 +-
 .../cpu/providers/default/custom_call/topk.cc |  11 +-
 .../cpu/providers/default/llvm/jit.cc         |   2 +-
 .../providers/default/math/elementwise_ops.cc |   2 +-
 .../providers/default/tensor_generate/fill.cc |   2 +-
 .../default/tensor_generate/rng_state.cc      |   4 +-
 .../backends/cuda/device/cuda_work_queue.cc   | 106 +++++++++---------
 .../cuda/providers/default/codegen/ptx.cc     |   3 +-
 .../cuda/providers/default/copy/copy.cc       |   4 +-
 .../providers/default/math/elementwise_ops.cc |   4 +-
 runtime/lib/core/framework/execution_plan.cc  |  17 ++-
 .../cuda/device/cuda_work_queue_test.cc       | 104 ++++++++++++++---
 .../test/backends/cuda/device/nvrtc_test.cc   |   4 +-
 runtime/test/backends/cuda/device/ptx_test.cc |  12 +-
 .../default/kernel/multi_stream_test.cc       |   6 +-
 runtime/test/external_kernels/cuda/kernels.cc |   2 +-
 25 files changed, 313 insertions(+), 194 deletions(-)

diff --git a/runtime/include/brt/backends/cpu/device/cpu_work_queue.h b/runtime/include/brt/backends/cpu/device/cpu_work_queue.h
index 633ad1d78..2260faac9 100644
--- a/runtime/include/brt/backends/cpu/device/cpu_work_queue.h
+++ b/runtime/include/brt/backends/cpu/device/cpu_work_queue.h
@@ -30,14 +30,13 @@ class CPUNaiveWorkQueue : public WorkQueue {
   explicit CPUNaiveWorkQueue(const std::string &name = "cpu_naive");
 
   common::Status AddTask(int /*task_type*/, const void * /*func*/,
-                         void ** /*args*/) override;
-
-  common::Status AddEventWait(mlir::Operation *,
-                              std::vector<mlir::Operation *>) override;
+                         void ** /*args*/, int /*op_id*/,
+                         const std::vector<int> & /*dependency*/) override;
 
   common::Status Sync() override;
 
-  common::Status AddHostTask(std::function<void(void)> &&task) override;
+  common::Status AddHostTask(const void *task, void **args, int op_id,
+                             const std::vector<int> &dependency) override;
 };
 
 // WorkQueue which runs host task lazily
@@ -46,14 +45,13 @@ class CPULazyWorkQueue : public WorkQueue {
   explicit CPULazyWorkQueue(const std::string &name = "cpu_lazy");
 
   common::Status AddTask(int /*task_type*/, const void * /*func*/,
-                         void ** /*args*/) override;
-
-  common::Status AddEventWait(mlir::Operation *,
-                              std::vector<mlir::Operation *>) override;
+                         void ** /*args*/, int /*op_id*/,
+                         const std::vector<int> & /*dependency*/) override;
 
   common::Status Sync() override;
 
-  common::Status AddHostTask(std::function<void(void)> &&task) override;
+  common::Status AddHostTask(const void *task, void **args, int op_id,
+                             const std::vector<int> &dependency) override;
 
 private:
   std::vector<std::function<void(void)>> tasks;
diff --git a/runtime/include/brt/backends/cuda/device/cuda_work_queue.h b/runtime/include/brt/backends/cuda/device/cuda_work_queue.h
index 93ba8a190..68627395d 100644
--- a/runtime/include/brt/backends/cuda/device/cuda_work_queue.h
+++ b/runtime/include/brt/backends/cuda/device/cuda_work_queue.h
@@ -56,18 +56,18 @@ class CUDAWorkQueue : public WorkQueue {
 
   // Enqueue a func call, thread-safe.
   // func is a stateless function
-  virtual common::Status AddTask(int task_type, const void *func,
-                                 void **args) override;
-
-  virtual common::Status AddEventWait(mlir::Operation *,
-                                      std::vector<mlir::Operation *>) override;
+  virtual common::Status AddTask(int task_type, const void *func, void **args,
+                                 int op_id,
+                                 const std::vector<int> &dependency) override;
   // Barrier
   virtual common::Status Sync() override;
 
   virtual CUstream_st *GetComputeStream() { return nullptr; }
 
-  common::Status AddHostTask(std::function<void(void)> &&task) override {
-    task();
+  common::Status AddHostTask(const void *task, void **args, int op_id,
+                             const std::vector<int> &dependency) override {
+    auto func = reinterpret_cast<const std::function<void(void)> *>(task);
+    (*func)();
     return common::Status::OK();
   }
 
@@ -93,10 +93,10 @@ class CUDASingleStreamWorkQueue final : public CUDAWorkQueue {
 
   // Enqueue a func call, thread-safe.
   // func is a stateless function
-  common::Status AddTask(int task_type, const void *func, void **args) override;
+  common::Status AddTask(int task_type, const void *func, void **args,
+                         int op_id,
+                         const std::vector<int> &dependency) override;
 
-  common::Status AddEventWait(mlir::Operation *,
-                              std::vector<mlir::Operation *>) override;
   // Barrier
   common::Status Sync() override;
 
@@ -110,42 +110,47 @@ class CUDASingleStreamWorkQueue final : public CUDAWorkQueue {
 };
 
 /**
- * CUDAOneComputeTwoTransferWorkQueue is a derived class of WorkQueue
+ * CUDAMultiStreamWorkQueue is a derived class of WorkQueue
  * that uses mutliple CUDA Stream's as a WorkQueue.
  *
  * A typical usage is using 3 streams: one for compute, two for bidirectional
- * data transfer.
+ * data transfer, one for host.
  */
-class CUDAOneComputeTwoTransferWorkQueue final : public CUDAWorkQueue {
+class CUDAMultiStreamWorkQueue final : public CUDAWorkQueue {
 public:
-  CUDAOneComputeTwoTransferWorkQueue(int device_id);
+  CUDAMultiStreamWorkQueue(int device_id);
 
   // Undefined what happens to pending work when destructor is called.
-  virtual ~CUDAOneComputeTwoTransferWorkQueue();
+  virtual ~CUDAMultiStreamWorkQueue();
 
   // Enqueue a func call, thread-safe.
   // func is a stateless function
-  common::Status AddTask(int task_type, const void *func, void **args) override;
+  common::Status AddTask(int task_type, const void *func, void **args,
+                         int op_id,
+                         const std::vector<int> &dependency) override;
+
+  common::Status AddHostTask(const void *task, void **args, int op_id,
+                             const std::vector<int> &dependency) override;
 
   // Barrier
   common::Status Sync() override;
 
-  size_t GetStreamIdx(mlir::Operation *op);
-
-  common::Status AddEventWait(mlir::Operation *,
-                              std::vector<mlir::Operation *>) override;
+  common::Status AddEventWait(size_t, std::vector<int>);
 
   CUstream_st *GetComputeStream() override { return streams_[0]; }
+  CUstream_st *GetH2DStream() { return streams_[1]; }
+  CUstream_st *GetD2HStream() { return streams_[2]; }
+  CUstream_st *GetHostStream() { return streams_[3]; }
 
 private:
-  CUstream_st *streams_[3]; // 0 for compute, 1 for h2d, 2 for d2h
+  // 0 for compute, 1 for h2d, 2 for d2h, 3 for host
+  CUstream_st *streams_[4];
 
   std::vector<CUevent_st *> events_;
 
-  CUDAOneComputeTwoTransferWorkQueue(
-      const CUDAOneComputeTwoTransferWorkQueue &) = delete;
-  CUDAOneComputeTwoTransferWorkQueue &
-  operator=(const CUDAOneComputeTwoTransferWorkQueue &) = delete;
+  CUDAMultiStreamWorkQueue(const CUDAMultiStreamWorkQueue &) = delete;
+  CUDAMultiStreamWorkQueue &
+  operator=(const CUDAMultiStreamWorkQueue &) = delete;
 };
 
 /**
@@ -165,7 +170,9 @@ class CUDAExternalStreamWorkQueue final : public CUDAWorkQueue {
 
   // Enqueue a func call, thread-safe.
   // func is a stateless function
-  common::Status AddTask(int task_type, const void *func, void **args) override;
+  common::Status AddTask(int task_type, const void *func, void **args,
+                         int op_id,
+                         const std::vector<int> &dependency) override;
 
   // Barrier
   common::Status Sync() override;
@@ -178,5 +185,4 @@ class CUDAExternalStreamWorkQueue final : public CUDAWorkQueue {
   CUDAExternalStreamWorkQueue &
   operator=(const CUDAExternalStreamWorkQueue &) = delete;
 };
-
 } // namespace brt
diff --git a/runtime/include/brt/core/context/work_queue.h b/runtime/include/brt/core/context/work_queue.h
index 9bb97f9c2..57bc8aef7 100644
--- a/runtime/include/brt/core/context/work_queue.h
+++ b/runtime/include/brt/core/context/work_queue.h
@@ -61,15 +61,13 @@ class WorkQueue {
 
   // Enqueue a func call, thread-safe.
   // func is a stateless function
-  virtual common::Status AddTask(int task_type, const void *func,
-                                 void **args) = 0;
+  virtual common::Status AddTask(int task_type, const void *func, void **args,
+                                 int op_id,
+                                 const std::vector<int> &dependency) = 0;
 
   // Enqueue a task on host side
-  virtual common::Status AddHostTask(std::function<void(void)> &&task) = 0;
-
-  // Add wait event if any
-  virtual common::Status AddEventWait(mlir::Operation *,
-                                      std::vector<mlir::Operation *>) = 0;
+  virtual common::Status AddHostTask(const void *task, void **args, int op_id,
+                                     const std::vector<int> &dependency) = 0;
 
   // Enqueue through a functor
   // Note, the functor is called immediately.
@@ -81,6 +79,9 @@ class WorkQueue {
   // Barrier
   virtual common::Status Sync() = 0;
 
+protected:
+  std::unordered_map<int, int> id_to_stream_map_;
+
 private:
   const std::string name_;
   WorkQueue(const WorkQueue &) = delete;
@@ -89,9 +90,10 @@ class WorkQueue {
 
 } // namespace brt
 
-#define DispatchHostTask(wq, stmt)                                             \
+#define DispatchHostTask(wq, op_id, dependency, stmt)                          \
   if (wq) {                                                                    \
-    wq->AddHostTask([=] { stmt });                                             \
+    std::function<void(void)> func = [=]() { stmt };                           \
+    wq->AddHostTask(&func, nullptr, op_id, dependency);                        \
   } else {                                                                     \
     do {                                                                       \
       stmt                                                                     \
diff --git a/runtime/include/brt/core/framework/op_kernel_impl_base.h b/runtime/include/brt/core/framework/op_kernel_impl_base.h
index 9d21dc20a..208522450 100644
--- a/runtime/include/brt/core/framework/op_kernel_impl_base.h
+++ b/runtime/include/brt/core/framework/op_kernel_impl_base.h
@@ -96,6 +96,21 @@ struct WorkQueue {
   }
 };
 
+struct OpId {
+  template <typename Impl>
+  static inline int Get(Impl *impl, const ExecutionContext &ctx) {
+    return impl->GetOpId();
+  }
+};
+
+struct Dependency {
+  template <typename Impl>
+  static inline const std::vector<int> &Get(Impl *impl,
+                                            const ExecutionContext &ctx) {
+    return impl->GetDependency();
+  }
+};
+
 // op kernel which has temporary workspace as argument, should implement
 // GetWorkspaceSize() interface
 struct Workspace : public PerFrameHookTrait {
@@ -196,6 +211,12 @@ struct NaiveOpKernelIfaceTraits : public OpKernelIfaceTraitsBase<Arguments...> {
       return OpAccessor(info_, ctx.exec_frame);
     }
 
+    int GetOpId() const { return info_.GetOpId(); }
+
+    const std::vector<int> &GetDependency() const {
+      return info_.GetDependency();
+    }
+
   private:
     const OpKernelInfo &info_;
   };
@@ -226,8 +247,10 @@ struct OpKernelWithWorkspaceIfaceTraitsT
 
 template <template <typename...> class Base, typename... Arguments>
 struct HostOpKernelIfaceTraitsT
-    : public Base<argument_type::WorkQueue, Arguments...> {
-  using BaseTraits = Base<argument_type::WorkQueue, Arguments...>;
+    : public Base<argument_type::WorkQueue, argument_type::OpId,
+                  argument_type::Dependency, Arguments...> {
+  using BaseTraits = Base<argument_type::WorkQueue, argument_type::OpId,
+                          argument_type::Dependency, Arguments...>;
 
   template <typename T>
   using ImplMixinBase = typename BaseTraits::template ImplMixin<T>;
@@ -238,8 +261,10 @@ struct HostOpKernelIfaceTraitsT
     using ImplMixinBase<ImplBase>::ImplMixinBase;
 
     template <typename... Args>
-    common::Status Execute(WorkQueue *work_queue, Args &&...args) {
-      DispatchHostTask(work_queue, { ImplBase::Execute(args...); });
+    common::Status Execute(WorkQueue *work_queue, int op_id,
+                           const std::vector<int> &dependency, Args &&...args) {
+      DispatchHostTask(work_queue, op_id, dependency,
+                       { ImplBase::Execute(args...); });
       return common::Status::OK();
     }
   };
diff --git a/runtime/include/brt/core/framework/op_kernel_info.h b/runtime/include/brt/core/framework/op_kernel_info.h
index d6c2a14b5..4784f5d0e 100644
--- a/runtime/include/brt/core/framework/op_kernel_info.h
+++ b/runtime/include/brt/core/framework/op_kernel_info.h
@@ -54,22 +54,21 @@ class OpKernelInfo {
 public:
   OpKernelInfo(
       const ExecutionProvider &provider, const ir::IRHandle &handle,
-      mlir::Operation *op,
+      mlir::Operation *op, int op_id,
       const std::unordered_map<std::string, std::unique_ptr<IAllocator>> &alloc,
       IAllocator *last_allc,
       const std::unordered_map<void *, size_t> &tensor_to_idx,
       const std::unordered_map<void *, size_t> &scalar_to_idx,
       const std::vector<AsyncValue> &weights, size_t intermediate_begin,
-      const std::string &ir_path,
-      const std::vector<mlir::Operation *> &dependency)
-      : provider_(provider), handle_(handle), op_(op), allocators_(alloc),
-        last_allocator_(last_allc), tensor_to_idx_(tensor_to_idx),
-        scalar_to_idx_(scalar_to_idx), weights_(weights),
-        intermediate_begin_(intermediate_begin), ir_path_(ir_path),
-        dependency_(dependency) {}
+      const std::string &ir_path, const std::vector<int> &dependency)
+      : provider_(provider), handle_(handle), op_(op), op_id_(op_id),
+        allocators_(alloc), last_allocator_(last_allc),
+        tensor_to_idx_(tensor_to_idx), scalar_to_idx_(scalar_to_idx),
+        weights_(weights), intermediate_begin_(intermediate_begin),
+        ir_path_(ir_path), dependency_(dependency) {}
 
   OpKernelInfo(const OpKernelInfo &other)
-      : OpKernelInfo(other.provider_, other.handle_, other.op_,
+      : OpKernelInfo(other.provider_, other.handle_, other.op_, other.op_id_,
                      other.allocators_, other.last_allocator_,
                      other.tensor_to_idx_, other.scalar_to_idx_, other.weights_,
                      other.intermediate_begin_, other.ir_path_,
@@ -81,6 +80,8 @@ class OpKernelInfo {
 
   mlir::Operation *GetOperation() const { return op_; }
 
+  int GetOpId() const { return op_id_; }
+
   const std::unordered_map<void *, size_t> &GetTensorToIndex() const {
     return tensor_to_idx_;
   }
@@ -91,9 +92,7 @@ class OpKernelInfo {
 
   const std::vector<AsyncValue> &GetWeights() const { return weights_; }
 
-  const std::vector<mlir::Operation *> &GetDependency() const {
-    return dependency_;
-  }
+  const std::vector<int> &GetDependency() const { return dependency_; }
 
   // const BrtMemoryInfo& GetMemoryInfo(int device_id, BrtMemType mem_type)
   // const;
@@ -119,6 +118,7 @@ class OpKernelInfo {
   const brt::ir::IRHandle &handle_;
 
   mlir::Operation *op_;
+  int op_id_;
 
   const std::unordered_map<std::string, std::unique_ptr<IAllocator>>
       &allocators_;
@@ -134,7 +134,7 @@ class OpKernelInfo {
 
   const std::string &ir_path_;
 
-  const std::vector<mlir::Operation *> &dependency_;
+  std::vector<int> dependency_;
 
   OpKernelInfo(OpKernelInfo &&) = delete;
   OpKernelInfo &operator=(OpKernelInfo &&) = delete;
diff --git a/runtime/lib/backends/cpu/device/cpu_work_queue.cc b/runtime/lib/backends/cpu/device/cpu_work_queue.cc
index 8f1542e58..a40544cb0 100644
--- a/runtime/lib/backends/cpu/device/cpu_work_queue.cc
+++ b/runtime/lib/backends/cpu/device/cpu_work_queue.cc
@@ -23,31 +23,30 @@ namespace cpu {
 CPUNaiveWorkQueue::CPUNaiveWorkQueue(const std::string &name)
     : WorkQueue(name) {}
 
-common::Status CPUNaiveWorkQueue::AddTask(int /*task_type*/,
-                                          const void * /*func*/,
-                                          void ** /*args*/) {
+common::Status
+CPUNaiveWorkQueue::AddTask(int /*task_type*/, const void * /*func*/,
+                           void ** /*args*/, int op_id,
+                           const std::vector<int> & /*dependency*/) {
   return common::Status(common::StatusCategory::BRT, common::StatusCode::FAIL,
                         "Use AddHostTask for cpu work queue");
 }
 
-common::Status CPUNaiveWorkQueue::AddEventWait(mlir::Operation *,
-                                               std::vector<mlir::Operation *>) {
-  return common::Status::OK();
-}
-
 common::Status CPUNaiveWorkQueue::Sync() { return common::Status::OK(); }
 
 common::Status
-CPUNaiveWorkQueue::AddHostTask(std::function<void(void)> &&task) {
-  task();
+CPUNaiveWorkQueue::AddHostTask(const void *task, void **args, int op_id,
+                               const std::vector<int> &dependency) {
+  auto func = reinterpret_cast<const std::function<void(void)> *>(task);
+  (*func)();
   return common::Status::OK();
 }
 
 CPULazyWorkQueue::CPULazyWorkQueue(const std::string &name) : WorkQueue(name) {}
 
-common::Status CPULazyWorkQueue::AddTask(int /*task_type*/,
-                                         const void * /*func*/,
-                                         void ** /*args*/) {
+common::Status
+CPULazyWorkQueue::AddTask(int /*task_type*/, const void * /*func*/,
+                          void ** /*args*/, int op_id,
+                          const std::vector<int> & /*dependency*/) {
   return common::Status(common::StatusCategory::BRT, common::StatusCode::FAIL,
                         "Use AddHostTask for cpu work queue");
 }
@@ -59,13 +58,11 @@ common::Status CPULazyWorkQueue::Sync() {
   return common::Status::OK();
 }
 
-common::Status CPULazyWorkQueue::AddEventWait(mlir::Operation *,
-                                              std::vector<mlir::Operation *>) {
-  return common::Status::OK();
-}
-
-common::Status CPULazyWorkQueue::AddHostTask(std::function<void(void)> &&task) {
-  tasks.push_back(std::move(task));
+common::Status
+CPULazyWorkQueue::AddHostTask(const void *task, void **args, int op_id,
+                              const std::vector<int> &dependency) {
+  auto func = reinterpret_cast<const std::function<void(void)> *>(task);
+  tasks.push_back(std::move(*func));
   return common::Status::OK();
 }
 
diff --git a/runtime/lib/backends/cpu/providers/default/custom_call/tf_equal.cc b/runtime/lib/backends/cpu/providers/default/custom_call/tf_equal.cc
index eb97a0514..c6529ca0b 100644
--- a/runtime/lib/backends/cpu/providers/default/custom_call/tf_equal.cc
+++ b/runtime/lib/backends/cpu/providers/default/custom_call/tf_equal.cc
@@ -49,7 +49,7 @@ common::Status TFEqual::RunImpl(const ExecutionContext &ctx) {
   case DTypeEnum::StringView: {
     StringView *ss0 = reinterpret_cast<StringView *>(src0),
                *ss1 = reinterpret_cast<StringView *>(src1);
-    DispatchHostTask(ctx.work_queue, {
+    DispatchHostTask(ctx.work_queue, info_.GetOpId(), info_.GetDependency(), {
       if (require_broadcast) {
         StringView scalar = arg0_length == 1 ? ss0[0] : ss1[0];
         StringView *ss = arg0_length == 1 ? ss1 : ss0;
diff --git a/runtime/lib/backends/cpu/providers/default/custom_call/tf_select.cc b/runtime/lib/backends/cpu/providers/default/custom_call/tf_select.cc
index 06b95a46e..491167608 100644
--- a/runtime/lib/backends/cpu/providers/default/custom_call/tf_select.cc
+++ b/runtime/lib/backends/cpu/providers/default/custom_call/tf_select.cc
@@ -71,7 +71,7 @@ common::Status TFSelect::RunImpl(const ExecutionContext &ctx) {
                *ss1 = reinterpret_cast<StringView *>(src1),
                *ds = reinterpret_cast<StringView *>(dest);
     auto length = accessor.GetNumElementsOfShape(arg1_shape);
-    DispatchHostTask(ctx.work_queue, {
+    DispatchHostTask(ctx.work_queue, info_.GetOpId(), info_.GetDependency(), {
       (TFSelectImpl(cond_shape, arg1_shape, length, cond, ss0, ss1, ds));
     });
     return common::Status::OK();
diff --git a/runtime/lib/backends/cpu/providers/default/custom_call/tf_string_to_number.cc b/runtime/lib/backends/cpu/providers/default/custom_call/tf_string_to_number.cc
index 9adc7e9ee..379b56fc0 100644
--- a/runtime/lib/backends/cpu/providers/default/custom_call/tf_string_to_number.cc
+++ b/runtime/lib/backends/cpu/providers/default/custom_call/tf_string_to_number.cc
@@ -90,7 +90,8 @@ void convertNumberFromStringParallel(int64_t length, StringView *ss, T *ds) {
 
 template <typename T>
 common::Status TFStringToNumberImpl(const OpAccessor &accessor,
-                                    WorkQueue *work_queue) {
+                                    WorkQueue *work_queue, int op_id,
+                                    const std::vector<int> &dependency) {
   DTypeEnum dtype = accessor.GetArgDTypeEnum(0);
   const auto src_shape = accessor.GetArgShape(0);
   const auto dest_shape = accessor.GetArgShape(1);
@@ -103,7 +104,7 @@ common::Status TFStringToNumberImpl(const OpAccessor &accessor,
     StringView *ss = reinterpret_cast<StringView *>(src);
     T *ds = reinterpret_cast<T *>(dest);
     auto length = accessor.GetNumElementsOfShape(src_shape);
-    DispatchHostTask(work_queue,
+    DispatchHostTask(work_queue, op_id, dependency,
                      { convertNumberFromStringParallel<T>(length, ss, ds); });
     return common::Status::OK();
   }
@@ -119,13 +120,17 @@ common::Status TFStringToNumber::RunImpl(const ExecutionContext &ctx) {
   std::string out_type = accessor.GetAttrAsString("out_type");
 
   if (out_type == "i32")
-    return TFStringToNumberImpl<int32_t>(accessor, ctx.work_queue);
+    return TFStringToNumberImpl<int32_t>(
+        accessor, ctx.work_queue, info_.GetOpId(), info_.GetDependency());
   else if (out_type == "i64")
-    return TFStringToNumberImpl<int64_t>(accessor, ctx.work_queue);
+    return TFStringToNumberImpl<int64_t>(
+        accessor, ctx.work_queue, info_.GetOpId(), info_.GetDependency());
   else if (out_type == "f32")
-    return TFStringToNumberImpl<float>(accessor, ctx.work_queue);
+    return TFStringToNumberImpl<float>(accessor, ctx.work_queue,
+                                       info_.GetOpId(), info_.GetDependency());
   else if (out_type == "f64")
-    return TFStringToNumberImpl<double>(accessor, ctx.work_queue);
+    return TFStringToNumberImpl<double>(accessor, ctx.work_queue,
+                                        info_.GetOpId(), info_.GetDependency());
   else
     return common::Status(common::StatusCategory::BRT,
                           common::StatusCode::NOT_IMPLEMENTED,
diff --git a/runtime/lib/backends/cpu/providers/default/custom_call/tf_where.cc b/runtime/lib/backends/cpu/providers/default/custom_call/tf_where.cc
index eccb56793..c6bc2204f 100644
--- a/runtime/lib/backends/cpu/providers/default/custom_call/tf_where.cc
+++ b/runtime/lib/backends/cpu/providers/default/custom_call/tf_where.cc
@@ -27,7 +27,8 @@ namespace brt {
 namespace cpu {
 
 template <typename T>
-void TFWhereImpl(const OpAccessor &accessor, WorkQueue *work_queue) {
+void TFWhereImpl(const OpAccessor &accessor, WorkQueue *work_queue, int op_id,
+                 const std::vector<int> &dependency) {
   const auto &shape = accessor.GetArgShape(0);
   const int64_t num_elements = accessor.GetNumElementsOfShape(shape);
   const int32_t rank = shape.size();
@@ -35,7 +36,7 @@ void TFWhereImpl(const OpAccessor &accessor, WorkQueue *work_queue) {
   T *data = static_cast<T *>(accessor.GetArgAsyncValueRef(0));
   int64_t *result = static_cast<int64_t *>(accessor.GetArgAsyncValueRef(1));
 
-  DispatchHostTask(work_queue, {
+  DispatchHostTask(work_queue, op_id, dependency, {
     for (int64_t i = 0, result_pos = 0; i < num_elements; ++i) {
       if (!data[i])
         continue;
@@ -60,8 +61,8 @@ common::Status TFWhere::RunImpl(const ExecutionContext &ctx) {
   auto data_dtype = accessor.GetArgDTypeEnum(0);
 #define HANDLE_DTYPE(DType)                                                    \
   if (data_dtype == DType) {                                                   \
-    TFWhereImpl<typename DTypeTraits<DType>::type_t>(accessor,                 \
-                                                     ctx.work_queue);          \
+    TFWhereImpl<typename DTypeTraits<DType>::type_t>(                          \
+        accessor, ctx.work_queue, info_.GetOpId(), info_.GetDependency());     \
     return common::Status::OK();                                               \
   }
   HANDLE_DTYPE(DTypeEnum::Float32)
diff --git a/runtime/lib/backends/cpu/providers/default/custom_call/topk.cc b/runtime/lib/backends/cpu/providers/default/custom_call/topk.cc
index c7fcdcd35..fed9d341b 100644
--- a/runtime/lib/backends/cpu/providers/default/custom_call/topk.cc
+++ b/runtime/lib/backends/cpu/providers/default/custom_call/topk.cc
@@ -270,8 +270,8 @@ static void FindTopKElements(const Tdata *input_data, const Shape &input_shape,
 }
 
 template <typename Tdata, typename Tidx>
-void TopKImpl(const OpAccessor &accessor, WorkQueue *work_queue,
-              int brt_omp_num_threads) {
+void TopKImpl(const OpAccessor &accessor, WorkQueue *work_queue, int op_id,
+              const std::vector<int> &dependency, int brt_omp_num_threads) {
   const auto &shape = accessor.GetArgShape(0);
   const int64_t num_elements = accessor.GetNumElementsOfShape(shape);
   // get input data
@@ -298,7 +298,7 @@ void TopKImpl(const OpAccessor &accessor, WorkQueue *work_queue,
   BRT_ENFORCE(output_values_size == output_size);
   const auto &output_shape = accessor.GetArgShape(1);
   // TODO: add support for thread pool
-  DispatchHostTask(work_queue, {
+  DispatchHostTask(work_queue, op_id, dependency, {
     (FindTopKElements<Tdata, Tidx, GreaterValueCmp<Tdata, Tidx>>(
         data, shape, values, indices, output_shape, k, sort_top_k, axis_value,
         brt_omp_num_threads));
@@ -316,8 +316,9 @@ common::Status TopK::RunImpl(const ExecutionContext &ctx) {
 #define HANDLE_DTYPE(DType, IType)                                             \
   if (data_dtype == DType && index_dtype == IType) {                           \
     TopKImpl<typename DTypeTraits<DType>::type_t,                              \
-             typename DTypeTraits<IType>::type_t>(accessor, ctx.work_queue,    \
-                                                  GetNumThreads());            \
+             typename DTypeTraits<IType>::type_t>(                             \
+        accessor, ctx.work_queue, info_.GetOpId(), info_.GetDependency(),      \
+        GetNumThreads());                                                      \
     return common::Status::OK();                                               \
   }
   HANDLE_DTYPE(DTypeEnum::Float16, DTypeEnum::Int32)
diff --git a/runtime/lib/backends/cpu/providers/default/llvm/jit.cc b/runtime/lib/backends/cpu/providers/default/llvm/jit.cc
index 126144565..f54e3cc4b 100644
--- a/runtime/lib/backends/cpu/providers/default/llvm/jit.cc
+++ b/runtime/lib/backends/cpu/providers/default/llvm/jit.cc
@@ -93,7 +93,7 @@ common::Status LLVMJITOpKernel::RunImpl(const ExecutionContext &ctx) {
   } else {
     BRT_ENFORCE(jit->LookupPacked(symbol_name, &symbol).IsOK());
   }
-  DispatchHostTask(ctx.work_queue, {
+  DispatchHostTask(ctx.work_queue, info_.GetOpId(), info_.GetDependency(), {
     std::vector<void *> args;
     args.reserve(nr_args);
     for (size_t i = 0; i < nr_args; ++i) {
diff --git a/runtime/lib/backends/cpu/providers/default/math/elementwise_ops.cc b/runtime/lib/backends/cpu/providers/default/math/elementwise_ops.cc
index 39a3d726c..c4bd72bc1 100644
--- a/runtime/lib/backends/cpu/providers/default/math/elementwise_ops.cc
+++ b/runtime/lib/backends/cpu/providers/default/math/elementwise_ops.cc
@@ -42,7 +42,7 @@ common::Status Add<T>::RunImpl(const ExecutionContext &ctx) {
   T *b = static_cast<T *>(accessor.GetArgAsyncValueRef(1));
   T *c = static_cast<T *>(accessor.GetArgAsyncValueRef(2));
 
-  DispatchHostTask(ctx.work_queue, {
+  DispatchHostTask(ctx.work_queue, info_.GetOpId(), info_.GetDependency(), {
     for (unsigned int i = 0; i < N; ++i) {
       c[i] = a[i] + b[i];
     }
diff --git a/runtime/lib/backends/cpu/providers/default/tensor_generate/fill.cc b/runtime/lib/backends/cpu/providers/default/tensor_generate/fill.cc
index eae59afea..4987ef2b0 100644
--- a/runtime/lib/backends/cpu/providers/default/tensor_generate/fill.cc
+++ b/runtime/lib/backends/cpu/providers/default/tensor_generate/fill.cc
@@ -31,7 +31,7 @@ common::Status Fill::RunImpl(const ExecutionContext &ctx) {
     // TODO: take the ownership of the underlying data of the
     // string_view which belongs to IRHandle
     auto value = accessor.GetAttrAsSplatValue<StringView>("value");
-    DispatchHostTask(ctx.work_queue, {
+    DispatchHostTask(ctx.work_queue, info_.GetOpId(), info_.GetDependency(), {
       std::fill_n(reinterpret_cast<StringView *>(p), length, value);
     });
     return common::Status::OK();
diff --git a/runtime/lib/backends/cpu/providers/default/tensor_generate/rng_state.cc b/runtime/lib/backends/cpu/providers/default/tensor_generate/rng_state.cc
index fca4f4b54..560034534 100644
--- a/runtime/lib/backends/cpu/providers/default/tensor_generate/rng_state.cc
+++ b/runtime/lib/backends/cpu/providers/default/tensor_generate/rng_state.cc
@@ -41,7 +41,7 @@ common::Status GetSeedOpKernel::RunImpl(const ExecutionContext &ctx) {
 #define CASE(D)                                                                \
   case DTypeEnum::D: {                                                         \
     using ctype = DTypeTraits<DTypeEnum::D>::type_t;                           \
-    DispatchHostTask(ctx.work_queue, {                                         \
+    DispatchHostTask(ctx.work_queue, info_.GetOpId(), info_.GetDependency(), { \
       *reinterpret_cast<ctype *>(device_p) = static_cast<ctype>(rngSeed);      \
     });                                                                        \
     return common::Status::OK();                                               \
@@ -75,7 +75,7 @@ common::Status NextOffsetOpKernel::RunImpl(const ExecutionContext &ctx) {
 #define CASE(D)                                                                \
   case DTypeEnum::D: {                                                         \
     using ctype = DTypeTraits<DTypeEnum::D>::type_t;                           \
-    DispatchHostTask(ctx.work_queue, {                                         \
+    DispatchHostTask(ctx.work_queue, info_.GetOpId(), info_.GetDependency(), { \
       *reinterpret_cast<ctype *>(device_p) = static_cast<ctype>(rngOffset);    \
     });                                                                        \
     return common::Status::OK();                                               \
diff --git a/runtime/lib/backends/cuda/device/cuda_work_queue.cc b/runtime/lib/backends/cuda/device/cuda_work_queue.cc
index dd7edfbd7..8eecda95e 100644
--- a/runtime/lib/backends/cuda/device/cuda_work_queue.cc
+++ b/runtime/lib/backends/cuda/device/cuda_work_queue.cc
@@ -80,6 +80,12 @@ inline common::Status ComputeDrv(const void *func, void **args,
                      (*block).z, *shared_size, stream, kernel_args, 0));
 }
 
+inline common::Status ComputeHost(const void *func, void **args,
+                                  CUstream_st *stream) {
+  return BRT_CUDA_CALL(cudaLaunchHostFunc(
+      stream, reinterpret_cast<CUhostFn>(const_cast<void *>(func)), *args));
+}
+
 inline common::Status RecordEvent(CUevent_st *event, CUstream_st *stream) {
   return BRT_CUDA_CALL(cudaEventRecord(event, stream));
 }
@@ -91,7 +97,8 @@ inline common::Status WaitEvent(CUevent_st *event, CUstream_st *stream) {
 } // namespace
 
 common::Status CUDAWorkQueue::AddTask(int task_type, const void *func,
-                                      void **args) {
+                                      void **args, int op_id,
+                                      const std::vector<int> &dependency) {
   GetCudaEnv().Activate();
   switch (task_type) {
   case CUDATaskType::kCompute:
@@ -111,12 +118,6 @@ common::Status CUDAWorkQueue::AddTask(int task_type, const void *func,
                 "unsupported task type " + std::to_string(task_type));
 }
 
-common::Status
-CUDAWorkQueue::AddEventWait(mlir::Operation *op,
-                            std::vector<mlir::Operation *> dependency_list) {
-  return Status::OK();
-}
-
 common::Status CUDAWorkQueue::Sync() {
   GetCudaEnv().Activate();
   return BRT_CUDA_CALL(cudaDeviceSynchronize());
@@ -133,9 +134,10 @@ CUDASingleStreamWorkQueue::~CUDASingleStreamWorkQueue() {
   BRT_CUDA_CHECK(cudaStreamDestroy(stream_));
 }
 
-common::Status CUDASingleStreamWorkQueue::AddTask(int task_type,
-                                                  const void *func,
-                                                  void **args) {
+common::Status
+CUDASingleStreamWorkQueue::AddTask(int task_type, const void *func, void **args,
+                                   int op_id,
+                                   const std::vector<int> &dependency) {
   GetCudaEnv().Activate();
 
   switch (task_type) {
@@ -156,32 +158,27 @@ common::Status CUDASingleStreamWorkQueue::AddTask(int task_type,
                 "unsupported task type " + std::to_string(task_type));
 }
 
-common::Status CUDASingleStreamWorkQueue::AddEventWait(
-    mlir::Operation *op, std::vector<mlir::Operation *> dependency_list) {
-  // single stream, no need to wait
-  return Status::OK();
-}
-
 common::Status CUDASingleStreamWorkQueue::Sync() {
   GetCudaEnv().Activate();
   return BRT_CUDA_CALL(cudaStreamSynchronize(stream_));
 }
 
-CUDAOneComputeTwoTransferWorkQueue::CUDAOneComputeTwoTransferWorkQueue(
-    int device_id)
-    : CUDAWorkQueue(device_id, "cuda_1_compute_2_copy_stream") {
+CUDAMultiStreamWorkQueue::CUDAMultiStreamWorkQueue(int device_id)
+    : CUDAWorkQueue(device_id, "cuda_1_compute_2_copy_1_host_stream") {
 
   GetCudaEnv().Activate();
   BRT_CUDA_CHECK(cudaStreamCreate(&streams_[0]));
   BRT_CUDA_CHECK(cudaStreamCreate(&streams_[1]));
   BRT_CUDA_CHECK(cudaStreamCreate(&streams_[2]));
+  BRT_CUDA_CHECK(cudaStreamCreate(&streams_[3]));
 }
 
-CUDAOneComputeTwoTransferWorkQueue::~CUDAOneComputeTwoTransferWorkQueue() {
+CUDAMultiStreamWorkQueue::~CUDAMultiStreamWorkQueue() {
   GetCudaEnv().Activate();
   BRT_CUDA_CHECK(cudaStreamDestroy(streams_[0]));
   BRT_CUDA_CHECK(cudaStreamDestroy(streams_[1]));
   BRT_CUDA_CHECK(cudaStreamDestroy(streams_[2]));
+  BRT_CUDA_CHECK(cudaStreamDestroy(streams_[3]));
 
   // destroy all events
   for (auto event : events_) {
@@ -211,25 +208,32 @@ inline CUstream_st *GetStream(void **args, CUstream_st **streams) {
 }
 } // namespace
 
-common::Status CUDAOneComputeTwoTransferWorkQueue::AddTask(int task_type,
-                                                           const void *func,
-                                                           void **args) {
+common::Status
+CUDAMultiStreamWorkQueue::AddTask(int task_type, const void *func, void **args,
+                                  int op_id,
+                                  const std::vector<int> &dependency) {
   GetCudaEnv().Activate();
 
   switch (task_type) {
   case CUDATaskType::kCompute:
+    id_to_stream_map_[op_id] = 0;
+    AddEventWait(0, dependency);
     return Compute(func, args, streams_[0]);
   case CUDATaskType::kComputeDrv:
+    id_to_stream_map_[op_id] = 0;
+    AddEventWait(0, dependency);
     return ComputeDrv(func, args, streams_[0]);
   case CUDATaskType::kH2D:
+    id_to_stream_map_[op_id] = 1;
+    AddEventWait(1, dependency);
     return CopyH2D(args, streams_[1]);
   case CUDATaskType::kD2H:
+    id_to_stream_map_[op_id] = 2;
+    AddEventWait(2, dependency);
     return CopyD2H(args, streams_[2]);
-  case CUDATaskType::kRecordEvent:
-    return RecordEvent(GetEvent(args, events_), GetStream(args, streams_));
-  case CUDATaskType::kWaitEvent:
-    return WaitEvent(GetEvent(args), GetStream(args, streams_));
   case CUDATaskType::kD2D:
+    id_to_stream_map_[op_id] = 0;
+    AddEventWait(0, dependency);
     return CopyD2D(args, streams_[0]);
   default:;
   }
@@ -238,43 +242,42 @@ common::Status CUDAOneComputeTwoTransferWorkQueue::AddTask(int task_type,
                 "unsupported task type " + std::to_string(task_type));
 }
 
-common::Status CUDAOneComputeTwoTransferWorkQueue::AddEventWait(
-    mlir::Operation *op, std::vector<mlir::Operation *> dependency_list) {
-  for (auto dep_op : dependency_list) {
-    size_t wait_stream = GetStreamIdx(dep_op);
-    size_t curr_stream = GetStreamIdx(op);
-    if (curr_stream == wait_stream)
+common::Status
+CUDAMultiStreamWorkQueue::AddHostTask(const void *func, void **args, int op_id,
+                                      const std::vector<int> &dependency) {
+  id_to_stream_map_[op_id] = 3;
+  AddEventWait(3, dependency);
+  return ComputeHost(func, args, streams_[3]);
+}
+
+common::Status
+CUDAMultiStreamWorkQueue::AddEventWait(size_t stream,
+                                       std::vector<int> wait_ids) {
+  for (auto wait_id : wait_ids) {
+    size_t wait_stream = id_to_stream_map_[wait_id];
+    if (stream == wait_stream)
       continue;
     void *record_args[] = {&wait_stream, nullptr /*placeholder for event*/};
-    this->AddTask(CUDATaskType::kRecordEvent, nullptr, record_args);
-    void *wait_args[] = {&curr_stream, record_args[1]};
-    this->AddTask(CUDATaskType::kWaitEvent, nullptr, wait_args);
+    RecordEvent(GetEvent(record_args, events_),
+                GetStream(record_args, streams_));
+    void *wait_args[] = {&stream, record_args[1]};
+    WaitEvent(GetEvent(wait_args), GetStream(wait_args, streams_));
   }
   return Status::OK();
 }
 
-common::Status CUDAOneComputeTwoTransferWorkQueue::Sync() {
+common::Status CUDAMultiStreamWorkQueue::Sync() {
   GetCudaEnv().Activate();
   return BRT_CUDA_CALL(cudaDeviceSynchronize());
 }
 
-size_t CUDAOneComputeTwoTransferWorkQueue::GetStreamIdx(mlir::Operation *op) {
-  auto byre_op = static_cast<mlir::byre::ByreOp>(op);
-  const std::string &key = ir::ByREHandle::GetKey(byre_op);
-  if (key == "cpu2cuda") {
-    return 1;
-  } else if (key == "cuda2cpu") {
-    return 2;
-  }
-  return 0;
-}
-
 CUDAExternalStreamWorkQueue::CUDAExternalStreamWorkQueue(CUstream_st *stream)
     : CUDAWorkQueue(stream, "cuda_external_stream"), stream_(stream) {}
 
-common::Status CUDAExternalStreamWorkQueue::AddTask(int task_type,
-                                                    const void *func,
-                                                    void **args) {
+common::Status
+CUDAExternalStreamWorkQueue::AddTask(int task_type, const void *func,
+                                     void **args, int op_id,
+                                     const std::vector<int> &dependency) {
   GetCudaEnv().Activate();
 
   switch (task_type) {
@@ -298,5 +301,4 @@ common::Status CUDAExternalStreamWorkQueue::Sync() {
   GetCudaEnv().Activate();
   return BRT_CUDA_CALL(cudaStreamSynchronize(stream_));
 }
-
 } // namespace brt
diff --git a/runtime/lib/backends/cuda/providers/default/codegen/ptx.cc b/runtime/lib/backends/cuda/providers/default/codegen/ptx.cc
index 39c77d5a3..25c14c9bf 100644
--- a/runtime/lib/backends/cuda/providers/default/codegen/ptx.cc
+++ b/runtime/lib/backends/cuda/providers/default/codegen/ptx.cc
@@ -221,7 +221,8 @@ common::Status PTXOpKernel::RunImpl(const ExecutionContext &ctx) {
   CUfunction func = impl_->GetOrCreateFunction(cuda_env.GetDeviceID());
 
   return work_queue->AddTask(5 /*CUDATaskType::kComputeDrv*/, (void *)func,
-                             args.data());
+                             args.data(), info_.GetOpId(),
+                             info_.GetDependency());
 }
 } // namespace cuda
 } // namespace brt
diff --git a/runtime/lib/backends/cuda/providers/default/copy/copy.cc b/runtime/lib/backends/cuda/providers/default/copy/copy.cc
index 35466aa0d..00900c21d 100644
--- a/runtime/lib/backends/cuda/providers/default/copy/copy.cc
+++ b/runtime/lib/backends/cuda/providers/default/copy/copy.cc
@@ -64,8 +64,8 @@ common::Status CopyOpKernel::RunImpl(const ExecutionContext &ctx) {
   args[1] = &src_value;
   args[2] = &byte_size;
   auto work_queue = static_cast<CUDAWorkQueue *>(ctx.work_queue);
-  work_queue->AddEventWait(info_.GetOperation(), info_.GetDependency());
-  return work_queue->AddTask(task_type, nullptr, args.data());
+  return work_queue->AddTask(task_type, nullptr, args.data(), info_.GetOpId(),
+                             info_.GetDependency());
 }
 
 } // namespace cuda
diff --git a/runtime/lib/backends/cuda/providers/default/math/elementwise_ops.cc b/runtime/lib/backends/cuda/providers/default/math/elementwise_ops.cc
index ef1e46902..899e06085 100644
--- a/runtime/lib/backends/cuda/providers/default/math/elementwise_ops.cc
+++ b/runtime/lib/backends/cuda/providers/default/math/elementwise_ops.cc
@@ -67,8 +67,8 @@ common::Status Add<T>::RunImpl(const ExecutionContext &ctx) {
   }
 
   args.push_back(&n); // n
-  ctx.work_queue->AddEventWait(info_.GetOperation(), info_.GetDependency());
-  return ctx.work_queue->AddTask(0, (void *)add_kernel<T>, args.data());
+  return ctx.work_queue->AddTask(0, (void *)add_kernel<T>, args.data(),
+                                 info_.GetOpId(), info_.GetDependency());
 }
 
 // instantiate
diff --git a/runtime/lib/core/framework/execution_plan.cc b/runtime/lib/core/framework/execution_plan.cc
index 154b2b006..846c56aa4 100644
--- a/runtime/lib/core/framework/execution_plan.cc
+++ b/runtime/lib/core/framework/execution_plan.cc
@@ -311,6 +311,7 @@ common::Status StaticBRTExecutionPlan::ProloguePerSession(
   if (!status_internal.IsOK())
     return status_internal;
 
+  std::unordered_map<Operation *, int> op_to_id_map;
   // create op kernel, generate tensor id and mapping IR value to it
   graph_.IterateNode([&](Operation *op) {
     if (auto byre_op = dyn_cast<byre::ByreOp>(op)) {
@@ -385,11 +386,23 @@ common::Status StaticBRTExecutionPlan::ProloguePerSession(
           frame_construct_info_.dependency_graph[op] = {};
         }
         // creat an OpKerenl based on the hitting provider
+        int op_id = op_kernels_.size();
+        op_to_id_map[op] = op_id;
+        std::vector<int> dependency_list;
+        for (Operation *dep_op : frame_construct_info_.dependency_graph[op]) {
+          if (op_to_id_map.find(dep_op) != op_to_id_map.end()) {
+            dependency_list.push_back(op_to_id_map[dep_op]);
+          } else {
+            status_internal = Status(BRT, FAIL, " op_to_id_map op not found ");
+            return WalkResult::interrupt();
+          }
+        }
+
         OpKernelInfo op_kernel_info(
-            *provider, graph_, op, allocators, last_alloc,
+            *provider, graph_, op, op_id, allocators, last_alloc,
             graph_info_.tensor_to_id, graph_info_.scalar_to_id,
             frame_construct_info_.weights, intermediate_begin,
-            graph_.GetIRPath(), frame_construct_info_.dependency_graph[op]);
+            graph_.GetIRPath(), dependency_list);
 
         auto op_ptr = (*registry)(key, op_kernel_info);
 
diff --git a/runtime/test/backends/cuda/device/cuda_work_queue_test.cc b/runtime/test/backends/cuda/device/cuda_work_queue_test.cc
index 54058d0d2..3b179a71b 100644
--- a/runtime/test/backends/cuda/device/cuda_work_queue_test.cc
+++ b/runtime/test/backends/cuda/device/cuda_work_queue_test.cc
@@ -65,10 +65,10 @@ TEST(CUDAWorkQueueTest, CUDAWorkQueue) {
   cudaDeviceSynchronize();
 
   void *args1[] = {&grid, &block, &shared_size, &arr1, &arr2, &n, &val1};
-  wq.AddTask(0, (void *)test_kernel, args1);
+  wq.AddTask(0, (void *)test_kernel, args1, 0, {});
 
   void *args2[] = {&grid, &block, &shared_size, &arr2, &arr3, &n, &val2};
-  wq.AddTask(0, (void *)test_kernel, args2);
+  wq.AddTask(0, (void *)test_kernel, args2, 0, {});
 
   wq.Sync();
 
@@ -112,13 +112,13 @@ TEST(CUDAWorkQueueTest, CUDASingleStreamWorkQueue) {
   cudaDeviceSynchronize();
 
   void *args0[] = {&arr1, &host1, &count};
-  wq.AddTask(1 /*h2d*/, nullptr, args0);
+  wq.AddTask(1 /*h2d*/, nullptr, args0, 0, {});
 
   void *args1[] = {&grid, &block, &shared_size, &arr1, &arr2, &n, &val1};
-  wq.AddTask(0, (void *)test_kernel, args1);
+  wq.AddTask(0, (void *)test_kernel, args1, 0, {});
 
   void *args2[] = {&grid, &block, &shared_size, &arr2, &arr3, &n, &val2};
-  wq.AddTask(0, (void *)test_kernel, args2);
+  wq.AddTask(0, (void *)test_kernel, args2, 0, {});
 
   wq.Sync();
 
@@ -130,8 +130,8 @@ TEST(CUDAWorkQueueTest, CUDASingleStreamWorkQueue) {
   BRT_CUDA_CHECK(cudaFree(arr3));
 }
 
-TEST(CUDAWorkQueueTest, CUDAOneComputeTwoTransferWorkQueue) {
-  CUDAOneComputeTwoTransferWorkQueue wq(0);
+TEST(CUDAWorkQueueTest, CUDAMultiStreamWorkQueue) {
+  CUDAMultiStreamWorkQueue wq(0);
 
   int gx = 4;
   int bx = 256;
@@ -161,21 +161,21 @@ TEST(CUDAWorkQueueTest, CUDAOneComputeTwoTransferWorkQueue) {
   cudaDeviceSynchronize();
 
   void *args0[] = {&arr1, &host1, &count};
-  wq.AddTask(1 /*H2D*/, nullptr, args0);
+  wq.AddTask(1 /*H2D*/, nullptr, args0, 0, {});
 
   size_t streamId1 = 1;
   void *args1[] = {&streamId1, nullptr /*placeholder for event*/};
-  wq.AddTask(3 /*RecordEvent*/, nullptr, args1);
+  wq.AddTask(3 /*RecordEvent*/, nullptr, args1, 0, {});
 
   size_t streamId2 = 0;
   void *args2[] = {&streamId2, args1[1]};
-  wq.AddTask(4 /*WaitEvent*/, nullptr, args2);
+  wq.AddTask(4 /*WaitEvent*/, nullptr, args2, 0, {});
 
   void *args3[] = {&grid, &block, &shared_size, &arr1, &arr2, &n, &val1};
-  wq.AddTask(0, (void *)test_kernel, args3);
+  wq.AddTask(0, (void *)test_kernel, args3, 0, {});
 
   void *args4[] = {&grid, &block, &shared_size, &arr2, &arr3, &n, &val2};
-  wq.AddTask(0, (void *)test_kernel, args4);
+  wq.AddTask(0, (void *)test_kernel, args4, 0, {});
 
   wq.Sync();
 
@@ -225,7 +225,7 @@ TEST(CUDAWorkQueueTest, CUDAExternalStreamWorkQueue) {
                    stream_external);
 
   void *args2[] = {&grid, &block, &shared_size, &arr2, &arr3, &n, &val2};
-  wq.AddTask(0, (void *)test_kernel, args2);
+  wq.AddTask(0, (void *)test_kernel, args2, 0, {});
 
   wq.Sync();
 
@@ -238,6 +238,75 @@ TEST(CUDAWorkQueueTest, CUDAExternalStreamWorkQueue) {
   BRT_CUDA_CHECK(cudaStreamDestroy(stream_external));
 }
 
+namespace {
+using HostArgs = std::tuple<void *, size_t>;
+
+void CUDART_CB host_memset(void *args) {
+  auto host_args = reinterpret_cast<HostArgs *>(args);
+  AssignCPUBuffer<float>(reinterpret_cast<float *>(std::get<0>(*host_args)),
+                         std::get<1>(*host_args) / sizeof(float), 1.0);
+}
+} // namespace
+
+TEST(CUDAWorkQueueTest, CUDAMultiStreamWithHostWorkQueue) {
+  CUDAMultiStreamWorkQueue wq(0);
+
+  int gx = 4;
+  int bx = 256;
+
+  dim3 grid(gx, 1, 1);
+  dim3 block(bx, 1, 1);
+  size_t shared_size = 0;
+
+  float *host1;
+  float *arr1;
+  float *arr2;
+  float *arr3;
+  int n = gx * bx;
+  float val1 = 1.0f;
+  float val2 = 2.0f;
+
+  size_t count = n * sizeof(float);
+
+  // Using a regular pageable host allocation (malloc) is likely to fail because
+  // the host function may not fire immediately (given it's on another thread),
+  // and the copy H2D may start before the host function complete.
+  // More details:
+  // https://stackoverflow.com/questions/50901819/cudastreamaddcallback-doesnt-block-later-cudamemcpyasync
+  cudaHostAlloc((void **)&host1, count, cudaHostAllocDefault);
+  memset(host1, 0, count);
+
+  BRT_CUDA_CHECK(cudaMalloc(&arr1, count));
+  BRT_CUDA_CHECK(cudaMemset(arr1, -1, n * sizeof(float)));
+  BRT_CUDA_CHECK(cudaMalloc(&arr2, count));
+  BRT_CUDA_CHECK(cudaMalloc(&arr3, count));
+
+  cudaDeviceSynchronize();
+
+  HostArgs *host_args = new HostArgs((void *)host1, count);
+  void *host_args_ptr = reinterpret_cast<void *>(host_args);
+  wq.AddHostTask((void *)host_memset, &host_args_ptr, 0, {});
+
+  void *args0[] = {&arr1, &host1, &count};
+  wq.AddTask(1 /*H2D*/, nullptr, args0, 1, {0}); // depend on Host
+
+  void *args3[] = {&grid, &block, &shared_size, &arr1, &arr2, &n, &val1};
+  wq.AddTask(0, (void *)test_kernel, args3, 2, {1}); // depend on H2D
+
+  void *args4[] = {&grid, &block, &shared_size, &arr2, &arr3, &n, &val2};
+  wq.AddTask(0, (void *)test_kernel, args4, 3, {2});
+
+  wq.Sync();
+
+  CheckResult(arr3, n, 4.0f);
+
+  cudaFreeHost(host1);
+  delete (host_args);
+  BRT_CUDA_CHECK(cudaFree(arr1));
+  BRT_CUDA_CHECK(cudaFree(arr2));
+  BRT_CUDA_CHECK(cudaFree(arr3));
+}
+
 namespace {
 void check_work_queue_multi_gpus(
     std::vector<std::unique_ptr<CUDAWorkQueue>> &work_queues) {
@@ -270,11 +339,11 @@ void check_work_queue_multi_gpus(
     dim3 block(bx, 1, 1);
     void *args1[] = {&grid,    &block, &shared_size, &arr1[i],
                      &arr2[i], &n,     &val1};
-    wq->AddTask(0, (void *)test_kernel, args1);
+    wq->AddTask(0, (void *)test_kernel, args1, 0, {});
 
     void *args2[] = {&grid,    &block, &shared_size, &arr2[i],
                      &arr3[i], &n,     &val2};
-    wq->AddTask(0, (void *)test_kernel, args2);
+    wq->AddTask(0, (void *)test_kernel, args2, 0, {});
 
     wq->Sync();
 
@@ -315,7 +384,7 @@ TEST(CUDAWorkQueueTest, CUDASingleStreamWorkQueueMultiGPUs) {
   check_work_queue_multi_gpus(work_queues);
 }
 
-TEST(CUDAWorkQueueTest, CUDAOneComputeTwoTransferWorkQueueMultiGPUs) {
+TEST(CUDAWorkQueueTest, CUDAMultiStreamWorkQueueMultiGPUs) {
   int nr_device;
   BRT_CUDA_CHECK(cudaGetDeviceCount(&nr_device));
   if (nr_device == 1)
@@ -323,8 +392,7 @@ TEST(CUDAWorkQueueTest, CUDAOneComputeTwoTransferWorkQueueMultiGPUs) {
 
   std::vector<std::unique_ptr<CUDAWorkQueue>> work_queues;
   for (int i = 0; i < nr_device; ++i) {
-    work_queues.emplace_back(
-        std::make_unique<CUDAOneComputeTwoTransferWorkQueue>(i));
+    work_queues.emplace_back(std::make_unique<CUDAMultiStreamWorkQueue>(i));
   }
   check_work_queue_multi_gpus(work_queues);
 }
diff --git a/runtime/test/backends/cuda/device/nvrtc_test.cc b/runtime/test/backends/cuda/device/nvrtc_test.cc
index 4b42902a7..92691ebae 100644
--- a/runtime/test/backends/cuda/device/nvrtc_test.cc
+++ b/runtime/test/backends/cuda/device/nvrtc_test.cc
@@ -81,10 +81,10 @@ TEST(NVRTCTest, Add) {
   cudaDeviceSynchronize();
 
   void *args1[] = {&grid, &block, &shared_size, &arr1, &arr2, &n, &val1};
-  wq.AddTask(5, (void *)func, args1);
+  wq.AddTask(5, (void *)func, args1, 0, {});
 
   void *args2[] = {&grid, &block, &shared_size, &arr2, &arr3, &n, &val2};
-  wq.AddTask(5, (void *)func, args2);
+  wq.AddTask(5, (void *)func, args2, 0, {});
 
   wq.Sync();
 
diff --git a/runtime/test/backends/cuda/device/ptx_test.cc b/runtime/test/backends/cuda/device/ptx_test.cc
index 159c438c9..2cbe00433 100644
--- a/runtime/test/backends/cuda/device/ptx_test.cc
+++ b/runtime/test/backends/cuda/device/ptx_test.cc
@@ -106,10 +106,10 @@ TEST(PTXTest, NVCCPTX) {
   cudaDeviceSynchronize();
 
   void* args1[] = { &grid, &block, &shared_size, &arr1, &arr2, &n, &val1 };
-  wq.AddTask(5, (void*)func, args1);
+  wq.AddTask(5, (void*)func, args1, 0, {});
 
   void* args2[] = { &grid, &block, &shared_size, &arr2, &arr3, &n, &val2 };
-  wq.AddTask(5, (void*)func, args2);
+  wq.AddTask(5, (void*)func, args2, 0, {});
 
   wq.Sync();
 
@@ -177,7 +177,7 @@ TEST(PTXTest, LLVMPTX) {
     InsertMemDescToArgs(desc2, args1);
 
     auto &wq = *work_queues[i];
-    wq.AddTask(5, (void *)funcs[i], args1.data());
+    wq.AddTask(5, (void *)funcs[i], args1.data(), 0, {});
 
     std::vector<void *> args2;
     args2.push_back(&grid);
@@ -187,7 +187,7 @@ TEST(PTXTest, LLVMPTX) {
     InsertMemDescToArgs(desc1, args2);
     InsertMemDescToArgs(desc3, args2);
 
-    wq.AddTask(5, (void *)funcs[i], args2.data());
+    wq.AddTask(5, (void *)funcs[i], args2.data(), 0, {});
 
     wq.Sync();
 
@@ -254,7 +254,7 @@ TEST(PTXTest, LLVMPTXBarePtr) {
     args1.push_back(&arr2[i]);
 
     auto &wq = *work_queues[i];
-    wq.AddTask(5, (void *)funcs[i], args1.data());
+    wq.AddTask(5, (void *)funcs[i], args1.data(), 0, {});
 
     std::vector<void *> args2;
     args2.push_back(&grid);
@@ -264,7 +264,7 @@ TEST(PTXTest, LLVMPTXBarePtr) {
     args2.push_back(&arr1[i]);
     args2.push_back(&arr3[i]);
 
-    wq.AddTask(5, (void *)funcs[i], args2.data());
+    wq.AddTask(5, (void *)funcs[i], args2.data(), 0, {});
 
     wq.Sync();
 
diff --git a/runtime/test/backends/cuda/providers/default/kernel/multi_stream_test.cc b/runtime/test/backends/cuda/providers/default/kernel/multi_stream_test.cc
index 210348218..d51023fcb 100644
--- a/runtime/test/backends/cuda/providers/default/kernel/multi_stream_test.cc
+++ b/runtime/test/backends/cuda/providers/default/kernel/multi_stream_test.cc
@@ -33,7 +33,7 @@ using namespace brt::test;
 
 static std::string test_file = "test/test_files/custom_add_cpu2cuda.mlir";
 
-TEST(CUDAWorkQueueTest, CUDAOneComputeTwoTransferWorkQueueAdd) {
+TEST(CUDAWorkQueueTest, CUDAMultiStreamWorkQueueAdd) {
   Session session;
   auto status_allocator = CUDAAllocatorFactory(&session);
   BRT_TEST_CHECK_STATUS(status_allocator);
@@ -49,8 +49,8 @@ TEST(CUDAWorkQueueTest, CUDAOneComputeTwoTransferWorkQueueAdd) {
   BRT_TEST_CHECK_STATUS(status_load);
 
   std::unique_ptr<RequestContext> request;
-  auto status_request = session.NewRequestContext(
-      &request, new CUDAOneComputeTwoTransferWorkQueue(0));
+  auto status_request =
+      session.NewRequestContext(&request, new CUDAMultiStreamWorkQueue(0));
   BRT_TEST_CHECK_STATUS(status_request);
 
   request->FinishIOBinding();
diff --git a/runtime/test/external_kernels/cuda/kernels.cc b/runtime/test/external_kernels/cuda/kernels.cc
index c87b07bb3..647fd9ea9 100644
--- a/runtime/test/external_kernels/cuda/kernels.cc
+++ b/runtime/test/external_kernels/cuda/kernels.cc
@@ -71,7 +71,7 @@ common::Status CustomAdd<T>::RunImpl(const ExecutionContext &ctx) {
   args.push_back(&n); // n
 
   ctx.work_queue->AddTask(0, (void *)external_kernels::add_kernel<T>,
-                          args.data());
+                          args.data(), info_.GetOpId(), info_.GetDependency());
 
   return Status::OK();
 }