From 7d143a48ab2fd6caab01cd6883d78c66dc739816 Mon Sep 17 00:00:00 2001
From: pangyoki <pangyoki@126.com>
Date: Thu, 7 Apr 2022 18:30:56 +0800
Subject: [PATCH] Cherry-pick-PR41407, fix device_id bug for final_state op in
 multiprocess testcase (#41407) (#41475)

Cherry-pick PR #41407
---
 .../final_state_generator/eager_gen.py                | 11 +++++++++++
 python/paddle/fluid/dygraph/math_op_patch.py          |  5 ++++-
 python/paddle/fluid/tests/unittests/test_inplace.py   |  4 +---
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/paddle/fluid/eager/auto_code_generator/final_state_generator/eager_gen.py b/paddle/fluid/eager/auto_code_generator/final_state_generator/eager_gen.py
index 19e42e1bdf..3b4c8f9621 100644
--- a/paddle/fluid/eager/auto_code_generator/final_state_generator/eager_gen.py
+++ b/paddle/fluid/eager/auto_code_generator/final_state_generator/eager_gen.py
@@ -194,6 +194,16 @@ FORWARD_FUNCTION_TEMPLATE = \
     
     // Get Input AutoGradMeta
 {}
+    // Set Device Id
+    auto place = egr::Controller::Instance().GetExpectedPlace();
+    if (paddle::platform::is_gpu_place(place)) {{
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+      phi::backends::gpu::SetDeviceId(place.device);
+#else
+      PADDLE_THROW(paddle::platform::errors::PreconditionNotMet(
+        "PaddlePaddle should compile with GPU if use CUDAPlace."));
+#endif
+    }}
     // Forward API Call
 {}
     // Get Outputs
@@ -284,6 +294,7 @@ FORWARD_CC_FILE_TEMPLATE = \
 #include "paddle/fluid/platform/profiler/event_tracing.h"
 #include "paddle/fluid/eager/amp_utils.h"
 #include "paddle/fluid/eager/eager_amp_auto_cast.h"
+#include "paddle/phi/backends/gpu/gpu_info.h"
 
 {}
 {}
diff --git a/python/paddle/fluid/dygraph/math_op_patch.py b/python/paddle/fluid/dygraph/math_op_patch.py
index 98de5949ba..8b80444fe9 100644
--- a/python/paddle/fluid/dygraph/math_op_patch.py
+++ b/python/paddle/fluid/dygraph/math_op_patch.py
@@ -270,7 +270,10 @@ def monkey_patch_math_varbase():
 
             # 4. calculation
             axis = -1
-            math_op = getattr(_C_ops, op_type)
+            if framework._in_eager_mode_ and op_type == 'elementwise_add':
+                math_op = getattr(_C_ops, 'final_state_add')
+            else:
+                math_op = getattr(_C_ops, op_type)
             return math_op(self, other_var, 'axis', axis)
 
         comment = OpProtoHolder.instance().get_op_proto(op_type).comment
diff --git a/python/paddle/fluid/tests/unittests/test_inplace.py b/python/paddle/fluid/tests/unittests/test_inplace.py
index ee0d5bcdde..c54d3f02d4 100644
--- a/python/paddle/fluid/tests/unittests/test_inplace.py
+++ b/python/paddle/fluid/tests/unittests/test_inplace.py
@@ -103,9 +103,7 @@ class TestInplace(unittest.TestCase):
 
             var_b[1:2] = 3  # var_b is modified inplace before using it
 
-            var_c = paddle.add(
-                var_b,
-                var_b)  # Here, the grad op of sum doesn't use the value of var_b
+            var_c = var_b + var_b  # Here, the grad op of sum doesn't use the value of var_b
             loss = var_c.sum()
 
             var_b[1:2] = 3  # var_b is modified inplace after using it
-- 
GitLab