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

namespace paddle {
namespace platform {

Q
qijun 已提交
17
template <>
Q
qijun 已提交
18 19 20
Eigen::DefaultDevice* DeviceContext::get_eigen_device<Eigen::DefaultDevice>()
    const {
  return reinterpret_cast<const CPUDeviceContext*>(this)->eigen_device();
Q
qijun 已提交
21 22
}

23 24 25 26 27 28 29 30 31 32 33 34
CPUDeviceContext::CPUDeviceContext() {
  eigen_device_.reset(new Eigen::DefaultDevice());
}

CPUDeviceContext::CPUDeviceContext(CPUPlace place) {
  eigen_device_.reset(new Eigen::DefaultDevice());
}

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

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

Q
qijun 已提交
37
#ifndef PADDLE_ONLY_CPU
38

Q
qijun 已提交
39
template <>
Q
qijun 已提交
40 41
Eigen::GpuDevice* DeviceContext::get_eigen_device<Eigen::GpuDevice>() const {
  return reinterpret_cast<const CUDADeviceContext*>(this)->eigen_device();
Q
qijun 已提交
42
}
43 44 45

CUDADeviceContext::CUDADeviceContext(GPUPlace place) : place_(place) {
  SetDeviceId(place_.device);
Q
qijun 已提交
46 47 48 49 50 51 52 53 54 55 56 57 58
  // TODO (qijun) Pass a created cuda stream to Eigen::CudaStreamDevice directly
  // here will cause segment fault. We must implement a class derived from
  // Eigen::StreamInterface, and reinitialize it with a cuda stream and a gpu id
  // later. Please refer to the implementation of class EigenCudaStreamDevice
  // in TensorFlow.
  //
  // We find that CUDA 7 introduces a new option, the per-thread default stream,
  // that has two effects. Please refer to https://devblogs.nvidia.com/
  // parallelforall/gpu-pro-tip-cuda-7-streams-simplify-concurrency/
  //
  // So, we decide to use default stream and add –default-stream per-thread nvcc
  // flag. Than, two threads with two CUDADeviceContexts will run parallelly.
  eigen_stream_.reset(new Eigen::CudaStreamDevice());
59 60 61 62 63
  eigen_device_.reset(new Eigen::GpuDevice(eigen_stream_.get()));
}

CUDADeviceContext::~CUDADeviceContext() {
  SetDeviceId(place_.device);
L
liaogang 已提交
64
  Wait();
65 66 67 68 69 70 71 72 73 74 75 76 77 78 79
  if (cublas_handle_) {
    PADDLE_ENFORCE(dynload::cublasDestroy(cublas_handle_));
  }

  if (cudnn_handle_) {
    PADDLE_ENFORCE(dynload::cudnnDestroy(cudnn_handle_));
  }

  if (curand_generator_) {
    PADDLE_ENFORCE(dynload::curandDestroyGenerator(curand_generator_));
  }
  eigen_stream_.reset();
  eigen_device_.reset();
}

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

L
liaogang 已提交
82
void CUDADeviceContext::Wait() const {
Q
qijun 已提交
83
  PADDLE_ENFORCE(cudaStreamSynchronize(0));
84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117
}

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

cublasHandle_t CUDADeviceContext::cublas_handle() {
  if (!cublas_handle_) {
    SetDeviceId(place_.device);
    PADDLE_ENFORCE(dynload::cublasCreate(&cublas_handle_));
  }
  return cublas_handle_;
}

cudnnHandle_t CUDADeviceContext::cudnn_handle() {
  if (!cudnn_handle_) {
    SetDeviceId(place_.device);
    PADDLE_ENFORCE(dynload::cudnnCreate(&cudnn_handle_));
  }
  return cudnn_handle_;
}

curandGenerator_t CUDADeviceContext::curand_generator() {
  if (!curand_generator_) {
    SetDeviceId(place_.device);
    PADDLE_ENFORCE(dynload::curandCreateGenerator(&curand_generator_,
                                                  CURAND_RNG_PSEUDO_DEFAULT));
    PADDLE_ENFORCE(
        dynload::curandSetPseudoRandomGeneratorSeed(curand_generator_, seed_));
  }
  return curand_generator_;
}

#endif  // PADDLE_ONLY_CPU
Q
qijun 已提交
118

Q
qijun 已提交
119
}  // namespace platform
Q
qijun 已提交
120
}  // namespace paddle