PaddlePaddle
diff --git a/‎paddle/fluid/eager/to_static/run_program_op_func.h
Copy file name to clipboard
-3Lines changed: 0 additions & 3 deletions b/‎paddle/fluid/eager/to_static/run_program_op_func.h
Copy file name to clipboard
-3Lines changed: 0 additions & 3 deletions
diff --git a/‎paddle/fluid/pybind/eager_utils.cc
Copy file name to clipboard
+198Lines changed: 198 additions & 0 deletions b/‎paddle/fluid/pybind/eager_utils.cc
Copy file name to clipboard
+198Lines changed: 198 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/eager_utils.h
Copy file name to clipboard
+18Lines changed: 18 additions & 0 deletions b/‎paddle/fluid/pybind/eager_utils.h
Copy file name to clipboard
+18Lines changed: 18 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/imperative.cc
Copy file name to clipboard
-10Lines changed: 0 additions & 10 deletions b/‎paddle/fluid/pybind/imperative.cc
Copy file name to clipboard
-10Lines changed: 0 additions & 10 deletions
diff --git a/‎paddle/fluid/pybind/pir.h
Copy file name to clipboard
+7Lines changed: 7 additions & 0 deletions b/‎paddle/fluid/pybind/pir.h
Copy file name to clipboard
+7Lines changed: 7 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/pybind.cc
Copy file name to clipboard
+1-12Lines changed: 1 addition & 12 deletions b/‎paddle/fluid/pybind/pybind.cc
Copy file name to clipboard
+1-12Lines changed: 1 addition & 12 deletions
@@ -165,7 +165,6 @@ inline void run_program_ad_func(
     auto x_names =
         PADDLE_GET_CONST(std::vector<std::string>, attrs.at("x_names"));
 
-    egr::EagerUtils::PassStopGradient(false, &p_autograd_outs);
     // Create GradOpNode (1 means [out_grad], 2 means [x_grad, paramx_grad])
     auto grad_node = std::make_shared<GradNodeRunProgram>(1, 2);
 
@@ -270,8 +269,6 @@ inline void pir_run_program_ad_func(
   PirRunProgramAPI(
       x, params, out, middles, step_scope, require_any_grad, attrs);
   if (!is_test && require_any_grad) {
-    egr::EagerUtils::PassStopGradient(false, &p_autograd_outs);
-
     // Set Attributes
     grad_node->SetAttrMap(attrs);
 
 
@@ -20,6 +20,7 @@ limitations under the License. */
 #include <string>
 #include <vector>
 
+#include "paddle/fluid/eager/accumulation/accumulation_node.h"
 #include "paddle/fluid/eager/api/all.h"
 #include "paddle/fluid/eager/autograd_meta.h"
 #include "paddle/fluid/eager/hooks.h"
@@ -30,9 +31,12 @@ limitations under the License. */
 #include "paddle/fluid/memory/allocation/allocator.h"
 #include "paddle/fluid/operators/py_func_op.h"
 #include "paddle/fluid/operators/utils.h"
+#include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
+#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
 #include "paddle/fluid/platform/enforce.h"
 #include "paddle/fluid/pybind/eager.h"
 #include "paddle/fluid/pybind/op_function_common.h"
+#include "paddle/fluid/pybind/pir.h"
 #include "paddle/fluid/pybind/tensor_py.h"
 #include "paddle/phi/api/ext/op_meta_info.h"
 #include "paddle/phi/common/data_type.h"
@@ -41,6 +45,7 @@ limitations under the License. */
 #include "paddle/phi/core/distributed/auto_parallel/placement_types.h"
 #include "paddle/phi/core/distributed/auto_parallel/process_mesh.h"
 #include "paddle/phi/core/flags.h"
+#include "paddle/pir/core/attribute.h"
 
 PHI_DECLARE_bool(check_nan_inf);
 PHI_DECLARE_int32(check_nan_inf_level);
@@ -1858,6 +1863,180 @@ std::vector<paddle::Tensor> GetTensorListFromPyObject(PyObject* obj,
 paddle::Tensor& UnSafeGetTensorFromPyObject(PyObject* obj) {
   return reinterpret_cast<TensorObject*>(obj)->tensor;
 }
+
+paddle::Tensor CreateTensorFromVarDesc(
+    const paddle::framework::VarDesc& var_desc) {
+  auto tensor = paddle::Tensor();
+
+  auto dtype = var_desc.GetDataType();
+  std::vector<int64_t> dims = var_desc.GetShape();
+
+  auto var_type = var_desc.GetType();
+
+  auto ddims = phi::make_ddim(dims);
+  tensor.set_name(var_desc.Name());
+  auto autograd_meta = egr::EagerUtils::autograd_meta(&tensor);
+  autograd_meta->SetPersistable(false);
+  autograd_meta->SetStopGradient(var_desc.StopGradient());
+
+  if (var_type == paddle::framework::proto::VarType::LOD_TENSOR) {
+    // TODO(jiabin): Maybe support LOD later
+    std::shared_ptr<phi::DenseTensor> dense_tensor = nullptr;
+    if (dims.size() == 1 && dims[0] == 0) {
+      std::shared_ptr<phi::Allocation> allocation_ptr = nullptr;
+      dense_tensor = std::make_shared<phi::DenseTensor>(
+          allocation_ptr,
+          phi::DenseTensorMeta(paddle::framework::TransToPhiDataType(dtype),
+                               ddims));
+    } else {
+      // TODO(dev): we need enhance check for ddims.
+      dense_tensor = std::make_shared<phi::DenseTensor>(
+          std::make_shared<phi::Allocation>(),
+          phi::DenseTensorMeta(paddle::framework::TransToPhiDataType(dtype),
+                               ddims));
+    }
+    tensor.set_impl(dense_tensor);
+  } else if (var_type == paddle::framework::proto::VarType::SELECTED_ROWS) {
+    std::shared_ptr<phi::SelectedRows> selected_rows_tensor =
+        std::make_shared<phi::SelectedRows>();
+    tensor.set_impl(selected_rows_tensor);
+  }
+
+  if (!autograd_meta->GetMutableGradNode()) {
+    autograd_meta->SetGradNode(
+        std::make_shared<egr::GradNodeAccumulation>(autograd_meta));
+  }
+
+  return tensor;
+}
+
+PyObject* GetEmpytyTensorsWithVarDesc(PyObject* self, PyObject* args) {
+  std::vector<paddle::Tensor> result;
+  std::unordered_map<std::string, paddle::Tensor> out_tensor_map;
+
+  auto var_desc_list = PyTuple_GetItem(args, 0);
+
+  if (PyList_Check(var_desc_list)) {
+    Py_ssize_t len = PyList_Size(var_desc_list);
+    for (Py_ssize_t i = 0; i < len; i++) {
+      auto var_desc = PyObjectCast<paddle::framework::VarDesc>(
+          PyList_GetItem(var_desc_list, i));
+      auto var_name = var_desc.Name();
+      if (out_tensor_map.find(var_name) == out_tensor_map.end()) {
+        paddle::Tensor tensor = CreateTensorFromVarDesc(var_desc);
+        out_tensor_map[var_name] = tensor;
+        result.emplace_back(tensor);
+      } else {
+        result.emplace_back(out_tensor_map[var_name]);
+      }
+    }
+  } else if (PyTuple_Check(var_desc_list)) {
+    Py_ssize_t len = PyTuple_Size(var_desc_list);
+    for (Py_ssize_t i = 0; i < len; i++) {
+      auto var_desc = PyObjectCast<paddle::framework::VarDesc>(
+          PyTuple_GetItem(var_desc_list, i));
+      auto var_name = var_desc.Name();
+      if (out_tensor_map.find(var_name) == out_tensor_map.end()) {
+        paddle::Tensor tensor = CreateTensorFromVarDesc(var_desc);
+        out_tensor_map[var_name] = tensor;
+        result.emplace_back(tensor);
+      } else {
+        result.emplace_back(out_tensor_map[var_name]);
+      }
+    }
+  } else if (var_desc_list != Py_None) {
+    PADDLE_THROW(platform::errors::InvalidArgument(
+        "Argument of CreateTensorsWithVarDesc must be list of VarDesc, but got "
+        "%s",
+        (reinterpret_cast<PyTypeObject*>(var_desc_list->ob_type))->tp_name));
+  }
+  return ToPyObject(result);
+}
+
+paddle::Tensor CreateTensorFromOpResult(const pir::OpResult& op_result) {
+  auto tensor = paddle::Tensor();
+
+  auto dims = phi::vectorize(GetOpResultDims(op_result));
+  auto ddims = phi::make_ddim(dims);
+  auto autograd_meta = egr::EagerUtils::autograd_meta(&tensor);
+  autograd_meta->SetPersistable(false);
+  autograd_meta->SetStopGradient(
+      GetOpResultBoolAttr(op_result, kAttrStopGradients));
+
+  if (op_result.type().isa<paddle::dialect::DenseTensorType>()) {
+    // TODO(jiabin): Maybe support LOD later
+    std::shared_ptr<phi::DenseTensor> dense_tensor = nullptr;
+    auto dtype = paddle::dialect::TransToPhiDataType(
+        op_result.type().dyn_cast<paddle::dialect::DenseTensorType>().dtype());
+
+    if (dims.size() == 1 && dims[0] == 0) {
+      std::shared_ptr<phi::Allocation> allocation_ptr = nullptr;
+      dense_tensor = std::make_shared<phi::DenseTensor>(
+          allocation_ptr, phi::DenseTensorMeta(dtype, ddims));
+    } else {
+      // TODO(dev): we need enhance check for ddims.
+      dense_tensor = std::make_shared<phi::DenseTensor>(
+          std::make_shared<phi::Allocation>(),
+          phi::DenseTensorMeta(dtype, ddims));
+    }
+    tensor.set_impl(dense_tensor);
+  } else if (op_result.type().isa<paddle::dialect::SelectedRowsType>()) {
+    std::shared_ptr<phi::SelectedRows> selected_rows_tensor =
+        std::make_shared<phi::SelectedRows>();
+    tensor.set_impl(selected_rows_tensor);
+  }
+
+  if (!autograd_meta->GetMutableGradNode()) {
+    autograd_meta->SetGradNode(
+        std::make_shared<egr::GradNodeAccumulation>(autograd_meta));
+  }
+
+  return tensor;
+}
+
+PyObject* GetEmpytyTensorsWithOpResult(PyObject* self, PyObject* args) {
+  std::vector<paddle::Tensor> result;
+  std::unordered_map<pir::OpResult, paddle::Tensor> out_tensor_map;
+
+  auto op_result_list = PyTuple_GetItem(args, 0);
+
+  if (PyList_Check(op_result_list)) {
+    Py_ssize_t len = PyList_Size(op_result_list);
+    for (Py_ssize_t i = 0; i < len; i++) {
+      auto op_result =
+          PyObjectCast<pir::OpResult>(PyList_GetItem(op_result_list, i));
+      if (out_tensor_map.find(op_result) == out_tensor_map.end()) {
+        paddle::Tensor tensor = CreateTensorFromOpResult(op_result);
+        out_tensor_map[op_result] = tensor;
+        result.emplace_back(tensor);
+      } else {
+        result.emplace_back(out_tensor_map[op_result]);
+      }
+    }
+  } else if (PyTuple_Check(op_result_list)) {
+    Py_ssize_t len = PyTuple_Size(op_result_list);
+    for (Py_ssize_t i = 0; i < len; i++) {
+      auto op_result =
+          PyObjectCast<pir::OpResult>(PyTuple_GetItem(op_result_list, i));
+      if (out_tensor_map.find(op_result) == out_tensor_map.end()) {
+        paddle::Tensor tensor = CreateTensorFromOpResult(op_result);
+        out_tensor_map[op_result] = tensor;
+        result.emplace_back(tensor);
+      } else {
+        result.emplace_back(out_tensor_map[op_result]);
+      }
+    }
+  } else if (op_result_list != Py_None) {
+    PADDLE_THROW(platform::errors::InvalidArgument(
+        "Argument of GetTensorsWithOpResultInArgs must be list of OpResult, "
+        "but got "
+        "%s",
+        (reinterpret_cast<PyTypeObject*>(op_result_list->ob_type))->tp_name));
+  }
+
+  return ToPyObject(result);
+}
+
 paddle::experimental::Scalar CastNumpy2Scalar(PyObject* obj,
                                               const std::string& op_type,
                                               ssize_t arg_pos) {
@@ -2484,5 +2663,24 @@ void DistTensorConverter::operator()(paddle::optional<std::vector<Tensor>>* x) {
   }
 }
 
+static PyMethodDef EagerUtilMethods[] = {
+    {"create_empty_tensors_with_var_descs",
+     (PyCFunction)(void (*)(void))GetEmpytyTensorsWithVarDesc,
+     METH_VARARGS,
+     "GetEmpytyTensorsWithVarDesc"},
+    {"create_empty_tensors_with_op_results",
+     (PyCFunction)(void (*)(void))GetEmpytyTensorsWithOpResult,
+     METH_VARARGS,
+     "GetEmpytyTensorsWithOpResult."},
+    {nullptr, nullptr, 0, nullptr}};
+
+void BindEagerUtils(PyObject* module) {
+  if (PyModule_AddFunctions(module, EagerUtilMethods) < 0) {
+    PADDLE_THROW(platform::errors::Fatal(
+        "Init Paddle error in BindEagerUtils(PyModule_AddFunctions)."));
+    return;
+  }
+}
+
 }  // namespace pybind
 }  // namespace paddle
@@ -54,6 +54,18 @@ namespace pybind {
 
 namespace py = ::pybind11;
 
+template <typename T>
+static T PyObjectCast(PyObject* obj) {
+  try {
+    return py::cast<T>(py::handle(obj));
+  } catch (py::cast_error&) {
+    PADDLE_THROW(platform::errors::InvalidArgument(
+        "Python object is not type of %s, the real type is %s",
+        typeid(T).name(),
+        obj->ob_type->tp_name));
+  }
+}
+
 int TensorDtype2NumpyDtype(phi::DataType dtype);
 
 bool PyObject_CheckLongOrConvertToLong(PyObject** obj);
@@ -381,6 +393,10 @@ std::vector<paddle::Tensor> GetTensorListFromPyObject(PyObject* obj,
                                                       bool allow_none = false);
 paddle::Tensor& UnSafeGetTensorFromPyObject(PyObject* obj);
 
+PyObject* GetEmpytyTensorsWithVarDesc(PyObject* self, PyObject* args);
+
+PyObject* GetEmpytyTensorsWithOpResult(PyObject* self, PyObject* args);
+
 // end of Slice related methods
 
 std::vector<paddle::framework::Scope*> GetScopePtrListFromArgs(
@@ -468,5 +484,7 @@ void ConvertAllInputsToDistTensor(const phi::distributed::ProcessMesh* mesh,
 }
 
 void ConvertToDistTensor(Tensor* x, const phi::distributed::ProcessMesh* mesh);
+void BindEagerUtils(PyObject* module);
+
 }  // namespace pybind
 }  // namespace paddle
@@ -72,16 +72,6 @@ std::atomic<int> VarBaseUniqueNameID{0};
 
 namespace py = ::pybind11;
 
-template <typename T>
-static T PyObjectCast(PyObject *obj) {
-  try {
-    return py::cast<T>(py::handle(obj));
-  } catch (py::cast_error &) {
-    PADDLE_THROW(platform::errors::InvalidArgument(
-        "Python object is not type of %s", typeid(T).name()));
-  }
-}
-
 class PyVariableWrapperHook : public imperative::VariableWrapperHook {
  public:
   explicit PyVariableWrapperHook(PyObject *func) : py_func_(func) {
 
@@ -15,9 +15,16 @@
 #pragma once
 
 #include <pybind11/pybind11.h>
+#include "paddle/phi/common/data_type.h"
+#include "paddle/phi/core/ddim.h"
+#include "paddle/pir/core/op_result.h"
 
 namespace paddle {
 namespace pybind {
+using pir::OpResult;
 void BindPir(pybind11::module *m);
+phi::DataType GetOpResultDtype(const OpResult &result);
+const phi::DDim &GetOpResultDims(const OpResult &result);
+bool GetOpResultBoolAttr(const OpResult &self, const std::string &attr_name);
 }  // namespace pybind
 }  // namespace paddle
@@ -526,18 +526,6 @@ static PyObject *GetPythonAttribute(PyObject *obj, const char *attr_name) {
   }
 }
 
-template <typename T>
-static T PyObjectCast(PyObject *obj) {
-  try {
-    return py::cast<T>(py::handle(obj));
-  } catch (py::cast_error &) {
-    PADDLE_THROW(platform::errors::InvalidArgument(
-        "Python object is not type of %s, the real type is %s",
-        typeid(T).name(),
-        obj->ob_type->tp_name));
-  }
-}
-
 using PyNameVarBaseMap = std::unordered_map<std::string, py::handle>;
 
 static std::vector<std::shared_ptr<imperative::VarBase>> GetVarBaseList(
@@ -814,6 +802,7 @@ PYBIND11_MODULE(libpaddle, m) {
   BindJit(&m);
   BindEvalFrame(&m);
   BindCustomDevicePy(&m);
+  BindEagerUtils(m.ptr());
 
   // Not used, just make sure cpu_info.cc is linked.
   phi::backends::cpu::CpuTotalPhysicalMemory();