[INTEL_HPU] re-enable async runner and do enhancement (#1934)

Author: Leo Zhao

backends/intel_hpu/kernels/hpu_operator.cc

@@ -31,6 +31,9 @@ FLAGS_DEFINE_bool(intel_hpu_sync_execute, false, "set sync execute mode");
 FLAGS_DEFINE_bool(intel_hpu_reciperunner_debug,
                   false,
                   "reciperunner debug log");
+FLAGS_DEFINE_bool(intel_hpu_async_runner_multithread,
+                  false,
+                  "enable async runner multi-threading support");
 
 typedef std::pair<synSectionHandle, bool> sectionWithFirstIndication;
 static std::unordered_map<std::string, sectionWithFirstIndication> sectionMap;
@@ -142,6 +145,47 @@ void RecipeRunner::prepareTensorInfo(synRecipeHandle recipe,
 }
 
 #ifdef ENABLE_ASYNC_RUN
+void GlobalWorkStreamExecutor::add_task(const synStreamHandle stream,
+                                        std::function<void()> task) {
+  std::shared_ptr<WorkerThread> worker;
+
+  synStreamHandle workstream = nullptr;
+  if (FLAGS_intel_hpu_async_runner_multithread) workstream = stream;
+
+  {
+    std::lock_guard<std::mutex> lock(workers_mutex_);
+    if (workers_.find(workstream) == workers_.end()) {
+      worker = std::make_shared<WorkerThread>();
+      worker->thread = std::thread([this, workstream, worker]() {
+        while (true) {
+          std::function<void()> task;
+          {
+            std::unique_lock<std::mutex> lock(worker->mutex);
+            worker->condition.wait(lock, [worker] {
+              return !worker->tasks.empty() || worker->stop;
+            });
+
+            if (worker->stop && worker->tasks.empty()) break;
+
+            task = std::move(worker->tasks.front());
+            worker->tasks.pop();
+          }
+          task();
+        }
+      });
+      workers_[workstream] = worker;
+    } else {
+      worker = workers_[workstream];
+    }
+  }
+
+  {
+    std::lock_guard<std::mutex> lock(worker->mutex);
+    worker->tasks.emplace(std::move(task));
+  }
+  worker->condition.notify_one();
+}
+
 void RecipeRunner::ExecuteRecipe(C_Stream stream,
                                  const std::map<std::string, uint64_t>& tensors,
                                  synRecipeHandle recipeHandle_) {

backends/intel_hpu/kernels/hpu_operator.h

@@ -33,7 +33,7 @@
 #include "paddle/phi/common/type_traits.h"
 #include "paddle/phi/extension.h"
 
-// #define ENABLE_ASYNC_RUN
+#define ENABLE_ASYNC_RUN
 
 class HpuOperator {
  public:
@@ -88,78 +88,64 @@ class GlobalWorkStreamExecutor {
   }
 
   template <typename R>
-  std::future<R> async(std::function<R()> func) {
+  std::future<R> async(synStreamHandle stream, std::function<R()> func) {
     auto task = std::make_shared<std::packaged_task<R()>>(std::move(func));
     std::future<R> res = task->get_future();
-    add_task([task]() { (*task)(); });
+    add_task(stream, [task]() { (*task)(); });
     return res;
   }
 
   template <typename F>
-  auto async(F&& func) -> std::future<decltype(func())> {
+  auto async(synStreamHandle stream, F&& func)
+      -> std::future<decltype(func())> {
     using R = decltype(func());
     auto task =
         std::make_shared<std::packaged_task<R()>>(std::forward<F>(func));
     std::future<R> res = task->get_future();
-    add_task([task]() { (*task)(); });
+    add_task(stream, [task]() { (*task)(); });
     return res;
   }
 
   template <typename R>
-  R sync(std::function<R()> func) {
-    return async(std::move(func)).get();
+  R sync(synStreamHandle stream, std::function<R()> func) {
+    return async(stream, std::move(func)).get();
   }
 
   template <typename F>
-  auto sync(F&& func) -> decltype(func()) {
-    return async(std::forward<F>(func)).get();
+  auto sync(synStreamHandle stream, F&& func) -> decltype(func()) {
+    return async(stream, std::forward<F>(func)).get();
   }
 
  private:
-  GlobalWorkStreamExecutor() {
-    worker_ = std::thread([this]() {
-      while (true) {
-        std::function<void()> task;
-        {
-          std::unique_lock<std::mutex> lock(queue_mutex_);
-          condition_.wait(lock, [this] { return !tasks_.empty() || stop_; });
-
-          if (stop_ && tasks_.empty()) break;
-
-          task = std::move(tasks_.front());
-          tasks_.pop();
-        }
-        task();  // 执行任务
-      }
-    });
-  }
-
+  struct WorkerThread {
+    std::thread thread;
+    std::queue<std::function<void()>> tasks;
+    std::mutex mutex;
+    std::condition_variable condition;
+    bool stop = false;
+  };
+
+  GlobalWorkStreamExecutor() = default;
   ~GlobalWorkStreamExecutor() {
-    {
-      std::lock_guard<std::mutex> lock(queue_mutex_);
-      stop_ = true;
+    for (auto& [stream, worker] : workers_) {
+      {
+        std::lock_guard<std::mutex> lock(worker->mutex);
+        worker->stop = true;
+      }
+      worker->condition.notify_all();
+      if (worker->thread.joinable()) {
+        worker->thread.join();
+      }
     }
-    condition_.notify_all();
-    if (worker_.joinable()) worker_.join();
   }
 
-  void add_task(std::function<void()> task) {
-    {
-      std::lock_guard<std::mutex> lock(queue_mutex_);
-      tasks_.emplace(std::move(task));
-    }
-    condition_.notify_one();
-  }
+  void add_task(const synStreamHandle stream, std::function<void()> task);
 
-  // 删除拷贝构造和赋值
   GlobalWorkStreamExecutor(const GlobalWorkStreamExecutor&) = delete;
   GlobalWorkStreamExecutor& operator=(const GlobalWorkStreamExecutor&) = delete;
 
-  std::thread worker_;
-  std::queue<std::function<void()>> tasks_;
-  std::mutex queue_mutex_;
-  std::condition_variable condition_;
-  bool stop_ = false;
+  std::unordered_map<synStreamHandle, std::shared_ptr<WorkerThread>> workers_;
+  std::mutex workers_mutex_;
 };
 #endif
 
@@ -175,6 +161,7 @@ class RecipeRunner {
   void Run(C_Stream stream, std::map<std::string, uint64_t> tensors) {
     synRecipeHandle recipehandle = this->recipeHandle_;
     auto future = GlobalWorkStreamExecutor::instance().async(
+        reinterpret_cast<synStreamHandle>(stream),
         [this, stream, tensors, recipehandle] {
           ExecuteRecipe(stream, tensors, recipehandle);
         });