Added EagerUtils to Eager Dygraph (#37479)

* Added EagerUtils to Eager Dygraph

* Purified include dependencies for global_utils

* Fixed merge conflicts

paddle/fluid/eager/CMakeLists.txt

-add_subdirectory(tests)
 add_subdirectory(api)
+add_subdirectory(tests)
 cc_library(grad_node_info SRCS grad_node_info.cc DEPS pten pten_api)
 cc_library(autograd_meta SRCS autograd_meta.cc DEPS pten pten_api)
-cc_library(utils SRCS utils.cc DEPS pten pten_api autograd_meta eager_api)
+cc_library(utils SRCS utils.cc DEPS pten pten_api global_utils layer proto_desc operator op_registry variable_helper memcpy scale_op autograd_meta)

paddle/fluid/eager/api/CMakeLists.txt

 add_subdirectory(utils)
-
-cc_library(eager_api SRCS all.cc DEPS global_utils)

paddle/fluid/eager/api/utils/CMakeLists.txt

-cc_library(global_utils SRCS global_utils.cc DEPS enforce)
+cc_library(global_utils SRCS global_utils.cc DEPS place)

paddle/fluid/eager/api/utils/global_utils.h

@@ -15,8 +15,9 @@
 
 #pragma once
 
-#include "paddle/fluid/eager/eager_tensor.h"
-#include "paddle/fluid/platform/enforce.h"
+#include <atomic>
+#include <memory>
+#include "paddle/fluid/platform/place.h"
 
 namespace egr {
 

paddle/fluid/eager/tests/CMakeLists.txt

-set(eager_deps pten pten_api)
+set(eager_deps pten pten_api pten_tensor utils global_utils autograd_meta grad_node_info)
 add_subdirectory(data_structure_tests)
 add_subdirectory(task_tests)

paddle/fluid/eager/tests/task_tests/CMakeLists.txt

-cc_test(test_egr_task_eager_utils SRCS eager_utils_test.cc DEPS ${eager_deps} grad_node_info autograd_meta utils)
+cc_test(test_egr_task_eager_utils SRCS eager_utils_test.cc DEPS ${eager_deps})

paddle/fluid/eager/tests/task_tests/eager_utils_test.cc

@@ -12,17 +12,96 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Eager Dygraph
+#include <sstream>
 
 #include "gtest/gtest.h"
 
-#include "paddle/fluid/eager/autograd_meta.h"
+#include "paddle/fluid/eager/eager_tensor.h"
 #include "paddle/fluid/eager/grad_node_info.h"
 #include "paddle/fluid/eager/tests/data_structure_tests/grad_node_test.h"
 #include "paddle/fluid/eager/tests/test_utils.h"
 #include "paddle/fluid/eager/utils.h"
 
+#include "paddle/pten/api/lib/utils/allocator.h"
+
+// TODO(jiabin): remove nolint here!!!
+using namespace egr;  // NOLINT
+
 namespace eager_test {
+
+TEST(EagerUtils, AutoGradMeta) {
+  // Construct Eager Tensor
+  pten::DenseTensorMeta meta = pten::DenseTensorMeta(
+      pten::DataType::FLOAT32, paddle::framework::make_ddim({1, 1}));
+  std::shared_ptr<pten::DenseTensor> dt0 = std::make_shared<pten::DenseTensor>(
+      std::make_shared<paddle::experimental::DefaultAllocator>(
+          paddle::platform::CPUPlace()),
+      meta);
+  dt0->mutable_data<float>()[0] = 10.0;
+  EagerTensor et0 = EagerTensor(dt0);
+
+  std::shared_ptr<pten::DenseTensor> dt1 = std::make_shared<pten::DenseTensor>(
+      std::make_shared<paddle::experimental::DefaultAllocator>(
+          paddle::platform::CPUPlace()),
+      meta);
+  dt1->mutable_data<float>()[0] = 20.0;
+  EagerTensor et1 = EagerTensor(dt1);
+
+  std::vector<EagerTensor> ets = {et0, et1};
+  auto test_node = std::make_shared<eager_test::GradTestNode>();
+
+  // unsafe_autograd_meta()
+  // autograd_meta()
+  // multi_autograd_meta()
+  AutogradMeta* autograd_meta0 = EagerUtils::autograd_meta(&et0);
+  AutogradMeta* autograd_meta1 = EagerUtils::autograd_meta(&et1);
+
+  AutogradMeta* unsafe_autograd_meta_after =
+      EagerUtils::unsafe_autograd_meta(et0);
+  CHECK_NOTNULL(unsafe_autograd_meta_after);
+
+  std::vector<AutogradMeta*> autograd_metas =
+      EagerUtils::multi_autograd_meta(&ets);
+  std::vector<AutogradMeta*> unsafe_autograd_metas =
+      EagerUtils::unsafe_autograd_meta(&ets);
+  CHECK_NOTNULL(unsafe_autograd_metas[0]);
+  CHECK_NOTNULL(unsafe_autograd_metas[1]);
+
+  // Set Autograd Meta
+  autograd_meta0->SetSingleOutRankWithSlot(0, 1);
+
+  autograd_meta0->SetGradNode(test_node);
+
+  // OutRankInfo()
+  std::pair<size_t, size_t> out_rank_info0 = EagerUtils::OutRankInfo(et0);
+  CHECK_EQ(static_cast<int>(out_rank_info0.first), 0);
+  CHECK_EQ(static_cast<int>(out_rank_info0.second), 1);
+
+  // grad_node()
+  std::shared_ptr<GradNodeBase> grad_node0 = EagerUtils::grad_node(et0);
+  CHECK_NOTNULL(grad_node0.get());
+
+  EagerUtils::SetHistory(autograd_meta1, test_node);
+  EagerUtils::SetHistory({autograd_meta1}, test_node);
+  std::shared_ptr<GradNodeBase> grad_node1 = EagerUtils::grad_node(et1);
+  CHECK_NOTNULL(grad_node1.get());
+
+  // SetOutRankWithSlot()
+  EagerUtils::SetOutRankWithSlot(autograd_meta1, 0);
+  std::pair<size_t, size_t> out_rank_info1 = EagerUtils::OutRankInfo(et1);
+  CHECK_EQ(static_cast<int>(out_rank_info1.first), 0);
+  CHECK_EQ(static_cast<int>(out_rank_info1.second), 0);
+
+  EagerUtils::SetOutRankWithSlot(&autograd_metas, 0);
+  std::pair<size_t, size_t> out_rank_info2 = EagerUtils::OutRankInfo(et0);
+  CHECK_EQ(static_cast<int>(out_rank_info2.first), 0);
+  CHECK_EQ(static_cast<int>(out_rank_info2.second), 0);
+
+  std::pair<size_t, size_t> out_rank_info3 = EagerUtils::OutRankInfo(et1);
+  CHECK_EQ(static_cast<int>(out_rank_info3.first), 0);
+  CHECK_EQ(static_cast<int>(out_rank_info3.second), 1);
+}
+
 template <typename T>
 egr::EagerTensor CreateTestCPUTensor(T val,
                                      const paddle::framework::DDim& ddim) {
@@ -40,7 +119,7 @@ egr::EagerTensor CreateTestCPUTensor(T val,
   tensor.set_impl(dt);
   return tensor;
 }
-}  // namespace eager_test
+
 TEST(EagerUtils, ComputeRequireGrad) {
   auto auto_grad0 = std::make_shared<egr::AutogradMeta>();
   auto auto_grad1 = std::make_shared<egr::AutogradMeta>();
@@ -200,3 +279,5 @@ TEST(EagerUtils, ConstructDuplicableOutput) {
   CHECK(outs[0]->defined() == false);
   CHECK(outs[0]->initialized() == false);
 }
+
+}  // namespace eager_test

paddle/fluid/eager/utils.cc

@@ -13,9 +13,94 @@
 // limitations under the License.
 
 #include "paddle/fluid/eager/utils.h"
-#include "paddle/fluid/eager/api/all.h"
+#include "paddle/fluid/eager/api/utils/global_utils.h"
+
+#include "paddle/pten/api/all.h"
+#include "paddle/pten/common/layout.h"
+#include "paddle/pten/core/tensor_meta.h"
+
+#include "paddle/fluid/framework/data_layout.h"
+#include "paddle/fluid/framework/pten_utils.h"
+#include "paddle/fluid/framework/variable.h"
 
 namespace egr {
+/**
+ * Implementation of Eager Utils.
+**/
+
+AutogradMeta* EagerUtils::autograd_meta(egr::EagerTensor* target) {
+  auto* p_autograd_meta = target->get_autograd_meta();
+  if (!p_autograd_meta) {
+    auto p_autograd_meta_ptr = std::make_shared<AutogradMeta>();
+    p_autograd_meta = p_autograd_meta_ptr.get();
+    target->set_autograd_meta(p_autograd_meta_ptr);
+  }
+  return static_cast<AutogradMeta*>(p_autograd_meta);
+}
+
+AutogradMeta* EagerUtils::unsafe_autograd_meta(const egr::EagerTensor& target) {
+  auto* p_autograd_meta = target.get_autograd_meta();
+  PADDLE_ENFORCE(p_autograd_meta,
+                 paddle::platform::errors::Fatal(
+                     "Null autograd_meta gotten from unsafe_autograd_meta()"));
+  return static_cast<AutogradMeta*>(p_autograd_meta);
+}
+
+std::vector<AutogradMeta*> EagerUtils::unsafe_autograd_meta(
+    std::vector<egr::EagerTensor>* targets) {
+  std::vector<AutogradMeta*> metas;
+  for (const egr::EagerTensor& t : *targets) {
+    metas.push_back(unsafe_autograd_meta(t));
+  }
+  return metas;
+}
+
+std::vector<AutogradMeta*> EagerUtils::multi_autograd_meta(
+    std::vector<egr::EagerTensor>* targets) {
+  std::vector<AutogradMeta*> ret;
+  ret.reserve(targets->size());
+
+  // for multi_autograd_meta we can tolerent it has nullptr.
+  for (auto& t : (*targets)) {
+    auto* p_autograd_meta = autograd_meta(&t);
+    ret.push_back(static_cast<AutogradMeta*>(p_autograd_meta));
+  }
+  return ret;
+}
+
+std::pair<size_t, size_t> EagerUtils::OutRankInfo(
+    const egr::EagerTensor& target) {
+  return unsafe_autograd_meta(target)->OutRankInfo();
+}
+
+std::shared_ptr<GradNodeBase> EagerUtils::grad_node(
+    const egr::EagerTensor& target) {
+  return unsafe_autograd_meta(target)->GetMutableGradNode();
+}
+
+void EagerUtils::SetHistory(std::vector<AutogradMeta*>* autograd_metas,
+                            const std::shared_ptr<GradNodeBase>& grad_node) {
+  for (const auto& autograd_meta : *autograd_metas) {
+    autograd_meta->SetGradNode(grad_node);
+  }
+}
+
+void EagerUtils::SetHistory(AutogradMeta* autograd_meta,
+                            const std::shared_ptr<GradNodeBase>& grad_node) {
+  autograd_meta->SetGradNode(grad_node);
+}
+
+void EagerUtils::SetOutRankWithSlot(std::vector<AutogradMeta*>* targets,
+                                    size_t slot_id) {
+  // Set OutRankInfo from 0 to size of targets
+  for (size_t i = 0; i < targets->size(); i++) {
+    (*targets)[i]->SetSingleOutRankWithSlot(slot_id, i);
+  }
+}
+void EagerUtils::SetOutRankWithSlot(AutogradMeta* target, size_t slot_id) {
+  target->SetSingleOutRankWithSlot(slot_id, 0);
+}
+
 /* ---- Tensor -> Var ---- */
 std::vector<std::shared_ptr<egr::EagerTensor>> EagerUtils::SyncToVars(
     const egr::EagerTensor& tensor) {
@@ -103,18 +188,4 @@ egr::EagerTensor EagerUtils::GetOutput(
   return EagerTensor((*(out.get())));
 }
 
-AutogradMeta* EagerUtils::unsafe_autograd_meta(const egr::EagerTensor& target) {
-  auto* p_autograd_meta = target.get_autograd_meta();
-  PADDLE_ENFORCE(p_autograd_meta,
-                 paddle::platform::errors::Fatal(
-                     "Null autograd_meta gotten from unsafe_autograd_meta(), "
-                     "if you are using unsafe_autograd_meta, please make sure "
-                     "your tensor's autograd_meta is set"));
-  return static_cast<AutogradMeta*>(p_autograd_meta);
-}
-
-std::pair<size_t, size_t> EagerUtils::OutRankInfo(
-    const egr::EagerTensor& target) {
-  return unsafe_autograd_meta(target)->OutRankInfo();
-}
 }  // namespace egr

paddle/fluid/eager/utils.h

@@ -87,6 +87,33 @@ class EagerUtils {
    * constructor (it's abstract class there)
    *
    * **/
+  static AutogradMeta* autograd_meta(egr::EagerTensor* target);
+
+  static std::vector<AutogradMeta*> multi_autograd_meta(
+      std::vector<egr::EagerTensor>* targets);
+
+  static std::pair<size_t, size_t> OutRankInfo(const egr::EagerTensor& target);
+
+  static std::shared_ptr<GradNodeBase> grad_node(
+      const egr::EagerTensor& target);
+
+  // Set history is used to set backward info during forward process, it will
+  // set forward var's autograd meta's grad node as current backward node.
+  static void SetHistory(std::vector<AutogradMeta*>* autograd_metas,
+                         const std::shared_ptr<GradNodeBase>& grad_node);
+  static void SetHistory(AutogradMeta* autograd_meta,
+                         const std::shared_ptr<GradNodeBase>& grad_node);
+
+  // This is used for Set vector of tensors' rank
+  static void SetOutRankWithSlot(std::vector<AutogradMeta*>* targets,
+                                 size_t slot_id);
+  static void SetOutRankWithSlot(AutogradMeta* target, size_t slot_id);
+
+  // This method will return an AutogradMeta pointer unsafely.
+  static AutogradMeta* unsafe_autograd_meta(const egr::EagerTensor& target);
+  static std::vector<AutogradMeta*> unsafe_autograd_meta(
+      std::vector<egr::EagerTensor>* targets);
+
   template <typename T, typename... Args>
   static bool ComputeRequireGrad(T trace_backward, Args&&... args) {
     if (!trace_backward) return false;
@@ -103,9 +130,6 @@ class EagerUtils {
     iter.SetStopGradient(stop_gradient);
     iter.apply(std::forward<Args>(args)...);
   }
-  static std::pair<size_t, size_t> OutRankInfo(const egr::EagerTensor& target);
-  // This method will return an AutogradMeta pointer unsafely.
-  static AutogradMeta* unsafe_autograd_meta(const egr::EagerTensor& target);
 
   // Intermidate needed remove this once we don't need legacy
   static std::vector<std::shared_ptr<egr::EagerTensor>> SyncToVars(