Tmp compute

oneflow/core/autograd/autograd_engine.cpp

@@ -19,8 +19,10 @@ limitations under the License.
 #include <queue>
 #include "oneflow/core/autograd/autograd_engine.h"
 #include "oneflow/core/autograd/autograd_meta.h"
+#include "oneflow/core/framework/stream.h"
 #include "oneflow/core/framework/tensor.h"
 #include "oneflow/core/framework/tensor_arg.h"
+#include "oneflow/core/framework/tensor_methods.h"
 #include "oneflow/core/framework/tensor_tuple.h"
 #include "oneflow/core/framework/tensor_rpc_util.h"
 #include "oneflow/core/autograd/autograd_mode.h"
@@ -29,6 +31,7 @@ limitations under the License.
 #include "oneflow/core/framework/global_param_grad_sync_mode.h"
 #include "oneflow/core/common/container_util.h"
 #include "oneflow/core/profiler/profiler.h"
+#include "oneflow/core/common/env_var/autograd.h"
 
 namespace oneflow {
 namespace one {
@@ -114,6 +117,32 @@ Maybe<void> CheckGlobalTensorsMeta(const TensorTuple& tensor_tuple) {
   return Maybe<void>::Ok();
 }
 
+Maybe<void> TouchInTmpComputeStream(const TensorTuple& inputs) {
+  for (auto input : inputs) {
+    if (input->is_global()) { input = JUST(input->cur_rank_phy_tensor()); }
+    if (input) {
+      Symbol<Device> device = JUST(input->device());
+      auto stream = JUST(Stream::New(device, StreamRole::kTmpCompute));
+      JUST(Touch(input, stream));
+    }
+  }
+  return Maybe<void>::Ok();
+}
+
+constexpr static int kSmallTensorThreshold = 1024;
+
+Maybe<TensorTuple> TryCopyForSmallTensor(const TensorTuple& inputs) {
+  auto outputs = std::make_shared<TensorTuple>();
+  outputs->reserve(inputs.size());
+  for (auto input : inputs) {
+    if (input->shape()->elem_cnt() <= kSmallTensorThreshold) {
+      input = JUST(functional::Identity(input));
+    }
+    outputs->push_back(input);
+  }
+  return outputs;
+}
+
 }  // namespace
 
 Maybe<void> AutogradEngine::RunBackwardAndSaveGrads4LeafTensorIf(const TensorTuple& outputs,
@@ -123,7 +152,16 @@ Maybe<void> AutogradEngine::RunBackwardAndSaveGrads4LeafTensorIf(const TensorTup
   JUST(CheckGlobalTensorsMeta(outputs));
   JUST(CheckGlobalTensorsMeta(out_grads));
   DisableCheckGlobalTensorMetaScope disable_meta_check;
-  return RunBackwardAndSaveGrads4LeafTensor(outputs, out_grads, retain_graph, create_graph);
+  if (ThreadLocalEnvBool<ONEFLOW_AD_PUT_LOSS_ON_TMP_COMPUTE_STREAM>()) {
+    // Put outputs into kTmpCompute stream for reducing blocking time of outputs[i].numpy() in main
+    // thread.
+    auto copied_outputs = JUST(TryCopyForSmallTensor(outputs));
+    JUST(TouchInTmpComputeStream(outputs));
+    return RunBackwardAndSaveGrads4LeafTensor(*copied_outputs, out_grads, retain_graph,
+                                              create_graph);
+  } else {
+    return RunBackwardAndSaveGrads4LeafTensor(outputs, out_grads, retain_graph, create_graph);
+  }
 }
 
 Maybe<TensorTuple> AutogradEngine::RunBackwardAndReturnInputsTensorGradIf(

oneflow/core/common/env_var/autograd.h

@@ -0,0 +1,27 @@
+/*
+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_AUTOGRAD_H_
+#define ONEFLOW_CORE_COMMON_ENV_VAR_AUTOGRAD_H_
+
+#include "oneflow/core/common/env_var/env_var.h"
+
+namespace oneflow {
+
+DEFINE_THREAD_LOCAL_ENV_BOOL(ONEFLOW_AD_PUT_LOSS_ON_TMP_COMPUTE_STREAM, true);
+
+}
+
+#endif  // ONEFLOW_CORE_COMMON_ENV_VAR_AUTOGRAD_H_

oneflow/core/common/stream_role.h

@@ -28,12 +28,14 @@ enum class StreamRole {
   kCompute,
   kHost2Device,
   kDevice2Host,
+  kAsyncedDevice2Host,
   kSyncedLaunchedCommNet,
   kAsyncedLaunchedCommNet,
   kBarrier,
   kCriticalSection,
   kLazyJobLauncher,
-  kPinnedCompute
+  kPinnedCompute,
+  kTmpCompute
 };
 
 template<typename DerivedT>
@@ -45,6 +47,8 @@ struct StreamRoleVisitor {
       case StreamRole::kCompute: return DerivedT::VisitCompute(std::forward<Args>(args)...);
       case StreamRole::kHost2Device: return DerivedT::VisitHost2Device(std::forward<Args>(args)...);
       case StreamRole::kDevice2Host: return DerivedT::VisitDevice2Host(std::forward<Args>(args)...);
+      case StreamRole::kAsyncedDevice2Host:
+        return DerivedT::VisitAsyncedDevice2Host(std::forward<Args>(args)...);
       case StreamRole::kSyncedLaunchedCommNet:
         return DerivedT::VisitSyncedLaunchedCommNet(std::forward<Args>(args)...);
       case StreamRole::kAsyncedLaunchedCommNet:
@@ -56,6 +60,7 @@ struct StreamRoleVisitor {
         return DerivedT::VisitLazyJobLauncher(std::forward<Args>(args)...);
       case StreamRole::kPinnedCompute:
         return DerivedT::VisitPinnedCompute(std::forward<Args>(args)...);
+      case StreamRole::kTmpCompute: return DerivedT::VisitTmpCompute(std::forward<Args>(args)...);
     }
     LOG(FATAL) << "invalid stream role";
   }

oneflow/core/framework/instructions_builder.cpp

@@ -409,12 +409,19 @@ Maybe<void> InstructionsBuilder::ReleaseTensor(
   return Maybe<void>::Ok();
 }
 
-Maybe<void> InstructionsBuilder::TouchTensors(const vm::EagerBlobObjectListPtr& eager_blob_object) {
+Maybe<void> InstructionsBuilder::TouchTensors(
+    const vm::EagerBlobObjectListPtr& eager_blob_objects) {
   Symbol<Device> device = JUST(Device::New("cpu"));
   Symbol<Stream> stream = JUST(GetDefaultStreamByDevice(device));
+  return TouchTensors(eager_blob_objects, stream);
+}
+
+Maybe<void> InstructionsBuilder::TouchTensors(const vm::EagerBlobObjectListPtr& eager_blob_objects,
+                                              Symbol<Stream> stream) {
+  JUST(SoftSyncStream(*eager_blob_objects, stream));
   auto instruction = intrusive::make_shared<vm::Instruction>(
       JUST(Singleton<VirtualMachine>::Get()->GetVmStream(stream)),
-      std::make_shared<vm::TouchTensorsInstructionPolicy>(*eager_blob_object));
+      std::make_unique<vm::TouchTensorsInstructionPolicy>(*eager_blob_objects));
   instruction_list_->EmplaceBack(std::move(instruction));
   return Maybe<void>::Ok();
 }

oneflow/core/framework/instructions_builder.h

@@ -76,7 +76,10 @@ class InstructionsBuilder : public std::enable_shared_from_this<InstructionsBuil
 
   Maybe<void> ReleaseTensor(const std::shared_ptr<vm::EagerBlobObject>& eager_blob_object);
 
-  Maybe<void> TouchTensors(const vm::EagerBlobObjectListPtr& eager_blob_object);
+  Maybe<void> TouchTensors(const vm::EagerBlobObjectListPtr& eager_blob_objects);
+
+  Maybe<void> TouchTensors(const vm::EagerBlobObjectListPtr& eager_blob_objects,
+                           Symbol<Stream> stream);
 
   template<typename T>
   Maybe<void> SyncAccessBlobByCallback(

oneflow/core/framework/stream_allocator_is_pinned.h

@@ -25,12 +25,14 @@ struct IsStreamAllocatorPinned : public StreamRoleVisitor<IsStreamAllocatorPinne
   static bool VisitCompute() { return false; }
   static bool VisitHost2Device() { return false; }
   static bool VisitDevice2Host() { return false; }
+  static bool VisitAsyncedDevice2Host() { return VisitDevice2Host(); }
   static bool VisitSyncedLaunchedCommNet() { return false; }
   static bool VisitAsyncedLaunchedCommNet() { return false; }
   static bool VisitBarrier() { return false; }
   static bool VisitCriticalSection() { return false; }
   static bool VisitLazyJobLauncher() { return false; }
   static bool VisitPinnedCompute() { return true; }
+  static bool VisitTmpCompute() { return false; }
 };
 
 }  // namespace oneflow

oneflow/core/framework/stream_get_stream_role_name.h

@@ -28,12 +28,14 @@ struct GetStreamRoleName : public StreamRoleVisitor<GetStreamRoleName> {
   static const char* VisitCompute() { return "compute"; }
   static const char* VisitHost2Device() { return "h2d"; }
   static const char* VisitDevice2Host() { return "d2h"; }
+  static const char* VisitAsyncedDevice2Host() { return "asynced_d2h"; }
   static const char* VisitSyncedLaunchedCommNet() { return "synced_launched_comm_net"; }
   static const char* VisitAsyncedLaunchedCommNet() { return "asynced_launched_comm_net"; }
   static const char* VisitBarrier() { return "barrier"; }
   static const char* VisitCriticalSection() { return "critical_section"; }
   static const char* VisitLazyJobLauncher() { return "lazy_job_launcher"; }
-  static const char* VisitPinnedCompute() { return "pin_memory"; }
+  static const char* VisitPinnedCompute() { return "pinned_compute"; }
+  static const char* VisitTmpCompute() { return "tmp_compute"; }
 };
 
 }  // namespace oneflow

oneflow/core/framework/stream_is_comm_net_stream.h

@@ -25,12 +25,14 @@ struct IsCommNetStream final : public StreamRoleVisitor<IsCommNetStream> {
   static bool VisitCompute() { return false; }
   static bool VisitHost2Device() { return false; }
   static bool VisitDevice2Host() { return false; }
+  static bool VisitAsyncedDevice2Host() { return VisitDevice2Host(); }
   static bool VisitSyncedLaunchedCommNet() { return true; }
   static bool VisitAsyncedLaunchedCommNet() { return true; }
   static bool VisitBarrier() { return false; }
   static bool VisitCriticalSection() { return false; }
   static bool VisitLazyJobLauncher() { return false; }
   static bool VisitPinnedCompute() { return VisitCompute(); }
+  static bool VisitTmpCompute() { return VisitCompute(); }
 };
 
 }  // namespace oneflow

oneflow/core/framework/stream_need_soft_sync.h

@@ -26,12 +26,16 @@ struct NeedSoftSync : public StreamRoleVisitor<NeedSoftSync> {
   static bool VisitCompute(DeviceType device_type) { return device_type != kCPU; }
   static bool VisitHost2Device(DeviceType) { return false; }
   static bool VisitDevice2Host(DeviceType) { return false; }
+  static bool VisitAsyncedDevice2Host(DeviceType device_type) {
+    return VisitDevice2Host(device_type);
+  }
   static bool VisitSyncedLaunchedCommNet(DeviceType device_type) { return false; }
   static bool VisitAsyncedLaunchedCommNet(DeviceType) { return false; }
   static bool VisitBarrier(DeviceType) { return false; }
   static bool VisitCriticalSection(DeviceType) { return false; }
   static bool VisitLazyJobLauncher(DeviceType) { return false; }
   static bool VisitPinnedCompute(DeviceType device_type) { return VisitCompute(device_type); }
+  static bool VisitTmpCompute(DeviceType device_type) { return false; }
 };
 
 }  // namespace oneflow

oneflow/core/framework/stream_on_independent_thread.h

@@ -25,12 +25,14 @@ struct StreamOnIndependentThread : public StreamRoleVisitor<StreamOnIndependentT
   static bool VisitCompute() { return false; }
   static bool VisitHost2Device() { return false; }
   static bool VisitDevice2Host() { return false; }
+  static bool VisitAsyncedDevice2Host() { return true; }
   static bool VisitSyncedLaunchedCommNet() { return false; }
   static bool VisitAsyncedLaunchedCommNet() { return false; }
   static bool VisitBarrier() { return false; }
   static bool VisitCriticalSection() { return true; }
   static bool VisitLazyJobLauncher() { return true; }
   static bool VisitPinnedCompute() { return VisitCompute(); }
+  static bool VisitTmpCompute() { return VisitCompute(); }
 };
 
 }  // namespace oneflow

oneflow/core/framework/tensor_methods.cpp

@@ -546,5 +546,16 @@ Maybe<Tensor> Diagonal(const std::shared_ptr<Tensor>& input, const int32_t offse
 }
 
 }  // namespace view
+
+Maybe<void> Touch(std::shared_ptr<Tensor> input, Symbol<Stream> stream) {
+  auto eager_blob_objects = std::make_shared<vm::EagerBlobObjectList>();
+  if (input->is_global()) { input = JUST(input->cur_rank_phy_tensor()); }
+  if (input) { eager_blob_objects->push_back(JUST(input->eager_blob_object())); }
+  JUST(PhysicalRun([&](InstructionsBuilder* builder) -> Maybe<void> {
+    return builder->TouchTensors(eager_blob_objects, stream);
+  }));
+  return Maybe<void>::Ok();
+}
+
 }  // namespace one
 }  // namespace oneflow

oneflow/core/framework/tensor_methods.h

@@ -20,6 +20,9 @@ limitations under the License.
 #include "oneflow/core/framework/tensor.h"
 
 namespace oneflow {
+
+class Stream;
+
 namespace one {
 
 class Tensor;
@@ -67,7 +70,11 @@ Maybe<Tensor> Diagonal(const std::shared_ptr<Tensor>& input, const int32_t offse
                        const int32_t dim1, const int32_t dim2);
 
 }  // namespace view
+
+Maybe<void> Touch(std::shared_ptr<Tensor> input, Symbol<Stream> stream);
+
 }  // namespace one
+
 }  // namespace oneflow
 
 #endif  // ONEFLOW_CORE_FRAMEWORK_TENSOR_METHOD_H_

oneflow/core/functional/impl/array_functor.cpp

@@ -1478,13 +1478,23 @@ class CopyFunctor {
     JUST(attrs.SetAttr<std::string>("device_type", device_type));
     JUST(attrs.SetAttr<int64_t>("device_id", device_id));
     JUST(attrs.SetAttr<bool>("pin_memory", pin_memory));
+    JUST(attrs.SetAttr<bool>("asynced_copy", JUST(GetAsyncedCopy(*x))));
 
 #ifdef WITH_CUDA
     if (device_type == "cuda") { InitCudaContextOnce(device_id); }
 #endif
     return OpInterpUtil::Dispatch<Tensor>(*op_, {x}, attrs);
   }
 
+  Maybe<bool> GetAsyncedCopy(const one::Tensor& x) const {
+    if (!x.is_eager()) { return false; }
+    if (!x.is_local()) { return false; }
+    const auto& eager_blob_object = JUST(x.eager_blob_object());
+    const auto& opt_stream = eager_blob_object->last_used_stream();
+    if (!opt_stream.has_value()) { return false; }
+    return JUST(opt_stream)->stream_role() == StreamRole::kTmpCompute;
+  }
+
  private:
   std::shared_ptr<OpExpr> op_;
 };

oneflow/core/vm/ep_record_event_instruction_policy.h

@@ -110,6 +110,11 @@ struct GetRecordEventInstructionPolicy : public StreamRoleVisitor<GetRecordEvent
         new vm::EpRecordEventInstructionPolicy(std::forward<Args>(args)...));
   }
   template<typename... Args>
+  static Maybe<vm::InstructionPolicy> VisitAsyncedDevice2Host(DeviceType device_type,
+                                                              Args&&... args) {
+    return VisitDevice2Host(device_type, std::forward<Args>(args)...);
+  }
+  template<typename... Args>
   static Maybe<vm::InstructionPolicy> VisitSyncedLaunchedCommNet(DeviceType device_type,
                                                                  Args&&... args) {
     return std::shared_ptr<vm::InstructionPolicy>(
@@ -140,6 +145,11 @@ struct GetRecordEventInstructionPolicy : public StreamRoleVisitor<GetRecordEvent
     return std::shared_ptr<vm::InstructionPolicy>(
         new vm::EpRecordEventInstructionPolicy(std::forward<Args>(args)...));
   }
+  template<typename... Args>
+  static Maybe<vm::InstructionPolicy> VisitTmpCompute(DeviceType device_type, Args&&... args) {
+    return std::shared_ptr<vm::InstructionPolicy>(
+        new vm::EpRecordEventInstructionPolicy(std::forward<Args>(args)...));
+  }
 };
 
 }  // namespace oneflow

oneflow/core/vm/release_tensor_instruction_policy.h

@@ -145,6 +145,11 @@ struct MakeReleaseTensorInstructionPolicy
       const Optional<vm::Stream*>& stream) {
     return Make(data_type, eager_blob_object, stream);
   }
+  static Maybe<vm::InstructionPolicy> VisitAsyncedDevice2Host(
+      DataType data_type, const std::shared_ptr<vm::EagerBlobObject>& eager_blob_object,
+      const Optional<vm::Stream*>& stream) {
+    return VisitDevice2Host(data_type, eager_blob_object, stream);
+  }
   static Maybe<vm::InstructionPolicy> VisitSyncedLaunchedCommNet(
       DataType data_type, const std::shared_ptr<vm::EagerBlobObject>& eager_blob_object,
       const Optional<vm::Stream*>& stream) {
@@ -178,6 +183,12 @@ struct MakeReleaseTensorInstructionPolicy
     return VisitCompute(data_type, eager_blob_object, stream);
   }
 
+  static Maybe<vm::InstructionPolicy> VisitTmpCompute(
+      DataType data_type, const std::shared_ptr<vm::EagerBlobObject>& eager_blob_object,
+      const Optional<vm::Stream*>& stream) {
+    return VisitCompute(data_type, eager_blob_object, stream);
+  }
+
  private:
   static Maybe<vm::InstructionPolicy> Make(
       DataType data_type, const std::shared_ptr<vm::EagerBlobObject>& eager_blob_object,

oneflow/core/vm/stream_get_stream_policy.h

@@ -40,6 +40,9 @@ struct CreateStreamPolicy final : public StreamRoleVisitor<CreateStreamPolicy> {
   static Maybe<vm::StreamPolicy> VisitDevice2Host(Symbol<Device> device) {
     return std::shared_ptr<vm::StreamPolicy>(new vm::EpD2HStreamPolicy(device));
   }
+  static Maybe<vm::StreamPolicy> VisitAsyncedDevice2Host(Symbol<Device> device) {
+    return VisitDevice2Host(device);
+  }
   static Maybe<vm::StreamPolicy> VisitSyncedLaunchedCommNet(Symbol<Device> device) {
     return std::shared_ptr<vm::StreamPolicy>(new vm::EventRecordedEpStreamPolicy(device));
   }
@@ -58,6 +61,9 @@ struct CreateStreamPolicy final : public StreamRoleVisitor<CreateStreamPolicy> {
   static Maybe<vm::StreamPolicy> VisitPinnedCompute(Symbol<Device> device) {
     return std::shared_ptr<vm::StreamPolicy>(new vm::PinnedEpStreamPolicy(device));
   }
+  static Maybe<vm::StreamPolicy> VisitTmpCompute(Symbol<Device> device) {
+    return std::shared_ptr<vm::StreamPolicy>(new vm::EventRecordedEpStreamPolicy(device));
+  }
 };
 
 }  // namespace oneflow

oneflow/ir/include/OneFlow/OneFlowUserOps.td

@@ -8601,7 +8601,8 @@ def OneFlow_CopyOp : OneFlow_BaseOp<"copy", [NoSideEffect, DeclareOpInterfaceMet
   let attrs = (ins
     StrAttr:$device_type,
     DefaultValuedAttr<SI64Attr, "0">:$device_id,
-    DefaultValuedAttr<BoolAttr, "false">:$pin_memory
+    DefaultValuedAttr<BoolAttr, "false">:$pin_memory,
+    DefaultValuedAttr<BoolAttr, "false">:$asynced_copy
   );
   let has_logical_tensor_desc_infer_fn = 1;
   let has_physical_tensor_desc_infer_fn = 1;