device_context.h 6.8 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
Q
QI JUN 已提交
2 3 4 5 6 7 8 9 10 11 12
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. */
#pragma once

D
dzhwinter 已提交
13
#include <memory>
Y
yuyang18 已提交
14
#include <mutex>  // NOLINT
15
#include <string>
D
dzhwinter 已提交
16
#include <unordered_map>
17
#include <vector>
L
liaogang 已提交
18

19
#ifdef PADDLE_WITH_CUDA
Y
Yi Wang 已提交
20 21 22
#include "paddle/fluid/platform/dynload/cublas.h"
#include "paddle/fluid/platform/dynload/cudnn.h"
#include "paddle/fluid/platform/gpu_info.h"
Q
QI JUN 已提交
23 24
#define EIGEN_USE_GPU
#endif
D
dzhwinter 已提交
25

T
tensor-tang 已提交
26
#ifdef PADDLE_WITH_MKLDNN
L
luotao1 已提交
27
#include "mkldnn.hpp"
T
tensor-tang 已提交
28 29
#endif

30 31
#include <map>
#include "glog/logging.h"
Y
Yi Wang 已提交
32 33
#include "paddle/fluid/platform/enforce.h"
#include "paddle/fluid/platform/place.h"
S
sneaxiy 已提交
34 35 36
#ifdef PADDLE_WITH_CUDA
#include "paddle/fluid/platform/stream_callback_manager.h"
#endif
Q
qijun 已提交
37
#include "unsupported/Eigen/CXX11/Tensor"
Q
QI JUN 已提交
38 39 40 41 42 43 44

namespace paddle {
namespace platform {

class DeviceContext {
 public:
  virtual ~DeviceContext() {}
L
liaogang 已提交
45
  virtual Place GetPlace() const = 0;
Q
QI JUN 已提交
46

47
  virtual void Wait() const {}
Q
QI JUN 已提交
48 49
};

Q
qijun 已提交
50 51
class CPUDeviceContext : public DeviceContext {
 public:
52
  CPUDeviceContext();
Q
qijun 已提交
53
  explicit CPUDeviceContext(CPUPlace place);
Q
qijun 已提交
54

55
  Eigen::DefaultDevice* eigen_device() const;
Q
qijun 已提交
56

L
liaogang 已提交
57
  Place GetPlace() const override;
Y
Yu Yang 已提交
58

Q
qijun 已提交
59
 private:
D
dzhwinter 已提交
60
  CPUPlace place_;
Q
qijun 已提交
61
  std::unique_ptr<Eigen::DefaultDevice> eigen_device_;
Q
QI JUN 已提交
62 63
};

Y
Yang Yu 已提交
64 65 66 67 68 69 70 71
template <typename Place>
struct DefaultDeviceContextType;

template <>
struct DefaultDeviceContextType<platform::CPUPlace> {
  using TYPE = CPUDeviceContext;
};

72
#ifdef PADDLE_WITH_CUDA
73

Q
qijun 已提交
74
class EigenCudaStreamDevice;
75
class CudnnHolder;
D
dongzhihong 已提交
76

77
class CUDADeviceContext : public DeviceContext {
Q
QI JUN 已提交
78
 public:
D
dzhwinter 已提交
79
  explicit CUDADeviceContext(CUDAPlace place);
80
  virtual ~CUDADeviceContext();
Q
QI JUN 已提交
81

82
  /*! \brief  Wait for all operations completion in the stream. */
83
  void Wait() const override;
Q
QI JUN 已提交
84

85
  /*! \brief  Return place in the device context. */
L
liaogang 已提交
86
  Place GetPlace() const override;
87

K
Kexin Zhao 已提交
88
  /*! \brief  Return compute capability in the device context. */
K
Kexin Zhao 已提交
89 90
  int GetComputeCapability() const;

91 92 93
  /*! \brief  Return the max physical thread count in the device context */
  int GetMaxPhysicalThreadCount() const;

94 95 96 97
  /*! \brief  Return eigen device in the device context. */
  Eigen::GpuDevice* eigen_device() const;

  /*! \brief  Return cublas handle in the device context. */
98
  cublasHandle_t cublas_handle() const;
99 100

  /*! \brief  Return cudnn  handle in the device context. */
101
  cudnnHandle_t cudnn_handle() const;
102

103 104 105 106 107
  /*! \brief  Run a cudnn function with the workspace provided by
   * CUDADeviceContext */
  void RunCudnnFuncWithWorkspace(const std::function<void(void*)>& cudnn_func,
                                 size_t workspace_len) const;

Q
init  
qijun 已提交
108
  /*! \brief  Return cuda stream in the device context. */
109
  cudaStream_t stream() const;
Q
QI JUN 已提交
110

Y
Yu Yang 已提交
111 112
  template <typename Callback>
  void RecordEvent(cudaEvent_t ev, Callback callback) {
Y
yuyang18 已提交
113
    std::lock_guard<std::mutex> guard(mtx_);
Y
Yu Yang 已提交
114 115 116 117
    callback();
    PADDLE_ENFORCE(cudaEventRecord(ev, stream_));
  }

S
sneaxiy 已提交
118 119 120 121 122 123 124 125 126 127 128
  template <typename Callback>
  void AddStreamCallback(Callback&& callback) const {
    std::lock_guard<std::mutex> guard(callback_mtx_);
    callback_manager_->AddCallback(callback);
  }

