follow comments

paddle/platform/cuda_device.h → paddle/platform/cuda_device_context.h

@@ -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};
 

paddle/platform/cuda_device_test.cc

-/* 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;
-  }
-}

paddle/platform/device.h

-/* 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

paddle/platform/device_context.h

@@ -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

paddle/platform/device_context_test.cc

@@ -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;
   }
-}
+}
\ 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;
+  }
+}