release np array in main thread to eliminate the time-consuming gil r…

…equest (#10050) Signed-off-by: daquexian <daquexian566@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Oneflow-Inc · Apr 1, 2023 · b7aac72 · b7aac72
1 parent 6514a17
commit b7aac72
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 7 deletions.
diff --git a/ci/fixed-dev-requirements.txt b/ci/fixed-dev-requirements.txt
@@ -1,4 +1,2 @@
 numpy==1.22.1 ; python_version >= "3.10"
-numpy==1.20.0 ; python_version >= "3.9" and python_version < "3.10"
-numpy==1.18.0 ; python_version >= "3.8" and python_version < "3.9"
-numpy==1.17.0 ; python_version >= "3.6" and python_version < "3.8"
+numpy==1.21.6 ; python_version >= "3.7" and python_version < "3.10"
diff --git a/dev-requirements.txt b/dev-requirements.txt
@@ -1,6 +1,6 @@
 black==19.10b0; python_version >= "3.6"
 click==8.0.0; python_version >= "3.6" # https://github.com/psf/black/issues/2964
-numpy>=1.17.0
+numpy>=1.21.6
 protobuf>=3.9.2, <4.0
 wheel
 tqdm

diff --git a/oneflow/api/python/utils/tensor_utils.h b/oneflow/api/python/utils/tensor_utils.h
@@ -138,10 +138,10 @@ inline Maybe<void> CopyBetweenLocalTensorAndNumpy(
   } else {
     Py_INCREF(array);
     NumPyArrayPtr array_ptr(array, [array]() {
-      CHECK_JUST(Singleton<ForeignLockHelper>::Get()->WithScopedAcquire([&]() -> Maybe<void> {
+      // release array in main thread to eliminate the time-consuming gil request
+      CHECK_JUST(SingletonMaybe<VirtualMachine>())->add_main_thread_pending_task([array]() {
         Py_DECREF(array);
-        return Maybe<void>::Ok();
-      }));
+      });
     });
 
     JUST(PhysicalRun([&](InstructionsBuilder* builder) -> Maybe<void> {

diff --git a/oneflow/core/vm/virtual_machine.cpp b/oneflow/core/vm/virtual_machine.cpp
@@ -208,6 +208,7 @@ Maybe<void> VirtualMachine::ShrinkAllMem() {
 
 VirtualMachine::~VirtualMachine() {
   if (!threads_closed_) { CHECK_JUST(CloseVMThreads()); }
+  RunMainThreadPendingTasks();
   CHECK(engine_->SchedulerEmpty());
   engine_.Reset();
 }
@@ -238,8 +239,17 @@ std::string VirtualMachine::GetBlockingDebugString() {
   return engine_->GetLivelyInstructionListDebugString(limit);
 }
 
+void VirtualMachine::RunMainThreadPendingTasks() {
+  std::unique_lock lock(main_thread_pending_tasks_mutex_);
+  for (const auto& main_thread_pending_task : main_thread_pending_tasks_) {
+    main_thread_pending_task();
+  }
+  main_thread_pending_tasks_.clear();
+}
+
 Maybe<void> VirtualMachine::Receive(vm::InstructionList* instruction_list) {
   SyncVmModeGuard guard(SyncVmMode::kEnable);
+  RunMainThreadPendingTasks();
   if (unlikely(pthread_fork::IsForkedSubProcess())) {
     INTRUSIVE_FOR_EACH_PTR(instruction, instruction_list) {
       const auto& device = instruction->stream().device();

diff --git a/oneflow/core/vm/virtual_machine.h b/oneflow/core/vm/virtual_machine.h
@@ -52,6 +52,11 @@ class VirtualMachine final {
 
   size_t flying_instruction_cnt() const { return engine().flying_instruction_cnt(); }
 
+  void add_main_thread_pending_task(std::function<void()> task) {
+    std::unique_lock lock(main_thread_pending_tasks_mutex_);
+    main_thread_pending_tasks_.push_back(std::move(task));
+  }
+
  private:
   friend class InstructionsBuilder;
 
@@ -80,6 +85,8 @@ class VirtualMachine final {
 
   Maybe<void> CloseWorkerThreads();
 
+  void RunMainThreadPendingTasks();
+
   bool multi_thread_;
   bool threads_closed_;
   bool scheduler_stopped_;
@@ -100,6 +107,9 @@ class VirtualMachine final {
 
   std::thread schedule_thread_;
   Notifier pending_notifier_;
+
+  std::mutex main_thread_pending_tasks_mutex_;
+  std::vector<std::function<void()>> main_thread_pending_tasks_;
 };
 
 }  // namespace oneflow