device_context.h 2.6 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
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

L
liaogang 已提交
14
#include "paddle/platform/enforce.h"
L
liaogang 已提交
15 16
#include "paddle/platform/place.h"

17 18 19
#ifndef PADDLE_ONLY_CPU
#include "paddle/platform/dynload/cublas.h"
#include "paddle/platform/dynload/cudnn.h"
L
liaogang 已提交
20
#include "paddle/platform/gpu_info.h"
21 22
#define EIGEN_USE_GPU
#endif
Q
qijun 已提交
23
#include <memory>
Q
qijun 已提交
24 25
#include "paddle/platform/place.h"
#include "unsupported/Eigen/CXX11/Tensor"
26 27 28 29 30 31 32

namespace paddle {
namespace platform {

class DeviceContext {
 public:
  virtual ~DeviceContext() {}
L
liaogang 已提交
33
  virtual Place GetPlace() const = 0;
34

Q
qijun 已提交
35
  template <typename DeviceType>
Q
qijun 已提交
36
  DeviceType* get_eigen_device() const;
37 38
};

Q
qijun 已提交
39 40
class CPUDeviceContext : public DeviceContext {
 public:
41
  CPUDeviceContext();
Q
qijun 已提交
42
  explicit CPUDeviceContext(CPUPlace place);
43
  virtual ~CPUDeviceContext() {}
Q
qijun 已提交
44

45
  Eigen::DefaultDevice* eigen_device() const;
Q
qijun 已提交
46

L
liaogang 已提交
47
  Place GetPlace() const override;
Y
Yu Yang 已提交
48

Q
qijun 已提交
49
 private:
Q
qijun 已提交
50
  std::unique_ptr<Eigen::DefaultDevice> eigen_device_;
51 52 53
};

#ifndef PADDLE_ONLY_CPU
Q
qijun 已提交
54
class EigenCudaStreamDevice;
55

56
class CUDADeviceContext : public DeviceContext {
57
 public:
Q
qijun 已提交
58
  explicit CUDADeviceContext(GPUPlace place);
59
  virtual ~CUDADeviceContext();
60

61
  /*! \brief  Wait for all operations completion in the stream. */
L
liaogang 已提交
62
  void Wait() const;
63

64
  /*! \brief  Return place in the device context. */
L
liaogang 已提交
65
  Place GetPlace() const override;
66 67 68 69 70 71

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

  // clang-format off
  /*! \brief  Return cublas handle in the device context. */
72
  cublasHandle_t    cublas_handle();
73 74

  /*! \brief  Return cudnn  handle in the device context. */
75
  cudnnHandle_t     cudnn_handle();
76

Q
init  
qijun 已提交
77 78
  /*! \brief  Return cuda stream in the device context. */
  cudaStream_t      stream();
79
  // clang-format on
80 81

 private:
82
  GPUPlace place_;
83

Q
qijun 已提交
84
  std::unique_ptr<Eigen::GpuDevice> eigen_device_;
Q
init  
qijun 已提交
85
  std::unique_ptr<EigenCudaStreamDevice> eigen_stream_;
86

87
  // clang-format off
Q
qijun 已提交
88
  cudaStream_t       stream_{nullptr};
Q
init  
qijun 已提交
89 90
  cudnnHandle_t      cudnn_handle_{nullptr};
  cublasHandle_t     cublas_handle_{nullptr};
91
  // clang-format on
92
};
Q
qijun 已提交
93

94
#endif
Q
qijun 已提交
95

96 97
}  // namespace platform
}  // namespace paddle
反馈
建议
客服 返回
顶部