device_context.h 7.1 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"
D
dzhwinter 已提交
34
#if defined(PADDLE_WITH_CUDA) && !defined(_WIN32)
S
sneaxiy 已提交
35 36
#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_));
  }

D
dzhwinter 已提交
118
#ifndef _WIN32
S
sneaxiy 已提交
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();
  }
D
dzhwinter 已提交
129 130 131 132 133 134 135 136 137 138
#else
  template <typename Callback>
  void AddStreamCallback(Callback&& callback) const {
    // ugly empty functor.
  }

  void WaitStreamCallback() const {
    // ugly empty functor.
  }
#endif
S
sneaxiy 已提交
139

Q
QI JUN 已提交
140
 private:
D
dzhwinter 已提交
141
  CUDAPlace place_;
Q
QI JUN 已提交
142

Q
qijun 已提交
143
  std::unique_ptr<Eigen::GpuDevice> eigen_device_;
Q
init  
qijun 已提交
144
  std::unique_ptr<EigenCudaStreamDevice> eigen_stream_;
145
  std::unique_ptr<CudnnHolder> cudnn_holder_;
146 147
  cudaStream_t stream_;
  cublasHandle_t cublas_handle_;
148

C
chengduo 已提交
149 150 151 152 153
  int compute_capability_;
  int runtime_version_;
  int driver_version_;
  int multi_process_;
  int max_threads_per_mp_;
Y
yuyang18 已提交
154

S
sneaxiy 已提交
155 156
  mutable std::mutex mtx_;

D
dzhwinter 已提交
157
#ifndef _WIN32
S
sneaxiy 已提交
158 159 160 161
  // 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_;
D
dzhwinter 已提交
162
#endif
Q
QI JUN 已提交
163
};
Q
qijun 已提交
164

Y
Yang Yu 已提交
165 166
template <>
struct DefaultDeviceContextType<platform::CUDAPlace> {
Y
Yang Yu 已提交
167
  using TYPE = CUDADeviceContext;
Y
Yang Yu 已提交
168 169
};

C
chengduoZH 已提交
170
// Currently, CUDAPinnedDeviceContext is only used to data copying.
C
chengduoZH 已提交
171 172 173 174 175 176
class CUDAPinnedDeviceContext : public DeviceContext {
 public:
  CUDAPinnedDeviceContext();
  explicit CUDAPinnedDeviceContext(CUDAPinnedPlace place);

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

C
chengduoZH 已提交
178 179 180 181 182 183 184 185 186 187 188
  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 已提交
189
#endif
Q
qijun 已提交
190

T
tensor-tang 已提交
191
#ifdef PADDLE_WITH_MKLDNN
S
Sylwester Fraczek 已提交
192 193 194 195 196 197
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 已提交
198 199 200 201 202
class MKLDNNDeviceContext : public CPUDeviceContext {
 public:
  explicit MKLDNNDeviceContext(CPUPlace place);

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

205 206
  // 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 已提交
207

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

 private:
212
  mkldnn::engine engine_;
213 214
  std::shared_ptr<BlobMap> p_blobmap_;
  std::shared_ptr<std::mutex> p_mutex_;
T
tensor-tang 已提交
215 216 217
};
#endif

D
dzhwinter 已提交
218 219 220 221 222
/*! \brief device context pool singleton */
class DeviceContextPool {
 public:
  explicit DeviceContextPool(const std::vector<platform::Place>& places);

Y
Yang Yu 已提交
223
  static DeviceContextPool& Instance() {
D
dzhwinter 已提交
224 225 226 227 228
    PADDLE_ENFORCE_NOT_NULL(pool, "Need to Create DeviceContextPool first!");
    return *pool;
  }

  /*! \brief  Create should only called by Init function */
Y
Yang Yu 已提交
229
  static DeviceContextPool& Init(const std::vector<platform::Place>& places) {
D
dzhwinter 已提交
230 231 232 233 234 235 236
    if (pool == nullptr) {
      pool = new DeviceContextPool(places);
    }
    return *pool;
  }

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

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

Y
Yang Yu 已提交
242 243 244 245 246 247 248
  template <typename Place>
  const typename DefaultDeviceContextType<Place>::TYPE* GetByPlace(
      const Place& place) {
    return reinterpret_cast<
        const typename DefaultDeviceContextType<Place>::TYPE*>(Get(place));
  }

249 250
  size_t size() const { return device_contexts_.size(); }

D
dzhwinter 已提交
251 252
 private:
  static DeviceContextPool* pool;
253
  std::map<Place, std::unique_ptr<DeviceContext>> device_contexts_;
D
dzhwinter 已提交
254 255 256
  DISABLE_COPY_AND_ASSIGN(DeviceContextPool);
};

Q
QI JUN 已提交
257 258
}  // namespace platform
}  // namespace paddle