RankTaskGraph

CMakeLists.txt

@@ -161,6 +161,10 @@ if(BUILD_PROFILER)
   add_definitions(-DOF_ENABLE_PROFILER)
 endif()
 
+if(BUILD_DEBUG_LAZY_RUNTIME)
+  add_definitions(-DOF_DEBUG_LAZY_RUNTIME)
+endif()
+
 if(OF_SOFTMAX_USE_FAST_MATH)
   add_definitions(-DOF_SOFTMAX_USE_FAST_MATH)
 endif()

oneflow/core/common/async_deallocate_context.cpp

@@ -0,0 +1,29 @@
+/*
+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.
+*/
+#include "oneflow/core/common/async_deallocate_context.h"
+#include "oneflow/core/thread/thread_pool.h"
+
+namespace oneflow {
+
+AsyncDeallocateContext::AsyncDeallocateContext() : thread_pool_(std::make_unique<ThreadPool>(1)) {}
+
+AsyncDeallocateContext::~AsyncDeallocateContext() {}
+
+void AsyncDeallocateContext::LazyDeallocate(std::function<void()> LazyDeallocator) {
+  thread_pool_->AddWork(LazyDeallocator);
+}
+
+}  // namespace oneflow

oneflow/core/common/async_deallocate_context.h

@@ -0,0 +1,38 @@
+/*
+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_COMMON_ASYNC_DEALLOCATE_CONTEXT_H_
+#define ONEFLOW_CORE_COMMON_ASYNC_DEALLOCATE_CONTEXT_H_
+
+#include "oneflow/core/common/deallocate_context.h"
+
+namespace oneflow {
+
+class ThreadPool;
+
+class AsyncDeallocateContext final : public DeallocateContext {
+ public:
+  AsyncDeallocateContext();
+  ~AsyncDeallocateContext() override;
+
+  void LazyDeallocate(std::function<void()> LazyDeallocator) override;
+
+ private:
+  std::shared_ptr<ThreadPool> thread_pool_;
+};
+
+}  // namespace oneflow
+
+#endif  // ONEFLOW_CORE_COMMON_ASYNC_DEALLOCATE_CONTEXT_H_

oneflow/core/common/balanced_splitter.cpp

@@ -21,8 +21,11 @@ BalancedSplitter::BalancedSplitter(int64_t total_num, int64_t split_num) {
   base_part_size_ = total_num / split_num;
   base_begin_idx_ = total_num % split_num;
   split_num_ = split_num;
+  CHECK_EQ(this->total_num(), total_num);
 }
 
+int64_t BalancedSplitter::total_num() const { return At(split_num_ - 1).end(); }
+
 Range BalancedSplitter::At(int64_t idx) const {
   CHECK_LT(idx, split_num_);
   int64_t left_bound = -1;
@@ -46,4 +49,27 @@ Range BalancedSplitter::At(int64_t first_idx, int64_t last_idx) const {
   return Range(first_range.begin(), last_range.end());
 }
 
+int64_t BalancedSplitter::RecursiveBinarySearchIndex(int64_t value) const {
+  CHECK_GE(value, 0);
+  CHECK_LT(value, total_num());
+  return RecursiveBinarySearchIndex(value, 0, split_num_);
+}
+
+int64_t BalancedSplitter::RecursiveBinarySearchIndex(int64_t value, int begin_idx,
+                                                     int end_idx) const {
+  CHECK_LT(begin_idx, end_idx);
+  if (begin_idx + 1 == end_idx) {
+    Range ret = At(begin_idx);
+    CHECK_GE(value, ret.begin());
+    CHECK_LT(value, ret.end());
+    return begin_idx;
+  }
+  int middle_idx = (begin_idx + end_idx) / 2;
+  if (value >= At(middle_idx).begin()) {
+    return RecursiveBinarySearchIndex(value, middle_idx, end_idx);
+  } else {
+    return RecursiveBinarySearchIndex(value, begin_idx, middle_idx);
+  }
+}
+
 }  // namespace oneflow

oneflow/core/common/balanced_splitter.h

@@ -42,7 +42,13 @@ class BalancedSplitter final {
   Range At(int64_t idx) const;
   Range At(int64_t first_idx, int64_t last_idx) const;
 
+  int64_t RecursiveBinarySearchIndex(int64_t value) const;
+  int64_t total_num() const;
+
  private:
+  // search in [begin_idx, end_idx)
+  int64_t RecursiveBinarySearchIndex(int64_t value, int begin_idx, int end_idx) const;
+
   int64_t base_part_size_;
   int64_t base_begin_idx_;
   int64_t split_num_;

oneflow/core/common/balanced_splitter_test.cpp

@@ -35,4 +35,17 @@ TEST(BalancedSplitter, split_2_to_3_part) {
   ASSERT_TRUE(splitter.At(2) == Range(2, 2));
 }
 
+TEST(BalancedSplitter, RecursiveBinarySearchIndex) {
+  const size_t total_num = 937;
+  const size_t split_num = 11;
+  BalancedSplitter bs(total_num, split_num);
+  ASSERT_TRUE(bs.total_num() == total_num);
+  for (size_t i = 0; i < split_num; ++i) {
+    Range range = bs.At(i);
+    for (size_t value = range.begin(); value < range.end(); ++value) {
+      ASSERT_TRUE(bs.RecursiveBinarySearchIndex(value) == i);
+    }
+  }
+}
+
 }  // namespace oneflow

oneflow/core/common/deallocate_context.h

@@ -0,0 +1,50 @@
+/*
+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_COMMON_DEALLOCATE_CONTEXT_H_
+#define ONEFLOW_CORE_COMMON_DEALLOCATE_CONTEXT_H_
+
+#include <functional>
+#include <memory>
+
+namespace oneflow {
+
+class DeallocateContext {
+ public:
+  DeallocateContext() = default;
+  virtual ~DeallocateContext() = default;
+
+  template<typename T>
+  void Deallocate(std::shared_ptr<T>&& ptr) {
+    std::shared_ptr<T> data = ptr;
+    ptr.reset();
+    LazyDeallocate([data] { const_cast<std::shared_ptr<T>*>(&data)->reset(); });
+    data.reset();
+  }
+
+  virtual void LazyDeallocate(std::function<void()> LazyDeallocator) = 0;
+};
+
+class NaiveDeallocateContext final : public DeallocateContext {
+ public:
+  NaiveDeallocateContext() = default;
+  ~NaiveDeallocateContext() = default;
+
+  void LazyDeallocate(std::function<void()> LazyDeallocator) { LazyDeallocator(); }
+};
+
+}  // namespace oneflow
+
+#endif  // ONEFLOW_CORE_COMMON_DEALLOCATE_CONTEXT_H_

oneflow/core/common/env_var/env_var.h

@@ -70,6 +70,17 @@ int64_t ThreadLocalEnvInteger();
 
 DEFINE_THREAD_LOCAL_ENV_INTEGER(ONEFLOW_THRAED_LOCAL_CACHED_SIZE, 128 * 1024);
 
+template<typename env_var>
+const std::string& ThreadLocalEnvString();
+
+#define DEFINE_THREAD_LOCAL_ENV_STRING(env_var, default_value)                                  \
+  struct env_var {};                                                                            \
+  template<>                                                                                    \
+  inline const std::string& ThreadLocalEnvString<env_var>() {                                   \
+    thread_local std::string value = GetStringFromEnv(OF_PP_STRINGIZE(env_var), default_value); \
+    return value;                                                                               \
+  }
+
 }  // namespace oneflow
 
 #endif  // ONEFLOW_CORE_COMMON_ENV_VAR_ENV_VAR_H_

oneflow/core/common/env_var/lazy.h

@@ -0,0 +1,29 @@
+/*
+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_COMMON_ENV_VAR_LAZY_H_
+#define ONEFLOW_CORE_COMMON_ENV_VAR_LAZY_H_
+
+#include <string>
+#include "oneflow/core/common/env_var/env_var.h"
+
+namespace oneflow {
+
+DEFINE_THREAD_LOCAL_ENV_STRING(ONEFLOW_LAZY_COMPILE_MODE, "naive");
+DEFINE_THREAD_LOCAL_ENV_INTEGER(ONEFLOW_LAZY_COMPILE_RPC_THREAD_NUM, 16);
+
+}  // namespace oneflow
+
+#endif  // ONEFLOW_CORE_COMMON_ENV_VAR_LAZY_H_

oneflow/core/common/id_pairs.h

@@ -0,0 +1,41 @@
+/*
+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_COMMON_ID_PAIRS_H_
+#define ONEFLOW_CORE_COMMON_ID_PAIRS_H_
+
+#include "oneflow/core/common/id_pairs.pb.h"
+#include <unordered_set>
+#include <utility>
+
+namespace oneflow {
+
+void InitIdPairs(const std::unordered_set<std::pair<int64_t, int64_t>>& pairs, IdPairs* proto) {
+  for (const auto& pair : pairs) {
+    auto* proto_pair = proto->mutable_int64_pair()->Add();
+    proto_pair->set_first(pair.first);
+    proto_pair->set_second(pair.second);
+  }
+}
+
+void MergeIdPairs(const IdPairs& id_pairs, std::unordered_set<std::pair<int64_t, int64_t>>* pairs) {
+  for (const auto& pair : id_pairs.int64_pair()) {
+    pairs->emplace(std::make_pair(pair.first(), pair.second()));
+  }
+}
+
+}  // namespace oneflow
+
+#endif  // ONEFLOW_CORE_COMMON_ID_PAIRS_H_

oneflow/core/common/id_pairs.proto

@@ -0,0 +1,11 @@
+syntax = "proto2";
+package oneflow;
+
+message Int64Pair {
+  required int64 first = 1;
+  required int64 second = 2;
+}
+
+message IdPairs {
+  repeated Int64Pair int64_pair = 1;
+}

oneflow/core/framework/instructions_builder.cpp

@@ -207,7 +207,7 @@ Maybe<void> InstructionsBuilder::SoftSyncNNGraphBuffers(
 namespace {
 
 int64_t NewSymbolId() {
-  static std::atomic<int64_t> cnt(0);
+  static std::atomic<int64_t> cnt(1);
   return cnt.fetch_add(1, std::memory_order_relaxed);
 }
 
@@ -225,10 +225,6 @@ Maybe<Scope> InstructionsBuilder::GetScopeSymbol(const ScopeProto& scope_proto)
   return Singleton<symbol::Storage<Scope>>::Get()->FindOrCreate(scope_proto, &NewSymbolId);
 }
 
-Maybe<OperatorConfSymbol> InstructionsBuilder::GetOpConfSymbol(const OperatorConf& op_conf) {
-  return Singleton<symbol::Storage<OperatorConfSymbol>>::Get()->FindOrCreate(op_conf, &NewSymbolId);
-}
-
 Maybe<Scope> InstructionsBuilder::BuildInitialScope(
     int64_t session_id, const JobConfigProto& job_conf, const std::string& device_tag,
     const std::vector<std::string>& machine_device_ids, const std::shared_ptr<Shape>& hierarchy,