From c774bcbd2d80c4bd3d4f0560a2a804d4236bce09 Mon Sep 17 00:00:00 2001
From: Yu Yang <yuyang18@baidu.com>
Date: Wed, 7 Nov 2018 16:11:49 +0800
Subject: [PATCH] Merge device_context

test=develop
---
 paddle/fluid/platform/device_context.cc | 13 +++++--------
 paddle/fluid/platform/device_context.h  | 25 ++++++++++++++++++++-----
 2 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/paddle/fluid/platform/device_context.cc b/paddle/fluid/platform/device_context.cc
index 36e7f2934..018e9d19b 100644
--- a/paddle/fluid/platform/device_context.cc
+++ b/paddle/fluid/platform/device_context.cc
@@ -160,29 +160,26 @@ class EigenCudaStreamDevice : public Eigen::StreamInterface {
 };
 
 CudnnHolder::CudnnHolder(const cudaStream_t* stream, const CUDAPlace& place)
-    : workspace_(nullptr), workspace_len_(0), stream_(stream), place_(place) {
+    : workspace_(nullptr), stream_(stream), place_(place) {
   PADDLE_ENFORCE(dynload::cudnnCreate(&cudnn_handle_));
   PADDLE_ENFORCE(dynload::cudnnSetStream(cudnn_handle_, *stream_));
 }
 
 CudnnHolder::~CudnnHolder() {
   PADDLE_ENFORCE(dynload::cudnnDestroy(cudnn_handle_));
-  if (workspace_ != nullptr) {
-    paddle::memory::Free(place_, workspace_);
-  }
 }
 
 void CudnnHolder::ReallocateWorkspace(size_t required_workspace_len) {
-  if (required_workspace_len <= workspace_len_) {
+  if (required_workspace_len <= WorkspaceSize()) {
     return;
   }
   if (workspace_ != nullptr) {
     // Maybe someone is using the current workspace
     PADDLE_ENFORCE(cudaStreamSynchronize(*stream_));
-    paddle::memory::Free(place_, workspace_);
+    workspace_.reset();
   }
-  workspace_ = paddle::memory::Alloc(place_, required_workspace_len);
-  workspace_len_ = required_workspace_len;
+  workspace_ = paddle::memory::Alloc(place_, required_workspace_len,
+                                     paddle::memory::Allocator::kScratchpad);
 }
 
 CUDADeviceContext::CUDADeviceContext(CUDAPlace place)
diff --git a/paddle/fluid/platform/device_context.h b/paddle/fluid/platform/device_context.h
index df248f9bb..0e7799833 100644
--- a/paddle/fluid/platform/device_context.h
+++ b/paddle/fluid/platform/device_context.h
@@ -16,7 +16,7 @@ limitations under the License. */
 #include <string>
 #include <unordered_map>
 #include <vector>
-
+#include "paddle/fluid/memory/malloc.h"
 #ifdef PADDLE_WITH_CUDA
 #include "paddle/fluid/platform/dynload/cublas.h"
 #include "paddle/fluid/platform/dynload/cudnn.h"
@@ -85,17 +85,32 @@ class CudnnHolder {
 
   template <typename Callback>
   void RunFuncImpl(Callback&& cudnn_func, size_t required_workspace_len) {
-    if (required_workspace_len > workspace_len_) {
+    if (required_workspace_len > WorkspaceSize()) {
       ReallocateWorkspace(required_workspace_len);
     }
-    cudnn_func(workspace_);
+    cudnn_func(WorkspacePtr());
+  }
+
+  inline void* WorkspacePtr() {
+    if (workspace_) {
+      return workspace_->ptr();
+    } else {
+      return nullptr;
+    }
+  }
+
+  inline size_t WorkspaceSize() {
+    if (workspace_) {
+      return workspace_->size();
+    } else {
+      return 0;
+    }
   }
 
   std::mutex& Mutex() { return mtx_; }
 
   cudnnHandle_t cudnn_handle_;
-  void* workspace_;
-  size_t workspace_len_;
+  std::unique_ptr<memory::Allocation> workspace_;
 
   const cudaStream_t* stream_;  // not owned;
   const CUDAPlace place_;
-- 
GitLab