Add GenericStreamContext (#8560)

oneflow/core/lazy/actor/actor.cpp

@@ -16,7 +16,7 @@ limitations under the License.
 #include "oneflow/core/lazy/actor/actor.h"
 #include "oneflow/core/control/global_process_ctx.h"
 #include "oneflow/core/job/runtime_job_descs.h"
-#include "oneflow/core/stream/include/stream_context.h"
+#include "oneflow/core/lazy/stream_context/include/stream_context.h"
 
 namespace oneflow {
 

oneflow/core/lazy/actor/actor_context.h

@@ -16,7 +16,7 @@ limitations under the License.
 #ifndef ONEFLOW_CORE_LAZY_ACTOR_ACTOR_CONTEXT_H_
 #define ONEFLOW_CORE_LAZY_ACTOR_ACTOR_CONTEXT_H_
 
-#include "oneflow/core/stream/include/stream_context.h"
+#include "oneflow/core/lazy/stream_context/include/stream_context.h"
 #include "oneflow/core/job/task.pb.h"
 
 namespace oneflow {

oneflow/core/lazy/actor/light_actor.cpp

@@ -27,7 +27,7 @@ limitations under the License.
 #include "oneflow/core/job/runtime_job_descs.h"
 #include "oneflow/core/common/util.h"
 #include "oneflow/core/kernel/user_kernel.h"
-#include "oneflow/core/stream/include/stream_context.h"
+#include "oneflow/core/lazy/stream_context/include/stream_context.h"
 
 #ifdef WITH_CUDA
 

oneflow/core/lazy/stream_context/common/generic_stream_context.cpp

@@ -0,0 +1,59 @@
+/*
+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/lazy/stream_context/include/generic_stream_context.h"
+#include "oneflow/core/job/global_for.h"
+#include "oneflow/core/ep/include/device_manager_registry.h"
+#include "oneflow/core/ep/include/active_device_guard.h"
+
+namespace oneflow {
+
+GenericStreamContext::GenericStreamContext(const StreamId& stream_id) : stream_(nullptr) {
+  device_ =
+      std::dynamic_pointer_cast<ep::Device>(Singleton<ep::DeviceManagerRegistry>::Get()->GetDevice(
+          stream_id.device_type(), stream_id.device_index()));
+  CHECK(device_);
+  ep::ActiveDeviceGuard guard(device_.get());
+  stream_ = dynamic_cast<ep::Stream*>(device_->CreateStream());
+  CHECK(stream_ != nullptr);
+  poller_thread_ = std::thread([this]() {
+    CHECK_JUST(stream_->OnExecutionContextSetup());
+    std::pair<ep::Event*, std::function<void()>> cb_event;
+    while (cb_event_chan_.Receive(&cb_event) == kChannelStatusSuccess) {
+      CHECK_JUST(cb_event.first->Sync());
+      cb_event.second();
+      device_->DestroyEvent(cb_event.first);
+    }
+    CHECK_JUST(stream_->OnExecutionContextTeardown());
+  });
+}
+
+GenericStreamContext::~GenericStreamContext() {
+  ep::ActiveDeviceGuard guard(device_.get());
+  cb_event_chan_.Close();
+  poller_thread_.join();
+  device_->DestroyStream(stream_);
+}
+
+Maybe<void> GenericStreamContext::AddCallback(std::function<void()> callback) {
+  ep::Event* event = device_->CreateEvent();
+  stream_->RecordEvent(event);
+  cb_event_chan_.Send(std::make_pair(event, std::move(callback)));
+  return Maybe<void>::Ok();
+}
+
+ep::Stream* GenericStreamContext::stream() { return stream_; }
+
+}  // namespace oneflow

oneflow/core/stream/cpu/cpu_stream_context.cpp → oneflow/core/lazy/stream_context/cpu/cpu_stream_context.cpp

@@ -13,7 +13,7 @@ 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/stream/include/stream_context.h"
+#include "oneflow/core/lazy/stream_context/include/stream_context.h"
 #include "oneflow/core/common/maybe.h"
 #include "oneflow/core/common/device_type.h"
 #include "oneflow/core/device/event_record.h"
@@ -30,7 +30,7 @@ class CpuStreamContext : public StreamContext, public KernelObserverProvider {
  public:
   OF_DISALLOW_COPY_AND_MOVE(CpuStreamContext);
   CpuStreamContext();
-  virtual ~CpuStreamContext();
+  ~CpuStreamContext() override;
 
   ep::Stream* stream() override;
   Maybe<void> AddCallback(std::function<void()> callback) override;

oneflow/core/stream/cuda/cuda_stream_context.cpp → oneflow/core/lazy/stream_context/cuda/cuda_stream_context.cpp

@@ -13,7 +13,7 @@ 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/stream/include/stream_context.h"
+#include "oneflow/core/lazy/stream_context/include/stream_context.h"
 #include "oneflow/core/profiler/profiler.h"
 #include "oneflow/core/job/global_for.h"
 #include "oneflow/core/common/device_type.h"
@@ -36,7 +36,7 @@ class CudaStreamContext : public StreamContext, public KernelObserverProvider {
  public:
   OF_DISALLOW_COPY_AND_MOVE(CudaStreamContext);
   explicit CudaStreamContext(int device_index);
-  virtual ~CudaStreamContext();
+  ~CudaStreamContext() override;
 
   Maybe<void> AddCallback(std::function<void()> callback) override;
   DeviceType device_type() const override { return DeviceType::kCUDA; }
@@ -53,8 +53,6 @@ class CudaStreamContext : public StreamContext, public KernelObserverProvider {
   std::shared_ptr<ep::CudaDevice> device_;
 };
 
-}  // namespace
-
 CudaStreamContext::CudaStreamContext(int device_index)
     : stream_(nullptr), device_index_(device_index) {
   CudaCurrentDeviceGuard guard(device_index_);
@@ -74,16 +72,16 @@ CudaStreamContext::CudaStreamContext(int device_index)
   kernel_observer_.reset(new ChainKernelObserver(kernel_observers));
 
   poller_thread_ = std::thread([this]() {
-    stream_->OnExecutionContextSetup();
+    CHECK_JUST(stream_->OnExecutionContextSetup());
     OF_PROFILER_NAME_THIS_HOST_THREAD("_cuda" + std::to_string(device_index_) + " Poller : ("
                                       + std::to_string(device_index_) + ")");
     std::pair<ep::Event*, std::function<void()>> cb_event;
     while (cb_event_chan_.Receive(&cb_event) == kChannelStatusSuccess) {
-      cb_event.first->Sync();
+      CHECK_JUST(cb_event.first->Sync());
       cb_event.second();
       device_->DestroyEvent(cb_event.first);
     }
-    stream_->OnExecutionContextTeardown();
+    CHECK_JUST(stream_->OnExecutionContextTeardown());
   });
 }
 
@@ -111,6 +109,8 @@ REGISTER_STREAM_CONTEXT_CREATOR_WITH_STREAM_ID(
       return new CudaStreamContext(stream_id.device_index());
     }));
 
+}  // namespace
+
 }  // namespace oneflow
 
 #endif

oneflow/core/lazy/stream_context/include/generic_stream_context.h

@@ -0,0 +1,48 @@
+/*
+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_LAZY_STREAM_CONTEXT_GENERIC_STREAM_CONTEXT_H_
+#define ONEFLOW_CORE_LAZY_STREAM_CONTEXT_GENERIC_STREAM_CONTEXT_H_
+
+#include "oneflow/core/lazy/stream_context/include/stream_context.h"
+#include "oneflow/core/common/device_type.h"
+#include "oneflow/core/graph/stream_id.h"
+#include "oneflow/core/ep/include/stream.h"
+#include "oneflow/core/ep/include/device.h"
+#include "oneflow/core/common/channel.h"
+
+namespace oneflow {
+
+class GenericStreamContext : public StreamContext {
+ public:
+  OF_DISALLOW_COPY_AND_MOVE(GenericStreamContext);
+  explicit GenericStreamContext(const StreamId& stream_id);
+  ~GenericStreamContext() override;
+
+  Maybe<void> AddCallback(std::function<void()> callback) override;
+  DeviceType device_type() const override { return stream_->device_type(); }
+
+  ep::Stream* stream() override;
+
+ private:
+  ep::Stream* stream_;
+  Channel<std::pair<ep::Event*, std::function<void()>>> cb_event_chan_;
+  std::thread poller_thread_;
+  std::shared_ptr<ep::Device> device_;
+};
+
+}  // namespace oneflow
+
+#endif  // ONEFLOW_CORE_LAZY_STREAM_CONTEXT_GENERIC_STREAM_CONTEXT_H_

oneflow/core/stream/include/stream_context.h → oneflow/core/lazy/stream_context/include/stream_context.h

@@ -13,8 +13,8 @@ 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_STREAM_STREAM_CONTEXT_H_
-#define ONEFLOW_CORE_STREAM_STREAM_CONTEXT_H_
+#ifndef ONEFLOW_CORE_LAZY_STREAM_CONTEXT_STREAM_CONTEXT_H_
+#define ONEFLOW_CORE_LAZY_STREAM_CONTEXT_STREAM_CONTEXT_H_
 
 #include "oneflow/core/common/util.h"
 #include "oneflow/core/common/auto_registration_factory.h"
@@ -39,4 +39,4 @@ class StreamContext {
 
 }  // namespace oneflow
 
-#endif  // ONEFLOW_CORE_STREAM_STREAM_CONTEXT_H_
+#endif  // ONEFLOW_CORE_LAZY_STREAM_CONTEXT_STREAM_CONTEXT_H_

oneflow/core/thread/thread.cpp

@@ -19,17 +19,23 @@ limitations under the License.
 #include "oneflow/core/lazy/actor/actor.h"
 #include "oneflow/core/lazy/actor/light_actor.h"
 #include "oneflow/core/profiler/profiler.h"
-#include "oneflow/core/stream/include/stream_context.h"
+#include "oneflow/core/lazy/stream_context/include/stream_context.h"
 #include "oneflow/core/framework/to_string.h"
+#include "oneflow/core/lazy/stream_context/include/generic_stream_context.h"
 
 namespace oneflow {
 
 Thread::Thread(const StreamId& stream_id) : thrd_id_(EncodeStreamIdToInt64(stream_id)) {
   local_msg_queue_enabled_ = ParseBooleanFromEnv("ONEFLOW_THREAD_ENABLE_LOCAL_MESSAGE_QUEUE", true);
   light_actor_enabled_ = ParseBooleanFromEnv("ONEFLOW_ACTOR_ENABLE_LIGHT_ACTOR", true);
-  StreamContext* stream_ctx =
-      NewObj<int, StreamContext, const StreamId&>(stream_id.device_id().device_type(), stream_id);
-  stream_ctx_.reset(stream_ctx);
+  if (IsClassRegistered<int, StreamContext, const StreamId&>(stream_id.device_id().device_type(),
+                                                             stream_id)) {
+    stream_ctx_.reset(NewObj<int, StreamContext, const StreamId&>(
+        stream_id.device_id().device_type(), stream_id));
+  } else {
+    stream_ctx_.reset(new GenericStreamContext(stream_id));
+  }
+
   actor_thread_ = std::thread([this, stream_id]() {
     OF_PROFILER_NAME_THIS_HOST_THREAD("_" + ToString(stream_id.device_id().device_type())
                                       + std::to_string(stream_id.device_id().device_index())