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

follow comments

上级 ef5f9deb
...@@ -20,7 +20,6 @@ limitations under the License. */ ...@@ -20,7 +20,6 @@ limitations under the License. */
#include "paddle/platform/dynload/cudnn.h" #include "paddle/platform/dynload/cudnn.h"
#include "paddle/platform/dynload/curand.h" #include "paddle/platform/dynload/curand.h"
#define EIGEN_USE_GPU #define EIGEN_USE_GPU
#include "paddle/platform/device.h"
#include "paddle/platform/place.h" #include "paddle/platform/place.h"
#include "unsupported/Eigen/CXX11/Tensor" #include "unsupported/Eigen/CXX11/Tensor"
...@@ -29,6 +28,13 @@ using DEVICE_GPU = Eigen::GpuDevice; ...@@ -29,6 +28,13 @@ using DEVICE_GPU = Eigen::GpuDevice;
namespace paddle { namespace paddle {
namespace platform { namespace platform {
class CUDADeviceContext;
template <>
DEVICE_GPU DeviceContext::get_eigen_device<DEVICE_GPU>() {
return static_cast<CUDADeviceContext*>(this)->eigen_handle();
}
class GPUPlaceGuard { class GPUPlaceGuard {
public: public:
explicit GPUPlaceGuard(GPUPlace new_place) : previous_(GetCurrentDeviceId()) { explicit GPUPlaceGuard(GPUPlace new_place) : previous_(GetCurrentDeviceId()) {
...@@ -43,8 +49,7 @@ class GPUPlaceGuard { ...@@ -43,8 +49,7 @@ class GPUPlaceGuard {
GPUPlace previous_; GPUPlace previous_;
}; };
template <> class CUDADeviceContext : public DeviceContext {
class Device<DEVICE_GPU> {
public: public:
explicit Device(const GPUPlace gpu_place) : gpu_place_(gpu_place) { explicit Device(const GPUPlace gpu_place) : gpu_place_(gpu_place) {
GPUPlaceGuard guard(gpu_place_); GPUPlaceGuard guard(gpu_place_);
...@@ -61,7 +66,7 @@ class Device<DEVICE_GPU> { ...@@ -61,7 +66,7 @@ class Device<DEVICE_GPU> {
cudaStream_t stream() { return stream_; } cudaStream_t stream() { return stream_; }
DEVICE_GPU eigen_device() { return *eigen_device_; } Eigen::GpuDevice eigen_device() { return *eigen_device_; }
cublasHandle_t cublas_handle() { cublasHandle_t cublas_handle() {
if (!blas_handle_) { if (!blas_handle_) {
...@@ -139,7 +144,7 @@ class Device<DEVICE_GPU> { ...@@ -139,7 +144,7 @@ class Device<DEVICE_GPU> {
cudaStream_t stream_; cudaStream_t stream_;
Eigen::CudaStreamDevice* eigen_stream_; Eigen::CudaStreamDevice* eigen_stream_;
DEVICE_GPU* eigen_device_; Eigen::GpuDevice* eigen_device_;
cublasHandle_t blas_handle_{nullptr}; 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 ...@@ -13,23 +13,39 @@ See the License for the specific language governing permissions and
limitations under the License. */ limitations under the License. */
#pragma once #pragma once
#include "paddle/framework/enforce.h"
#include "paddle/platform/device.h"
#include "unsupported/Eigen/CXX11/Tensor" #include "unsupported/Eigen/CXX11/Tensor"
#ifndef PADDLE_ONLY_CPU
#include "paddle/platform/cuda_device.h" using DEVICE_CPU = Eigen::DefaultDevice;
#endif
namespace paddle { namespace paddle {
namespace platform { namespace platform {
struct DeviceContext { class CPUDeviceContext;
void* device_context{nullptr};
class DeviceContext {
public:
virtual ~DeviceContext() {}
template <typename DeviceType> template <typename DeviceType>
inline paddle::platform::Device<DeviceType>* device_context() { DeviceType get_eigen_device();
return static_cast<paddle::platform::Device<DeviceType>*>(device_context); };
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 } // namespace platform
......
...@@ -12,19 +12,34 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. ...@@ -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 See the License for the specific language governing permissions and
limitations under the License. */ limitations under the License. */
#include "paddle/platform/device_context.h"
#include "gtest/gtest.h" #include "gtest/gtest.h"
#include "paddle/platform/cuda_device.h"
TEST(DeviceContext, Init) { TEST(Device, Init) {
int count = paddle::platform::GetDeviceCount(); int count = paddle::platform::GetDeviceCount();
for (int i = 0; i < count; i++) { for (int i = 0; i < count; i++) {
paddle::platform::Device<DEVICE_GPU>* device = paddle::platform::DeviceContext* device_context =
new paddle::platform::Device<DEVICE_GPU>(i); new paddle::platform::CUDADeviceContext(i);
paddle::platform::DeviceContext context;
context.device_context = device;
Eigen::GpuDevice gpu_device = Eigen::GpuDevice gpu_device =
context.device_context<DEVICE_GPU>->eigen_device(); device_context->get_eigen_device<DEVICE_GPU>();
ASSERT_NE(nullptr, gpu_device.stream()); ASSERT_NE(nullptr, gpu_device.stream());
delete device; delete device_context;
} }
} }
\ No newline at end of file
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.
先完成此消息的编辑!
想要评论请 注册