Caffe2 - C++ API
A deep learning, cross platform ML framework
net_async_dag_gpu.h
1 
17 #ifndef CAFFE2_CORE_NET_ASYNC_DAG_GPU_H_
18 #define CAFFE2_CORE_NET_ASYNC_DAG_GPU_H_
19 
20 #include "caffe2/core/common.h"
21 #include "caffe2/core/net_dag.h"
22 #include "caffe2/core/workspace.h"
23 #include "caffe2/proto/caffe2.pb.h"
24 
25 namespace caffe2 {
26 
27 // Run an event-driven graph - before each operator chain, wait on each parent
28 // operator for the chain source, then execute each operator. Due to the chain
29 // construction mechanism, operators in the same chain implicitly runs on the
30 // same stream.
31 // AsyncDAGNet is only registered in gpu mode, because CPU code is always sync
32 // and a CPU only AsyncDAG net is essentially a DAG net.
33 class AsyncDAGNet : public DAGNetBase {
34  public:
35  AsyncDAGNet(const std::shared_ptr<const NetDef>& net_def, Workspace* ws);
36  bool SupportsAsync() override {
37  return true;
38  }
39  bool RunAt(int chain_id, const std::vector<int>& chain) override;
40 
41  protected:
42  bool DoRunAsync() override;
43 
44  // Tracks whether a given op has had an event recorded in each
45  // RunAt() iteration.
46  std::vector<int32_t> eventRecorded_;
47 
48  int stream(const DeviceOption& device_option);
49  static thread_local std::vector<int> stream_counters_;
50 
51  DISABLE_COPY_AND_ASSIGN(AsyncDAGNet);
52 };
53 
54 } // namespace caffe2
55 
56 #endif // CAFFE2_CORE_NET_ASYNC_DAG_GPU_H_
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
Definition: workspace.h:63
Copyright (c) 2016-present, Facebook, Inc.