Sync vm mode guard

oneflow/core/common/thread_local_guard.h

@@ -18,42 +18,28 @@ limitations under the License.
 
 #include <memory>
 #include <glog/logging.h>
+#include "oneflow/core/common/optional.h"
 
 namespace oneflow {
 
-// Interfaces:
-//   - ThreadLocalGuard::CurrentValue()
-//   - ThreadLocalGuard::HasCurrentValue()
 template<typename T>
-class ThreadLocalGuard;
-
-template<>
-class ThreadLocalGuard<bool> {
+class ThreadLocalGuard {
  public:
-  explicit ThreadLocalGuard(bool value) {
+  explicit ThreadLocalGuard(const T& value) {
     old_value_ = *MutThreadLocalValue();
-    *MutThreadLocalValue() = int(value);
+    *MutThreadLocalValue() = Optional<T>(value);
   }
   ~ThreadLocalGuard() { *MutThreadLocalValue() = old_value_; }
 
-  static bool CurrentValue() {
-    int value = *MutThreadLocalValue();
-    CHECK_GE(value, 0);
-    return value > 0;
-  }
-
-  static bool HasCurrentValue() { return *MutThreadLocalValue() >= 0; }
+  static const Optional<T>& Current() { return *MutThreadLocalValue(); }
 
  private:
-  static int* MutThreadLocalValue() {
-    static thread_local int value = -1;
+  static Optional<T>* MutThreadLocalValue() {
+    static thread_local Optional<T> value{};
     return &value;
   }
 
-  // -1: not exists.
-  // 0: false.
-  // 1: true.
-  int old_value_;
+  Optional<T> old_value_;
 };
 
 }  // namespace oneflow

oneflow/core/common/thread_local_guard_test.cpp

@@ -20,40 +20,42 @@ limitations under the License.
 namespace oneflow {
 namespace test {
 
-template<typename T>
-void AssertCurrentValue(const T& value) {
-  ThreadLocalGuard<T> guard(value);
-  ASSERT_TRUE(ThreadLocalGuard<T>::HasCurrentValue());
-  ASSERT_EQ(ThreadLocalGuard<T>::CurrentValue(), value);
-}
-
 template<typename T>
 void Assert(const T& value0, const T& value1) {
-  ASSERT_FALSE(ThreadLocalGuard<T>::HasCurrentValue());
+  ASSERT_FALSE(ThreadLocalGuard<T>::Current().has_value());
   {
     ThreadLocalGuard<T> guard(value0);
-    ASSERT_TRUE(ThreadLocalGuard<T>::HasCurrentValue());
+    ASSERT_TRUE(ThreadLocalGuard<T>::Current().has_value());
   }
   {
     ThreadLocalGuard<T> guard(value0);
-    ASSERT_TRUE(ThreadLocalGuard<T>::HasCurrentValue());
-    ASSERT_EQ(ThreadLocalGuard<T>::CurrentValue(), value0);
+    ASSERT_TRUE(ThreadLocalGuard<T>::Current().has_value());
+    T value = CHECK_JUST(ThreadLocalGuard<T>::Current());
+    ASSERT_EQ(value, value0);
   }
   {
     ThreadLocalGuard<T> guard(value1);
-    ASSERT_TRUE(ThreadLocalGuard<T>::HasCurrentValue());
-    ASSERT_EQ(ThreadLocalGuard<T>::CurrentValue(), value1);
+    ASSERT_TRUE(ThreadLocalGuard<T>::Current().has_value());
+    const auto& value = CHECK_JUST(ThreadLocalGuard<T>::Current());
+    ASSERT_EQ(value, value1);
   }
   {
     ThreadLocalGuard<T> guard(value0);
-    ASSERT_TRUE(ThreadLocalGuard<T>::HasCurrentValue());
-    ASSERT_EQ(ThreadLocalGuard<T>::CurrentValue(), value0);
+    ASSERT_TRUE(ThreadLocalGuard<T>::Current().has_value());
+    {
+      const auto& value = CHECK_JUST(ThreadLocalGuard<T>::Current());
+      ASSERT_EQ(value, value0);
+    }
     {
       ThreadLocalGuard<T> nested_guard(value1);
-      ASSERT_TRUE(ThreadLocalGuard<T>::HasCurrentValue());
-      ASSERT_EQ(ThreadLocalGuard<T>::CurrentValue(), value1);
+      ASSERT_TRUE(ThreadLocalGuard<T>::Current().has_value());
+      const auto& value = CHECK_JUST(ThreadLocalGuard<T>::Current());
+      ASSERT_EQ(value, value1);
+    }
+    {
+      const auto& value = CHECK_JUST(ThreadLocalGuard<T>::Current());
+      ASSERT_EQ(value, value0);
     }
-    ASSERT_EQ(ThreadLocalGuard<T>::CurrentValue(), value0);
   }
 }
 

oneflow/core/platform/lib/pthread_fork.cpp

@@ -17,6 +17,7 @@ limitations under the License.
 #include "oneflow/core/common/util.h"
 #include "oneflow/core/vm/virtual_machine.h"
 #include "oneflow/core/vm/vm_util.h"
+#include "oneflow/core/vm/sync_vm_mode_guard.h"
 
 namespace oneflow {
 
@@ -29,6 +30,7 @@ static void SetIsForkedSubProcess() { is_fork = true; }
 
 namespace {
 void CurrentRankVmSync() {
+  if (SyncVmModeGuard::IsCurrentSyncVmMode()) { return; }
   // Instructions in forked subprocesses are not dispatched to vm,
   // so no need to sync vm in these processes.
   if (!is_fork && Singleton<VirtualMachine>::Get() != nullptr) {

oneflow/core/thread/thread_pool.cpp

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 #include "oneflow/core/thread/thread_pool.h"
+#include "oneflow/core/vm/sync_vm_mode_guard.h"
 
 namespace oneflow {
 
@@ -22,6 +23,7 @@ ThreadPool::ThreadPool(int32_t thread_num)
   FOR_RANGE(int32_t, i, 0, thread_num) {
     Channel<std::function<void()>>* chan = &(work_chans_.at(i));
     threads_[i] = std::thread([chan]() {
+      SyncVmModeGuard guard(SyncVmMode::kEnable);
       std::function<void()> work;
       while (chan->Receive(&work) == kChannelStatusSuccess) { work(); }
     });

oneflow/core/vm/sync_vm_mode_guard.h

@@ -0,0 +1,42 @@
+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+#ifndef ONEFLOW_CORE_VM_SYNC_VM_MODE_GUARD_H_
+#define ONEFLOW_CORE_VM_SYNC_VM_MODE_GUARD_H_
+
+#include "oneflow/core/common/thread_local_guard.h"
+
+namespace oneflow {
+
+enum class SyncVmMode {
+  kInvalid = 0,
+  kEnable = 1,
+  kDisable = 2,
+};
+
+class SyncVmModeGuard final : public ThreadLocalGuard<SyncVmMode> {
+ public:
+  using ThreadLocalGuard<SyncVmMode>::ThreadLocalGuard;
+  ~SyncVmModeGuard() = default;
+
+  static bool IsCurrentSyncVmMode() {
+    const auto& opt_sync_mode = Current();
+    return opt_sync_mode.has_value() && CHECK_JUST(opt_sync_mode) == SyncVmMode::kEnable;
+  }
+};
+
+}  // namespace oneflow
+
+#endif  // ONEFLOW_CORE_VM_SYNC_VM_MODE_GUARD_H_

oneflow/core/vm/virtual_machine.cpp

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 #include <typeinfo>
+#include "oneflow/core/vm/sync_vm_mode_guard.h"
 #include "oneflow/core/vm/barrier_instruction_policy.h"
 #include "oneflow/core/vm/caching_allocator.h"
 #include "oneflow/core/vm/global_sync_instruction_policy.h"
@@ -67,6 +68,7 @@ void GetSchedulerThreadInitializer(std::function<void()>* Initializer) {
 }
 
 void WorkerLoop(vm::ThreadCtx* thread_ctx, const std::function<void(vm::ThreadCtx*)>& Initializer) {
+  SyncVmModeGuard guard(SyncVmMode::kEnable);
   Initializer(thread_ctx);
   while (thread_ctx->mut_notifier()->WaitAndClearNotifiedCnt() == kNotifierStatusSuccess) {
     while (thread_ctx->TryReceiveAndRun()) {}
@@ -220,6 +222,7 @@ std::string VirtualMachine::GetBlockingDebugString() {
 }
 
 Maybe<void> VirtualMachine::Receive(vm::InstructionList* instruction_list) {
+  SyncVmModeGuard guard(SyncVmMode::kEnable);
   if (unlikely(pthread_fork::IsForkedSubProcess())) {
     INTRUSIVE_FOR_EACH_PTR(instruction, instruction_list) {
       const auto& device = instruction->stream().device();
@@ -286,6 +289,7 @@ class MultiThreadScheduleCtx : public vm::ScheduleCtx {
 }  // namespace
 
 void VirtualMachine::ScheduleLoop(const std::function<void()>& Initializer) {
+  SyncVmModeGuard guard(SyncVmMode::kEnable);
   Initializer();
   MultiThreadScheduleCtx schedule_ctx{};
   while (pending_notifier_.WaitAndClearNotifiedCnt() == kNotifierStatusSuccess) {