device_context.cc 4.5 KB
Newer Older
Q
qijun 已提交
1 2 3 4 5 6 7 8 9 10 11 12
/* 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. */

#include "paddle/platform/device_context.h"
Q
init  
qijun 已提交
13
#include "paddle/memory/memory.h"
Q
qijun 已提交
14 15 16 17

namespace paddle {
namespace platform {

18 19 20 21
CPUDeviceContext::CPUDeviceContext() {
  eigen_device_.reset(new Eigen::DefaultDevice());
}

Q
qijun 已提交
22
CPUDeviceContext::CPUDeviceContext(CPUPlace place) {
23 24 25 26 27 28 29
  eigen_device_.reset(new Eigen::DefaultDevice());
}

Eigen::DefaultDevice* CPUDeviceContext::eigen_device() const {
  return eigen_device_.get();
}

L
liaogang 已提交
30
Place CPUDeviceContext::GetPlace() const { return CPUPlace(); }
31

32
#ifdef PADDLE_WITH_CUDA
33

Q
init  
qijun 已提交
34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53
class EigenCudaStreamDevice : public Eigen::StreamInterface {
 public:
  EigenCudaStreamDevice() : scratch_(nullptr), semaphore_(nullptr) {
    Eigen::initializeDeviceProp();
  }
  ~EigenCudaStreamDevice() override {}

  void Reinitialize(const cudaStream_t* cuda_stream, GPUPlace place) {
    stream_ = cuda_stream;
    place_ = place;
    device_prop_ = &Eigen::m_deviceProperties[place.device];
  }

  const cudaStream_t& stream() const override { return *stream_; }

  const cudaDeviceProp& deviceProperties() const override {
    return *device_prop_;
  }

  void* allocate(size_t num_bytes) const override {
Q
qijun 已提交
54
    return paddle::memory::Alloc(place_, num_bytes);
Q
init  
qijun 已提交
55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82
  }

  void deallocate(void* buffer) const override {
    paddle::memory::Free(place_, buffer);
  }

  void* scratchpad() const override {
    if (scratch_ == NULL) {
      scratch_ = allocate(Eigen::kCudaScratchSize + sizeof(unsigned int));
    }
    return scratch_;
  }

  unsigned int* semaphore() const override {
    if (semaphore_ == NULL) {
      char* scratch =
          static_cast<char*>(scratchpad()) + Eigen::kCudaScratchSize;
      semaphore_ = reinterpret_cast<unsigned int*>(scratch);
      PADDLE_ENFORCE(
          cudaMemsetAsync(semaphore_, 0, sizeof(unsigned int), *stream_));
    }
    return semaphore_;
  }

 private:
  GPUPlace place_;
  const cudaStream_t* stream_;         // not owned;
  const cudaDeviceProp* device_prop_;  // not owned;
Q
qijun 已提交
83
  mutable void* scratch_;
Q
init  
qijun 已提交
84 85 86
  mutable unsigned int* semaphore_;
};

Q
qijun 已提交
87
CUDADeviceContext::CUDADeviceContext(GPUPlace place) : place_(place) {
88
  SetDeviceId(place_.device);
Q
init  
qijun 已提交
89 90 91
  PADDLE_ENFORCE(cudaStreamCreate(&stream_));
  eigen_stream_.reset(new EigenCudaStreamDevice());
  eigen_stream_->Reinitialize(&stream_, place);
92
  eigen_device_.reset(new Eigen::GpuDevice(eigen_stream_.get()));
93 94 95 96
  PADDLE_ENFORCE(dynload::cublasCreate(&cublas_handle_));
  PADDLE_ENFORCE(dynload::cublasSetStream(cublas_handle_, stream_));
  PADDLE_ENFORCE(dynload::cudnnCreate(&cudnn_handle_));
  PADDLE_ENFORCE(dynload::cudnnSetStream(cudnn_handle_, stream_));
97 98 99 100
}

CUDADeviceContext::~CUDADeviceContext() {
  SetDeviceId(place_.device);
L
liaogang 已提交
101
  Wait();
102 103
  PADDLE_ENFORCE(dynload::cublasDestroy(cublas_handle_));
  PADDLE_ENFORCE(dynload::cudnnDestroy(cudnn_handle_));
104 105
  eigen_stream_.reset();
  eigen_device_.reset();
Q
init  
qijun 已提交
106
  PADDLE_ENFORCE(cudaStreamDestroy(stream_));
107 108
}

L
liaogang 已提交
109
Place CUDADeviceContext::GetPlace() const { return place_; }
110

L
liaogang 已提交
111
void CUDADeviceContext::Wait() const {
Q
init  
qijun 已提交
112
  PADDLE_ENFORCE(cudaStreamSynchronize(stream_));
113 114 115
  PADDLE_ENFORCE(cudaGetLastError());
}

116 117 118 119
Eigen::GpuDevice* CUDADeviceContext::eigen_device() const {
  return eigen_device_.get();
}

120
cublasHandle_t CUDADeviceContext::cublas_handle() const {
121 122 123
  return cublas_handle_;
}

124
cudnnHandle_t CUDADeviceContext::cudnn_handle() const { return cudnn_handle_; }
125

126
cudaStream_t CUDADeviceContext::stream() const { return stream_; }
Q
qijun 已提交
127

D
dzhwinter 已提交
128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143
CudnnDeviceContext::CudnnDeviceContext(CudnnPlace place)
    : CUDADeviceContext(place), place_(place) {
  PADDLE_ENFORCE(dynload::cudnnCreate(&cudnn_handle_));
  PADDLE_ENFORCE(dynload::cudnnSetStream(cudnn_handle_, stream()));
}

CudnnDeviceContext::~CudnnDeviceContext() {
  SetDeviceId(place_.device);
  Wait();
  PADDLE_ENFORCE(dynload::cudnnDestroy(cudnn_handle_));
}

Place CudnnDeviceContext::GetPlace() const { return CudnnPlace(); }

cudnnHandle_t CudnnDeviceContext::cudnn_handle() const { return cudnn_handle_; }

L
Luo Tao 已提交
144
#endif
Q
qijun 已提交
145

Q
qijun 已提交
146
}  // namespace platform
Q
qijun 已提交
147
}  // namespace paddle