  void WaitStreamCallback() const {
    std::lock_guard<std::mutex> guard(callback_mtx_);
    callback_manager_->Wait();
  }

Q
QI JUN 已提交
129
 private:
D
dzhwinter 已提交
130
  CUDAPlace place_;
Q
QI JUN 已提交
131

Q
qijun 已提交
132
  std::unique_ptr<Eigen::GpuDevice> eigen_device_;
Q
init  
qijun 已提交
133
  std::unique_ptr<EigenCudaStreamDevice> eigen_stream_;
134
  std::unique_ptr<CudnnHolder> cudnn_holder_;
135 136
  cudaStream_t stream_;
  cublasHandle_t cublas_handle_;
137

C
chengduo 已提交
138 139 140 141 142
  int compute_capability_;
  int runtime_version_;
  int driver_version_;
  int multi_process_;
  int max_threads_per_mp_;
Y
yuyang18 已提交
143

S
sneaxiy 已提交
144 145 146 147 148 149
  mutable std::mutex mtx_;

  // This lock is only used by callback
  // If we use mtx_ for StreamCallbackManager, deadlock may occur sometimes
  mutable std::mutex callback_mtx_;
  std::unique_ptr<StreamCallbackManager> callback_manager_;
Q
QI JUN 已提交
150
};
Q
qijun 已提交
151

Y
Yang Yu 已提交
152 153
template <>
struct DefaultDeviceContextType<platform::CUDAPlace> {
Y
Yang Yu 已提交
154
  using TYPE = CUDADeviceContext;
Y
Yang Yu 已提交
155 156
};

C
chengduoZH 已提交
157
// Currently, CUDAPinnedDeviceContext is only used to data copying.
C
chengduoZH 已提交
158 159 160 161 162 163
class CUDAPinnedDeviceContext : public DeviceContext {
 public:
  CUDAPinnedDeviceContext();
  explicit CUDAPinnedDeviceContext(CUDAPinnedPlace place);

  Place GetPlace() const override;
C
chengduoZH 已提交
164

C
chengduoZH 已提交
165 166 167 168 169 170 171 172 173 174 175
  Eigen::DefaultDevice* eigen_device() const;

 private:
  CUDAPinnedPlace place_;
  std::unique_ptr<Eigen::DefaultDevice> eigen_device_;
};

template <>
struct DefaultDeviceContextType<platform::CUDAPinnedPlace> {
  using TYPE = CUDAPinnedDeviceContext;
};
Q
QI JUN 已提交
176
#endif
Q
qijun 已提交
177

T
tensor-tang 已提交
178
#ifdef PADDLE_WITH_MKLDNN
S
Sylwester Fraczek 已提交
179 180 181 182 183 184
using KeyBlob = std::unordered_map<std::string, std::shared_ptr<void>>;
using BlobMap = std::unordered_map<int, std::shared_ptr<KeyBlob>>;

void set_cur_thread_id(int);
int get_cur_thread_id(void);

T
tensor-tang 已提交
185 186 187 188 189
class MKLDNNDeviceContext : public CPUDeviceContext {
 public:
  explicit MKLDNNDeviceContext(CPUPlace place);

  /* \brief  Get the active engine */
190
  const mkldnn::engine& GetEngine() const { return engine_; }
T
tensor-tang 已提交
191

192 193
  // Set data to blob (i.e. name/data pair). Create blob if not existing
  void SetBlob(const std::string& name, std::shared_ptr<void> data) const;
T
tensor-tang 已提交
194

195 196
  // Find a saved blob. Return nullptr if not found
  std::shared_ptr<void> GetBlob(const std::string& name) const;
T
tensor-tang 已提交
197 198

 private:
199
  mkldnn::engine engine_;
200 201
  std::shared_ptr<BlobMap> p_blobmap_;
  std::shared_ptr<std::mutex> p_mutex_;
T
tensor-tang 已提交
202 203 204
};
#endif

D
dzhwinter 已提交
205 206 207 208 209
/*! \brief device context pool singleton */
class DeviceContextPool {
 public:
  explicit DeviceContextPool(const std::vector<platform::Place>& places);

Y
Yang Yu 已提交
210
  static DeviceContextPool& Instance() {
D
dzhwinter 已提交
211 212 213 214 215
    PADDLE_ENFORCE_NOT_NULL(pool, "Need to Create DeviceContextPool first!");
    return *pool;
  }

  /*! \brief  Create should only called by Init function */
Y
Yang Yu 已提交
216
  static DeviceContextPool& Init(const std::vector<platform::Place>& places) {
D
dzhwinter 已提交
217 218 219 220 221 222 223
    if (pool == nullptr) {
      pool = new DeviceContextPool(places);
    }
    return *pool;
  }

  /*! \brief  Return handle of single device context. */
Y
Yu Yang 已提交
224
  platform::DeviceContext* Get(const platform::Place& place);
D
dzhwinter 已提交
225

C
chengduozh 已提交
226 227 228
  /*! \brief  Return all the device contexts. */
  const std::vector<const DeviceContext*> GetAllDeviceContexts() const;

Y
Yang Yu 已提交
229 230 231 232 233 234 235
  template <typename Place>
  const typename DefaultDeviceContextType<Place>::TYPE* GetByPlace(
      const Place& place) {
    return reinterpret_cast<
        const typename DefaultDeviceContextType<Place>::TYPE*>(Get(place));
  }

236 237
  size_t size() const { return device_contexts_.size(); }

D
dzhwinter 已提交
238 239
 private:
  static DeviceContextPool* pool;
240
  std::map<Place, std::unique_ptr<DeviceContext>> device_contexts_;
D
dzhwinter 已提交
241 242 243
  DISABLE_COPY_AND_ASSIGN(DeviceContextPool);
};

Q
QI JUN 已提交
244 245
}  // namespace platform
}  // namespace paddle