From 95122ebe8664a8d93558f07bf8c42c198a4b4653 Mon Sep 17 00:00:00 2001
From: Yiqun Liu <liuyiqun01@baidu.com>
Date: Fri, 9 Apr 2021 13:52:27 +0800
Subject: [PATCH] Advoid CPU -> CPU memory copy when start, end, step is
 already on CPU. (#29088)

---
 paddle/fluid/operators/range_op.cu | 24 ++++--------------------
 paddle/fluid/operators/utils.h     | 13 +++++++++++++
 2 files changed, 17 insertions(+), 20 deletions(-)

diff --git a/paddle/fluid/operators/range_op.cu b/paddle/fluid/operators/range_op.cu
index f2c78e0f70b..6250d68730e 100644
--- a/paddle/fluid/operators/range_op.cu
+++ b/paddle/fluid/operators/range_op.cu
@@ -15,6 +15,7 @@ limitations under the License. */
 #include <algorithm>
 #include "paddle/fluid/framework/op_registry.h"
 #include "paddle/fluid/operators/range_op.h"
+#include "paddle/fluid/operators/utils.h"
 #include "paddle/fluid/platform/cuda_primitives.h"
 
 namespace paddle {
@@ -34,26 +35,9 @@ class CUDARangeKernel : public framework::OpKernel<T> {
     auto* step_t = context.Input<framework::Tensor>("Step");
     auto* out = context.Output<framework::Tensor>("Out");
 
-    T start, end, step;
-    framework::Tensor n;
-    if (::paddle::platform::is_cpu_place(start_t->place())) {
-      start = start_t->data<T>()[0];
-    } else {
-      framework::TensorCopy(*start_t, platform::CPUPlace(), &n);
-      start = n.data<T>()[0];
-    }
-    if (::paddle::platform::is_cpu_place(end_t->place())) {
-      end = end_t->data<T>()[0];
-    } else {
-      framework::TensorCopy(*end_t, platform::CPUPlace(), &n);
-      end = n.data<T>()[0];
-    }
-    if (::paddle::platform::is_cpu_place(step_t->place())) {
-      step = step_t->data<T>()[0];
-    } else {
-      framework::TensorCopy(*step_t, platform::CPUPlace(), &n);
-      step = n.data<T>()[0];
-    }
+    T start = GetValue<T>(start_t);
+    T end = GetValue<T>(end_t);
+    T step = GetValue<T>(step_t);
 
     int64_t size = 0;
     GetSize(start, end, step, &size);
diff --git a/paddle/fluid/operators/utils.h b/paddle/fluid/operators/utils.h
index 985c3512761..912d538d5e9 100644
--- a/paddle/fluid/operators/utils.h
+++ b/paddle/fluid/operators/utils.h
@@ -108,5 +108,18 @@ inline framework::DDim GetShape(const framework::ExecutionContext& ctx) {
   return framework::make_ddim(vec_shape);
 }
 
+template <typename T>
+inline T GetValue(const framework::Tensor* x) {
+  T value = static_cast<T>(0);
+  if (!platform::is_cpu_place(x->place())) {
+    framework::Tensor cpu_x;
+    framework::TensorCopy(*x, platform::CPUPlace(), &cpu_x);
+    value = cpu_x.data<T>()[0];
+  } else {
+    value = x->data<T>()[0];
+  }
+  return value;
+}
+
 }  // namespace operators
 }  // namespace paddle
-- 
GitLab