提交 4d336d90 编写于 作者: Q qijun

follow comments

上级 ef5f9deb
......@@ -20,7 +20,6 @@ limitations under the License. */
#include "paddle/platform/dynload/cudnn.h"
#include "paddle/platform/dynload/curand.h"
#define EIGEN_USE_GPU
#include "paddle/platform/device.h"
#include "paddle/platform/place.h"
#include "unsupported/Eigen/CXX11/Tensor"
......@@ -29,6 +28,13 @@ using DEVICE_GPU = Eigen::GpuDevice;
namespace paddle {
namespace platform {
class CUDADeviceContext;
template <>
DEVICE_GPU DeviceContext::get_eigen_device<DEVICE_GPU>() {
return static_cast<CUDADeviceContext*>(this)->eigen_handle();
}
class GPUPlaceGuard {
public:
explicit GPUPlaceGuard(GPUPlace new_place) : previous_(GetCurrentDeviceId()) {
......@@ -43,8 +49,7 @@ class GPUPlaceGuard {
GPUPlace previous_;
};
template <>
class Device<DEVICE_GPU> {
class CUDADeviceContext : public DeviceContext {
public:
explicit Device(const GPUPlace gpu_place) : gpu_place_(gpu_place) {
GPUPlaceGuard guard(gpu_place_);
......@@ -61,7 +66,7 @@ class Device<DEVICE_GPU> {
cudaStream_t stream() { return stream_; }
DEVICE_GPU eigen_device() { return *eigen_device_; }
Eigen::GpuDevice eigen_device() { return *eigen_device_; }
cublasHandle_t cublas_handle() {
if (!blas_handle_) {
......@@ -139,7 +144,7 @@ class Device<DEVICE_GPU> {
cudaStream_t stream_;
Eigen::CudaStreamDevice* eigen_stream_;
DEVICE_GPU* eigen_device_;
Eigen::GpuDevice* eigen_device_;
cublasHandle_t blas_handle_{nullptr};
......
/* 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/cuda_device.h"
#include "gtest/gtest.h"
TEST(Device, Init) {
int count = paddle::platform::GetDeviceCount();
for (int i = 0; i < count; i++) {
paddle::platform::Device<DEVICE_GPU>* device =
new paddle::platform::Device<DEVICE_GPU>(i);
Eigen::GpuDevice gpu_device = device->eigen_device();
ASSERT_NE(nullptr, gpu_device.stream());
cudnnHandle_t cudnn_handle = device->cudnn_handle();
ASSERT_NE(nullptr, cudnn_handle);
cublasHandle_t cublas_handle = device->cublas_handle();
ASSERT_NE(nullptr, cublas_handle);
curandGenerator_t curand_handle = device->curand_generator();
ASSERT_NE(nullptr, curand_handle);
delete device;
}
}
/* 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
#include "unsupported/Eigen/CXX11/Tensor"
using DEVICE_CPU = Eigen::DefaultDevice;
namespace paddle {
namespace platform {
template <typename DeviceType>
class Device;
template <>
class Device<DEVICE_CPU> {
public:
DEVICE_CPU eigen_handle() {
if (!eigen_handle_) {
eigen_handle_ = new Eigen::DefaultDevice();
}
return *eigen_handle_;
}
private:
DEVICE_CPU* eigen_handle_{nullptr};
};
} // namespace platform
} // namespace paddle
......@@ -13,23 +13,39 @@ See the License for the specific language governing permissions and
limitations under the License. */
#pragma once
#include "paddle/framework/enforce.h"
#include "paddle/platform/device.h"
#include "unsupported/Eigen/CXX11/Tensor"
#ifndef PADDLE_ONLY_CPU
#include "paddle/platform/cuda_device.h"
#endif
using DEVICE_CPU = Eigen::DefaultDevice;
namespace paddle {
namespace platform {
struct DeviceContext {
void* device_context{nullptr};
class CPUDeviceContext;
class DeviceContext {
public:
virtual ~DeviceContext() {}
template <typename DeviceType>
inline paddle::platform::Device<DeviceType>* device_context() {
return static_cast<paddle::platform::Device<DeviceType>*>(device_context);
DeviceType get_eigen_device();
};
template <>
DEVICE_CPU DeviceContext::get_eigen_device<DEVICE_CPU>() {
return static_cast<CPUDeviceContext*>(this)->eigen_handle();
}
class CPUDeviceContext : public DeviceContext {
public:
Eigen::DefaultDevice eigen_handle() {
if (!eigen_handle_) {
eigen_handle_ = new Eigen::DefaultDevice();
}
return *eigen_handle_;
}
private:
Eigen::DefaultDevice* eigen_handle_{nullptr};
};
} // namespace platform
......
......@@ -12,19 +12,34 @@ 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"
#include "gtest/gtest.h"
#include "paddle/platform/cuda_device.h"
TEST(DeviceContext, Init) {
TEST(Device, Init) {
int count = paddle::platform::GetDeviceCount();
for (int i = 0; i < count; i++) {
paddle::platform::Device<DEVICE_GPU>* device =
new paddle::platform::Device<DEVICE_GPU>(i);
paddle::platform::DeviceContext context;
context.device_context = device;
paddle::platform::DeviceContext* device_context =
new paddle::platform::CUDADeviceContext(i);
Eigen::GpuDevice gpu_device =
context.device_context<DEVICE_GPU>->eigen_device();
device_context->get_eigen_device<DEVICE_GPU>();
ASSERT_NE(nullptr, gpu_device.stream());
delete device;
delete device_context;
}
}
TEST(Device, CUDADeviceContext) {
int count = paddle::platform::GetDeviceCount();
for (int i = 0; i < count; i++) {
paddle::platform::CUDADeviceContext* device_context =
new paddle::platform::CUDADeviceContext(i);
Eigen::GpuDevice gpu_device = device_context->eigen_device();
ASSERT_NE(nullptr, gpu_device.stream());
cudnnHandle_t cudnn_handle = device_context->cudnn_handle();
ASSERT_NE(nullptr, cudnn_handle);
cublasHandle_t cublas_handle = device_context->cublas_handle();
ASSERT_NE(nullptr, cublas_handle);
curandGenerator_t curand_handle = device_context->curand_generator();
ASSERT_NE(nullptr, curand_handle);
delete device_context;
}
}
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册