device_context.h 5.6 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
27
#include <mkldnn.hpp>
T
tensor-tang 已提交
28 29
#endif

Y
Yi Wang 已提交
30 31
#include "paddle/fluid/platform/enforce.h"
#include "paddle/fluid/platform/place.h"
Q
qijun 已提交
32
#include "unsupported/Eigen/CXX11/Tensor"
Q
QI JUN 已提交
33

D
dzhwinter 已提交
34 35
#include "glog/logging.h"

Q
QI JUN 已提交
36 37 38 39 40 41
namespace paddle {
namespace platform {

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

44
  virtual void Wait() const {}
Q
QI JUN 已提交
45 46
};

Q
qijun 已提交
47 48
class CPUDeviceContext : public DeviceContext {
 public:
49
  CPUDeviceContext();
Q
qijun 已提交
50
  explicit CPUDeviceContext(CPUPlace place);
Q
qijun 已提交
51

52
  Eigen::DefaultDevice* eigen_device() const;
Q
qijun 已提交
53

L
liaogang 已提交
54
  Place GetPlace() const override;
Y
Yu Yang 已提交
55

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

Y
Yang Yu 已提交
61 62 63 64 65 66 67 68
template <typename Place>
struct DefaultDeviceContextType;

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

69
#ifdef PADDLE_WITH_CUDA
70

Q
qijun 已提交
71
class EigenCudaStreamDevice;
D
dongzhihong 已提交
72

73
class CUDADeviceContext : public DeviceContext {
Q
QI JUN 已提交
74
 public:
D
dzhwinter 已提交
75
  explicit CUDADeviceContext(CUDAPlace place);
76
  virtual ~CUDADeviceContext();
Q
QI JUN 已提交
77

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

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

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

87 88 89
  /*! \brief  Return the max physical thread count in the device context */
  int GetMaxPhysicalThreadCount() const;

90 91 92 93
  /*! \brief  Return eigen device in the device context. */
  Eigen::GpuDevice* eigen_device() const;

  /*! \brief  Return cublas handle in the device context. */
94
  cublasHandle_t cublas_handle() const;
95 96

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

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

Y
Yu Yang 已提交
102 103
  template <typename Callback>
  void RecordEvent(cudaEvent_t ev, Callback callback) {
Y
yuyang18 已提交
104
    std::lock_guard<std::mutex> guard(mtx_);
Y
Yu Yang 已提交
105 106 107 108
    callback();
    PADDLE_ENFORCE(cudaEventRecord(ev, stream_));
  }

Q
QI JUN 已提交
109
 private:
D
dzhwinter 已提交
110
  CUDAPlace place_;
Q
QI JUN 已提交
111

Q
qijun 已提交
112
  std::unique_ptr<Eigen::GpuDevice> eigen_device_;
Q
init  
qijun 已提交
113
  std::unique_ptr<EigenCudaStreamDevice> eigen_stream_;
114 115 116
  cudaStream_t stream_;
  cudnnHandle_t cudnn_handle_;
  cublasHandle_t cublas_handle_;
117

K
Kexin Zhao 已提交
118
  int compute_capability;
119 120
  int multi_process;
  int max_threads_per_mp;
Y
yuyang18 已提交
121 122

  std::mutex mtx_;
Q
QI JUN 已提交
123
};
Q
qijun 已提交
124

Y
Yang Yu 已提交
125 126
template <>
struct DefaultDeviceContextType<platform::CUDAPlace> {
Y
Yang Yu 已提交
127
  using TYPE = CUDADeviceContext;
Y
Yang Yu 已提交
128 129
};

C
chengduoZH 已提交
130
// Currently, CUDAPinnedDeviceContext is only used to data copying.
C
chengduoZH 已提交
131 132 133 134 135 136
class CUDAPinnedDeviceContext : public DeviceContext {
 public:
  CUDAPinnedDeviceContext();
  explicit CUDAPinnedDeviceContext(CUDAPinnedPlace place);

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

C
chengduoZH 已提交
138 139 140 141 142 143 144 145 146 147 148
  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 已提交
149
#endif
Q
qijun 已提交
150

T
tensor-tang 已提交
151 152 153 154 155 156
#ifdef PADDLE_WITH_MKLDNN
class MKLDNNDeviceContext : public CPUDeviceContext {
 public:
  explicit MKLDNNDeviceContext(CPUPlace place);

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

159 160
  // 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 已提交
161

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

 private:
166 167 168
  mkldnn::engine engine_;
  std::shared_ptr<std::unordered_map<std::string, std::shared_ptr<void>>>
      p_blobs_;
T
tensor-tang 已提交
169 170 171
};
#endif

D
dzhwinter 已提交
172 173 174 175 176
/*! \brief device context pool singleton */
class DeviceContextPool {
 public:
  explicit DeviceContextPool(const std::vector<platform::Place>& places);

Y
Yang Yu 已提交
177
  static DeviceContextPool& Instance() {
D
dzhwinter 已提交
178 179 180 181 182
    PADDLE_ENFORCE_NOT_NULL(pool, "Need to Create DeviceContextPool first!");
    return *pool;
  }

  /*! \brief  Create should only called by Init function */
Y
Yang Yu 已提交
183
  static DeviceContextPool& Init(const std::vector<platform::Place>& places) {
D
dzhwinter 已提交
184 185 186 187 188 189 190
    if (pool == nullptr) {
      pool = new DeviceContextPool(places);
    }
    return *pool;
  }

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

Y
Yang Yu 已提交
193 194 195 196 197 198 199
  template <typename Place>
  const typename DefaultDeviceContextType<Place>::TYPE* GetByPlace(
      const Place& place) {
    return reinterpret_cast<
        const typename DefaultDeviceContextType<Place>::TYPE*>(Get(place));
  }

200 201
  size_t size() const { return device_contexts_.size(); }

D
dzhwinter 已提交
202 203
 private:
  static DeviceContextPool* pool;
Y
Yu Yang 已提交
204 205
  std::unordered_map<const platform::Place,
                     std::unique_ptr<platform::DeviceContext>, PlaceHash>
D
dzhwinter 已提交
206 207 208 209
      device_contexts_;
  DISABLE_COPY_AND_ASSIGN(DeviceContextPool);
};

Q
QI JUN 已提交
210 211
}  // namespace platform
}  // namespace paddle