From 0f1fde51021e1c9deae099ee0c875c53128687b4 Mon Sep 17 00:00:00 2001
From: Leo Chen <chenqiuliang@baidu.com>
Date: Wed, 24 Feb 2021 19:28:52 +0800
Subject: [PATCH] fix the modification of set_expected_place (#31177)

* revert the modification of set_expected_place

* set device before op run

* add ut
---
 paddle/fluid/imperative/tests/test_tracer.cc | 15 +++++----
 paddle/fluid/imperative/tracer.cc            | 33 ++++++++++----------
 2 files changed, 24 insertions(+), 24 deletions(-)

diff --git a/paddle/fluid/imperative/tests/test_tracer.cc b/paddle/fluid/imperative/tests/test_tracer.cc
index e3b5ff67036..9e3b0ea5df6 100644
--- a/paddle/fluid/imperative/tests/test_tracer.cc
+++ b/paddle/fluid/imperative/tests/test_tracer.cc
@@ -72,6 +72,13 @@ TEST(test_tracer, test_trace_op) {
   framework::AttributeMap mul_attr_map;
   mul_attr_map["use_mkldnn"] = false;
   tracer.TraceOp("mul", ins, outs, mul_attr_map, place, true);
+
+#ifndef PADDLE_WITH_XPU
+  ASSERT_THROW(tracer.TraceOp("mul", ins, outs, mul_attr_map,
+                              platform::XPUPlace(0), true);
+               , platform::EnforceNotMet);
+#endif
+
   const auto& out_tensor = vout->Var().Get<framework::LoDTensor>();
   for (int i = 0; i < vout->Var().Get<framework::LoDTensor>().numel(); i++) {
     ASSERT_EQ(out_tensor.data<float>()[i], 20.0);
@@ -311,10 +318,6 @@ TEST(test_tracer, test_expected_place) {
     platform::CUDAPlace gpu_place(0);
     tracer.SetExpectedPlace(gpu_place);
     ASSERT_EQ(platform::is_gpu_place(tracer.ExpectedPlace()), true);
-
-    // assert throw
-    platform::XPUPlace xpu_place(0);
-    ASSERT_THROW(tracer.SetExpectedPlace(xpu_place), platform::EnforceNotMet);
 #endif
   }
   {
@@ -323,10 +326,6 @@ TEST(test_tracer, test_expected_place) {
     platform::XPUPlace xpu_place(0);
     tracer.SetExpectedPlace(xpu_place);
     ASSERT_EQ(platform::is_xpu_place(tracer.ExpectedPlace()), true);
-
-    // assert throw
-    platform::CUDAPlace cuda_place(0);
-    ASSERT_THROW(tracer.SetExpectedPlace(cuda_place), platform::EnforceNotMet);
 #endif
   }
 }
diff --git a/paddle/fluid/imperative/tracer.cc b/paddle/fluid/imperative/tracer.cc
index 3c20c1f647a..608cc407d5b 100644
--- a/paddle/fluid/imperative/tracer.cc
+++ b/paddle/fluid/imperative/tracer.cc
@@ -162,6 +162,23 @@ void Tracer::TraceOp(const std::string& type, const NameVarBaseMap& ins,
   }
 
   try {
+    if (platform::is_gpu_place(place)) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+      platform::SetDeviceId(BOOST_GET_CONST(platform::CUDAPlace, place).device);
+#else
+      PADDLE_THROW(platform::errors::PreconditionNotMet(
+          "PaddlePaddle should compile with GPU if use CUDAPlace."));
+#endif
+    } else if (platform::is_xpu_place(place)) {
+#ifdef PADDLE_WITH_XPU
+      platform::SetXPUDeviceId(
+          BOOST_GET_CONST(platform::XPUPlace, place).device);
+#else
+      PADDLE_THROW(platform::errors::PreconditionNotMet(
+          "PaddlePaddle should compile with XPU if use XPUPlace."));
+#endif
+    }
+
     OpBase::Run(*op, new_ins, outs, attrs, place);
   } catch (platform::EnforceNotMet& exception) {
     framework::AppendErrorOpHint(type, &exception);
@@ -199,22 +216,6 @@ void Tracer::TraceOp(const std::string& type, const NameVarBaseMap& ins,
 }
 
 void Tracer::SetExpectedPlace(platform::Place place) {
-  // NOTE(wangxi): set device id before launch device kernel
-  if (platform::is_gpu_place(place)) {
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-    platform::SetDeviceId(BOOST_GET_CONST(platform::CUDAPlace, place).device);
-#else
-    PADDLE_THROW(platform::errors::PreconditionNotMet(
-        "PaddlePaddle should compile with GPU if use CUDAPlace."));
-#endif
-  } else if (platform::is_xpu_place(place)) {
-#ifdef PADDLE_WITH_XPU
-    platform::SetXPUDeviceId(BOOST_GET_CONST(platform::XPUPlace, place).device);
-#else
-    PADDLE_THROW(platform::errors::PreconditionNotMet(
-        "PaddlePaddle should compile with XPU if use XPUPlace."));
-#endif
-  }
   expected_place_ = place;
 }
 
-- 
GitLab