From e1a792fe1565f825c8e4cf77592a7ab96a345b37 Mon Sep 17 00:00:00 2001
From: HongyuJia <jiahongyu@baidu.com>
Date: Fri, 10 Feb 2023 15:17:51 +0800
Subject: [PATCH] [Bug Fix] Fix NLP-Bert model performance loss (#50333)

* fix NLP-Bert model performance loss

* fix windows compile error
---
 paddle/fluid/framework/op_kernel_type.h     | 31 ++++++++++++++++++---
 paddle/fluid/imperative/prepared_operator.h |  4 +--
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/paddle/fluid/framework/op_kernel_type.h b/paddle/fluid/framework/op_kernel_type.h
index eb969a94d82..43b383aecb0 100644
--- a/paddle/fluid/framework/op_kernel_type.h
+++ b/paddle/fluid/framework/op_kernel_type.h
@@ -20,7 +20,9 @@ limitations under the License. */
 #include "paddle/fluid/framework/data_type.h"
 #include "paddle/fluid/framework/library_type.h"
 #include "paddle/fluid/platform/place.h"
+#include "paddle/phi/core/dense_tensor.h"
 #include "paddle/phi/core/device_context.h"
+#include "paddle/phi/core/enforce.h"
 #include "paddle/phi/core/kernel_factory.h"
 
 namespace paddle {
@@ -131,10 +133,31 @@ inline bool backends_are_same_class(const phi::Backend& l,
   return phi::TransToPhiPlace(l) == phi::TransToPhiPlace(r);
 }
 
-inline bool NeedTransform(const phi::KernelKey& l, const phi::KernelKey& r) {
-  return !backends_are_same_class(l.backend(), r.backend()) ||
-         NeedTransformDataType(l, r) ||
-         NeedTransformLayout(l.layout(), r.layout());
+inline bool NeedTransformBackend(const phi::Backend& type_for_var_backend,
+                                 const phi::Backend& expected_backend,
+                                 const phi::DenseTensor& tensor) {
+  // NOTE(jiahongyu): KernelKey does not hold place information, so we need to
+  // explicitly transform CUDAPinnedPlace->CUDAPlace
+  if (type_for_var_backend != phi::Backend::ALL_BACKEND &&
+      paddle::platform::is_cuda_pinned_place(tensor.place()) &&
+      expected_backend != phi::Backend::CPU) {
+    VLOG(3) << "Transform Variable " << tensor.name() << " from "
+            << tensor.place() << " to "
+            << phi::TransToPhiPlace(expected_backend);
+    return true;
+  }
+  return !backends_are_same_class(type_for_var_backend, expected_backend);
+}
+
+inline bool NeedTransform(const phi::KernelKey& kernel_type_for_var,
+                          const phi::KernelKey& expected_kernel_key,
+                          const phi::DenseTensor& tensor) {
+  return NeedTransformBackend(kernel_type_for_var.backend(),
+                              expected_kernel_key.backend(),
+                              tensor) ||
+         NeedTransformDataType(kernel_type_for_var, expected_kernel_key) ||
+         NeedTransformLayout(kernel_type_for_var.layout(),
+                             expected_kernel_key.layout());
 }
 
 }  // namespace framework
diff --git a/paddle/fluid/imperative/prepared_operator.h b/paddle/fluid/imperative/prepared_operator.h
index fb36a03e018..00e059572d2 100644
--- a/paddle/fluid/imperative/prepared_operator.h
+++ b/paddle/fluid/imperative/prepared_operator.h
@@ -87,8 +87,8 @@ std::shared_ptr<NameVarMap<VarType>> PrepareData(
       if (tensor && tensor->IsInitialized() && (tensor->memory_size() != 0)) {
         auto kernel_type_for_var = op.GetKernelTypeForVar(
             name_pair.first, *tensor, expected_kernel_key);
-        if (!framework::NeedTransform(kernel_type_for_var,
-                                      expected_kernel_key)) {
+        if (!framework::NeedTransform(
+                kernel_type_for_var, expected_kernel_key, *tensor)) {
           continue;
         } else {
           VLOG(3) << "Transform Variable " << GetNameFromVar(template_var)
-- 
GitLab