/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */
// disable numpy compile error
#include <Python.h>

#include <set>
#include <string>
#include <vector>

#pragma GCC diagnostic ignored "-Wattributes"
#include "paddle/fluid/eager/accumulation/accumulation_node.h"
#ifdef PADDLE_WITH_CUDA
#include "paddle/fluid/eager/activation_offloader.h"
#endif
#include "paddle/fluid/eager/api/all.h"
#include "paddle/fluid/eager/autograd_meta.h"
#include "paddle/fluid/eager/pylayer/py_layer_node.h"
#include "paddle/fluid/eager/utils.h"
#include "paddle/fluid/framework/convert_utils.h"
#include "paddle/fluid/platform/enforce.h"
#include "paddle/fluid/pybind/eager.h"
#include "paddle/fluid/pybind/eager_utils.h"
#include "paddle/fluid/pybind/exception.h"
#include "paddle/phi/common/data_type.h"
#include "paddle/phi/core/compat/convert_utils.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/core/memory/allocation/allocator.h"
#include "paddle/phi/core/memory/memcpy.h"
#include "pybind11/detail/internals.h"
#include "pybind11/pytypes.h"
#pragma GCC diagnostic ignored "-Wwrite-strings"
#pragma GCC diagnostic ignored "-Wmissing-field-initializers"
COMMON_DECLARE_bool(check_cuda_error);
COMMON_DECLARE_bool(check_nan_inf);
COMMON_DECLARE_int32(call_stack_level);
COMMON_DECLARE_int64(offload_retry_times);
COMMON_DECLARE_bool(enable_unique_name);
using egr::ConvertToDistTensor;
namespace paddle::pybind {

PyTypeObject* p_pylayer_type;
extern PyTypeObject* p_tensor_type;

std::set<Tensor*> GetTensorsFromPyObject(PyObject* obj) {
  std::set<Tensor*> result;
  if (obj == nullptr) {
    return result;
  }
  if (PyCheckTensor(obj)) {
    result.insert(&reinterpret_cast<TensorObject*>(obj)->tensor);  // NOLINT
  } else if (PyList_Check(obj)) {
    Py_ssize_t len = PyList_Size(obj);
    for (Py_ssize_t i = 0; i < len; i++) {
      if (PyCheckTensor(PyList_GetItem(obj, i))) {
        result.insert(
            &reinterpret_cast<TensorObject*>(PyList_GetItem(obj, i))  // NOLINT
                 ->tensor);
      }
    }
  } else if (PyTuple_Check(obj)) {
    Py_ssize_t len = PyTuple_Size(obj);
    for (Py_ssize_t i = 0; i < len; i++) {
      if (PyCheckTensor(PyTuple_GetItem(obj, i))) {
        result.insert(
            &reinterpret_cast<TensorObject*>(PyTuple_GetItem(obj, i))  // NOLINT
                 ->tensor);
      }
    }
  }
  return result;
}

PyObject* PyLayerNew(PyTypeObject* type, PyObject* args, PyObject* kwargs) {
  PyObject* obj = type->tp_alloc(type, 0);
  if (obj) {
    auto v = reinterpret_cast<PyLayerObject*>(obj);
    v->container = nullptr;
    v->materialize_grads = true;
    v->container_be_packed = false;
    v->grad_in_dtype_consistent = true;
    new (&v->grad_node) std::weak_ptr<egr::GradNodePyLayer>();
    new (&v->forward_input_tensor_is_duplicable) std::vector<bool>();
    new (&v->forward_output_tensor_is_duplicable) std::vector<bool>();
    new (&v->tensor_hold_helper)
        std::vector<std::shared_ptr<phi::DenseTensor>>();
    v->closure_obj = nullptr;
    new (&v->closure_tensor_hold_helper)
        std::vector<std::shared_ptr<phi::TensorBase>>();
#ifdef PADDLE_WITH_CUDA
    new (&v->reload_functors) std::vector<egr::ReloadFunctor>();
#endif
  }
  return obj;
}

static void PyLayerDealloc(PyLayerObject* self) {
  if (self->container) {
    Py_DECREF(self->container);
  }
  if (self->non_differentiable) {
    Py_DECREF(self->non_differentiable);
  }
  if (self->not_inplace_tensors) {
    Py_DECREF(self->not_inplace_tensors);
  }
  self->grad_node.~weak_ptr<egr::GradNodePyLayer>();
  self->unpack_hook = nullptr;
  self->forward_input_tensor_is_duplicable.~vector();
  self->forward_output_tensor_is_duplicable.~vector();
  self->tensor_hold_helper.~vector();
  Py_XDECREF(self->closure_obj);
  self->closure_obj = nullptr;
  self->closure_tensor_hold_helper.~vector();
#ifdef PADDLE_WITH_CUDA
  self->reload_functors.~vector();
#endif
  Py_TYPE(self)->tp_free(reinterpret_cast<PyObject*>(self));
}

PyObject* pylayer_method_name(PyObject* self, PyObject* noargs) {
  EAGER_TRY
  return ToPyObject(
      reinterpret_cast<PyLayerObject*>(self)->grad_node.lock()->name());
  EAGER_CATCH_AND_THROW_RETURN_NULL
}

PyObject* new_tensor_with_impl(Tensor* tensor) {
  PyObject* obj = p_tensor_type->tp_alloc(p_tensor_type, 0);
  if (obj) {
    auto v = reinterpret_cast<TensorObject*>(obj);
    new (&(v->tensor)) Tensor();
    v->tensor.set_impl(tensor->impl());
    v->tensor.set_name(egr::Controller::Instance().GenerateUniqueName());
    egr::EagerUtils::autograd_meta(&v->tensor)
        ->SetStopGradient(
            egr::EagerUtils::autograd_meta(tensor)->StopGradient());
  } else {
    PADDLE_THROW(
        common::errors::Fatal("tp_alloc return null, can not new a PyObject."));
  }
  return obj;
}

#ifdef PADDLE_WITH_CUDA
template <typename Callback>
static void GetTensorWithCallbackRecursively(PyObject* obj,
                                             const Callback& callback) {
  if (obj == nullptr || obj == Py_None) {
    return;
  } else if (paddle::pybind::PyCheckTensor(obj)) {
    const auto& tensor =
        reinterpret_cast<paddle::pybind::TensorObject*>(obj)->tensor;
    callback(tensor);
  } else if (PyTuple_Check(obj)) {
    Py_ssize_t n = PyTuple_GET_SIZE(obj);
    for (Py_ssize_t i = 0; i < n; ++i) {
      auto* item = PyTuple_GET_ITEM(obj, i);
      GetTensorWithCallbackRecursively(item, callback);
    }
  } else if (PyList_Check(obj)) {
    Py_ssize_t n = PyList_GET_SIZE(obj);
    for (Py_ssize_t i = 0; i < n; ++i) {
      auto* item = PyList_GET_ITEM(obj, i);
      GetTensorWithCallbackRecursively(item, callback);
    }
  }
}

static void PyLayerAddOffloadActivation(PyLayerObject* ctx,
                                        const std::string& name) {
  PADDLE_ENFORCE_NOT_NULL(
      ctx,
      common::errors::InvalidArgument("PyLayerObject should not be nullptr."));
  if (ctx->container_be_packed) {
    VLOG(10) << "Return directly because of packed value";
    return;
  }

  auto add_functor = [ctx, &name](const Tensor& t) {
    VLOG(10) << "Add offload tensor to PyLayer starts: " << name;
    auto reload_functor = egr::ActivationOffloader::Instance()->Add(t);
    if (const auto* rf_ptr = reload_functor.get_ptr()) {
      ctx->reload_functors.push_back(*rf_ptr);
    }
    VLOG(10) << "Add offload tensor to PyLayer ends: " << name;
  };

  GetTensorWithCallbackRecursively(ctx->container, add_functor);
}
#endif

PyObject* pylayer_method_apply(PyObject* cls,
                               PyObject* args,
                               PyObject* kwargs) {
  EAGER_TRY
  SetPythonStack();
  std::string classname =
      std::string(reinterpret_cast<PyTypeObject*>(cls)->tp_name);
  std::string forward_stack;
  VLOG(3) << classname << ":Running PyLayer Apply ";
  if (VLOG_IS_ON(2)) egr::LogIndent::Instance().IncreaseIndentLevel();
  if (FLAGS_check_nan_inf || FLAGS_call_stack_level == 3) {
    // record the forward stack
    forward_stack = egr::Controller::Instance().GetPythonStack();
  }
  std::string unique_api_name;
  if (VLOG_IS_ON(3) || FLAGS_enable_unique_name) {
    static int64_t call_count = 0;
    call_count++;
    unique_api_name = egr::GenerateUniqueApiName(classname, call_count);
  }
  VLOG(4) << classname << ":"
          << "Construct PyLayerContext";
  static PyObject* kBackwardFunctionAttr =
      PyUnicode_InternFromString("_backward_function");
  PyObject* backward_function = PyObject_GetAttr(cls, kBackwardFunctionAttr);
  if (!backward_function) [[unlikely]] {
    PADDLE_THROW(
        common::errors::InvalidArgument("Get _backward_function failed."));
  }
  PyLayerObject* ctx = reinterpret_cast<PyLayerObject*>(
      PyObject_CallFunctionObjArgs(backward_function, nullptr));
  if (!ctx) [[unlikely]] {
    PADDLE_THROW(
        common::errors::External(pybind11::detail::error_string().c_str()));
  }
  VLOG(6) << "PyLayer construct PyLayerContext finish...";
  if (FLAGS_check_cuda_error) [[unlikely]] {
    egr::CUDAErrorCheck("pylayer_method_apply " +
                        std::string(Py_TYPE(ctx)->tp_name) + " begin");
  }

  bool require_any_grad = false;

  size_t inputs_size = 0;
  size_t args_size = 0;
  size_t kwargs_size = 0;
  PyObject* forward_args = nullptr;
  PyObject* kwargs_value_list = nullptr;
  if (kwargs) {
    kwargs_size = PyDict_Size(kwargs);
    kwargs_value_list = PyDict_Values(kwargs);
  }
  if (args) {
    args_size = PyTuple_GET_SIZE(args);
  }
  inputs_size = kwargs_size + args_size;
  forward_args = PyTuple_New(args_size + 1);  // NOLINT
  Py_INCREF(ctx);
  PyTuple_SET_ITEM(forward_args, 0, reinterpret_cast<PyObject*>(ctx));
  VLOG(6) << classname << ":Prepare Pylayer forward args ";
  VLOG(6) << classname << ":Input size is " << inputs_size;
  paddle::small_vector<std::vector<egr::AutogradMeta*>> inputs_autograd_meta;
  inputs_autograd_meta.reserve(inputs_size);
  paddle::small_vector<std::vector<Tensor*>> inputs_tensor;
  inputs_tensor.reserve(inputs_size);
  ctx->forward_input_tensor_is_duplicable.clear();
  ctx->forward_input_tensor_is_duplicable.reserve(inputs_size);
  std::set<phi::TensorBase*> input_tensorbases;

  const phi::distributed::ProcessMesh* mesh = nullptr;
  auto TrySetMeshFromTensor = [&mesh](PyObject* obj) {
    if (!PyCheckTensor(obj)) return false;
    paddle::Tensor& tensor = reinterpret_cast<TensorObject*>(obj)->tensor;
    if (!tensor.defined() || !tensor.is_dist_tensor()) return false;
    mesh =
        &(std::static_pointer_cast<phi::distributed::DistTensor>(tensor.impl())
              ->dist_attr()
              .process_mesh());
    return true;
  };

  for (int64_t i = inputs_size - 1; i >= 0; --i) {
    PyObject* obj = nullptr;
    if (i >= static_cast<int64_t>(args_size)) {
      obj = PyList_GetItem(kwargs_value_list,
                           i - static_cast<int64_t>(args_size));  // NOLINT
    } else {
      obj = PyTuple_GET_ITEM(args, i);
    }
    bool mesh_found = false;
    if (PyList_Check(obj)) {
      int64_t len = PyList_Size(obj);
      for (int64_t j = len - 1; j >= 0; --j) {
        PyObject* o = PyList_GetItem(obj, j);
        mesh_found |= TrySetMeshFromTensor(o);
      }
    } else if (PyTuple_Check(obj)) {
      int64_t len = PyTuple_Size(obj);
      for (int64_t j = len - 1; j >= 0; --j) {
        PyObject* o = PyTuple_GetItem(obj, j);
        mesh_found |= TrySetMeshFromTensor(o);
      }
    } else {
      mesh_found |= TrySetMeshFromTensor(obj);
    }
    if (mesh_found) break;
  }

  auto HandleSingleTensorObj = [&mesh,
                                &input_tensorbases,
                                &inputs_autograd_meta,
                                &inputs_tensor,
                                &require_any_grad,
                                &ctx](TensorObject* t_obj) {
    auto& t = t_obj->tensor;
    if (mesh) {
      ConvertToDistTensor(&t, mesh);
    }
    input_tensorbases.insert(t.impl().get());
    auto* autograd_meta = egr::EagerUtils::nullable_autograd_meta(t);
    inputs_autograd_meta.push_back({autograd_meta});

    inputs_tensor.push_back({&t});  // NOLINT

    bool stop_gradient =
        autograd_meta == nullptr ? true : autograd_meta->StopGradient();
    if (!stop_gradient) {
      require_any_grad = true;
    }
    ctx->forward_input_tensor_is_duplicable.push_back(false);
  };

  auto CollectOneTensor = [&mesh, &input_tensorbases](
                              TensorObject* t_obj,
                              std::vector<paddle::Tensor*>* tensors) {
    auto& t = t_obj->tensor;

    if (mesh) {
      ConvertToDistTensor(&t, mesh);
    }

    input_tensorbases.insert(t.impl().get());
    tensors->push_back(&t);
  };

  auto FinalizeTensorList = [&require_any_grad,
                             &inputs_autograd_meta,
                             &inputs_tensor,
                             &ctx](std::vector<paddle::Tensor*>& tensors) {
    if (tensors.empty()) return;

    auto autograd_meta = egr::EagerUtils::nullable_autograd_meta(tensors);
    for (auto* m : autograd_meta) {
      bool stop_gradient = (m == nullptr) ? true : m->StopGradient();
      if (!stop_gradient) {
        require_any_grad = true;
      }
    }

    inputs_autograd_meta.push_back(autograd_meta);
    inputs_tensor.push_back(tensors);
    ctx->forward_input_tensor_is_duplicable.push_back(true);
  };

  for (size_t i = 0; i < inputs_size; ++i) {
    PyObject* obj = nullptr;
    if (i >= args_size) {
      obj = PyList_GetItem(kwargs_value_list, i - args_size);  // NOLINT
    } else {
      obj = PyTuple_GET_ITEM(args, i);
    }

    if (PyCheckTensor(obj)) {
      HandleSingleTensorObj(reinterpret_cast<TensorObject*>(obj));
    } else if (PyList_Check(obj)) {
      std::vector<Tensor*> tensors;
      Py_ssize_t len = PyList_Size(obj);
      for (Py_ssize_t j = 0; j < len; ++j) {
        PyObject* o = PyList_GetItem(obj, j);  // borrowed
        if (PyCheckTensor(o)) {
          CollectOneTensor(reinterpret_cast<TensorObject*>(o), &tensors);
        }
      }
      FinalizeTensorList(tensors);
    } else if (PyTuple_Check(obj)) {
      std::vector<Tensor*> tensors;
      Py_ssize_t len = PyTuple_Size(obj);
      for (Py_ssize_t j = 0; j < len; ++j) {
        PyObject* o = PyTuple_GetItem(obj, j);  // borrowed
        if (PyCheckTensor(o)) {
          CollectOneTensor(reinterpret_cast<TensorObject*>(o), &tensors);
        }
      }
      FinalizeTensorList(tensors);
    }

    if (i < args_size) {
      Py_INCREF(obj);
      PyTuple_SET_ITEM(forward_args, i + 1, obj);
    }
  }

  // Check LeafTensor if its GradNodeAccumulation TensorMeta is consistent with
  // its TensorMeta
  egr::CheckGradNodeAccumulation(inputs_tensor);

  VLOG(6)
      << classname << ":"
      << "PyLayer forward args is ready, begin call user's forward function...";
  // call forward
  static PyObject* kForwardAttr = PyUnicode_InternFromString("forward");
  auto forward_fn = PyObject_GetAttr(cls, kForwardAttr);
  if (!forward_fn) [[unlikely]] {
    PADDLE_THROW(
        common::errors::InvalidArgument("Get forward function failed."));
  }
  bool trace_backward = egr::Controller::Instance().HasGrad();
  egr::Controller::Instance().SetHasGrad(false);
  auto outputs = PyObject_Call(forward_fn, forward_args, kwargs);
  egr::Controller::Instance().SetHasGrad(trace_backward);
  if (!outputs) {
    Py_XDECREF(forward_args);
    Py_XDECREF(kwargs_value_list);
    Py_XDECREF(backward_function);
    Py_XDECREF(forward_fn);
    return nullptr;
  }

  PyObject* outputs_tuple = nullptr;
  if (PyTuple_Check(outputs)) {
    outputs_tuple = outputs;
  } else if (PyList_Check(outputs)) {
    outputs_tuple = PyList_AsTuple(outputs);
  } else {
    outputs_tuple = PyTuple_New(1);
    Py_INCREF(outputs);
    PyTuple_SET_ITEM(outputs_tuple, 0, outputs);
  }

  std::set<Tensor*> inplace_tensors;
  std::set<phi::TensorBase*> not_inplace_tensorbases;
  auto not_inplace_tensors = GetTensorsFromPyObject(ctx->not_inplace_tensors);
  for (auto it : not_inplace_tensors) {
    not_inplace_tensorbases.insert(it->impl().get());
  }

  auto outputs_size = PyTuple_GET_SIZE(outputs_tuple);
  paddle::small_vector<std::vector<Tensor*>> outputs_tensor;
  outputs_tensor.reserve(outputs_size);
  paddle::small_vector<std::vector<egr::AutogradMeta*>> outputs_autograd_meta;
  outputs_autograd_meta.reserve(outputs_size);
  ctx->forward_output_tensor_is_duplicable.clear();
  ctx->forward_output_tensor_is_duplicable.reserve(outputs_size);
  for (Py_ssize_t i = 0; i < outputs_size; i++) {
    PyObject* obj = PyTuple_GET_ITEM(outputs_tuple, i);
    if (PyCheckTensor(obj)) {
      outputs_tensor.push_back(
          {&(reinterpret_cast<TensorObject*>(obj)->tensor)});
      outputs_autograd_meta.push_back({egr::EagerUtils::autograd_meta(
          &(reinterpret_cast<TensorObject*>(obj)->tensor))});
      ctx->forward_output_tensor_is_duplicable.push_back(false);
      if (input_tensorbases.count(
              reinterpret_cast<TensorObject*>(obj)->tensor.impl().get())) {
        if (not_inplace_tensorbases.count(
                reinterpret_cast<TensorObject*>(obj)->tensor.impl().get())) {
          PyTuple_SET_ITEM(outputs_tuple,
                           i,
                           new_tensor_with_impl(&(
                               reinterpret_cast<TensorObject*>(obj)->tensor)));
        } else {
          inplace_tensors.insert(
              &(reinterpret_cast<TensorObject*>(obj)->tensor));
        }
      }
    } else if (PyList_Check(obj)) {
      std::vector<Tensor*> tensors;
      Py_ssize_t len = PyList_Size(obj);
      for (Py_ssize_t j = 0; j < len; j++) {
        PyObject* o = PyList_GetItem(obj, j);
        if (PyCheckTensor(o)) {
          tensors.push_back(&(reinterpret_cast<TensorObject*>(o)->tensor));
          if (input_tensorbases.count(
                  reinterpret_cast<TensorObject*>(o)->tensor.impl().get())) {
            if (not_inplace_tensorbases.count(
                    reinterpret_cast<TensorObject*>(o)->tensor.impl().get())) {
              PyList_SetItem(obj,
                             j,
                             new_tensor_with_impl(&(
                                 reinterpret_cast<TensorObject*>(o)->tensor)));
            } else {
              inplace_tensors.insert(
                  &(reinterpret_cast<TensorObject*>(o)->tensor));
            }
          }
        }
      }
      if (!tensors.empty()) {
        outputs_tensor.push_back(tensors);
        outputs_autograd_meta.push_back(
            egr::EagerUtils::autograd_meta(&tensors));
        ctx->forward_output_tensor_is_duplicable.push_back(true);
      }
    } else if (PyTuple_Check(obj)) {
      std::vector<Tensor*> tensors;
      Py_ssize_t len = PyTuple_Size(obj);
      for (Py_ssize_t j = 0; j < len; j++) {
        PyObject* o = PyTuple_GetItem(obj, j);
        if (PyCheckTensor(o)) {
          tensors.push_back(&(reinterpret_cast<TensorObject*>(o)->tensor));
          if (input_tensorbases.count(
                  reinterpret_cast<TensorObject*>(o)->tensor.impl().get())) {
            if (not_inplace_tensorbases.count(
                    reinterpret_cast<TensorObject*>(o)->tensor.impl().get())) {
              PyTuple_SetItem(obj,
                              j,
                              new_tensor_with_impl(&(
                                  reinterpret_cast<TensorObject*>(o)->tensor)));
            } else {
              inplace_tensors.insert(
                  &(reinterpret_cast<TensorObject*>(o)->tensor));
            }
          }
        }
      }
      if (!tensors.empty()) {
        outputs_tensor.push_back(tensors);
        outputs_autograd_meta.push_back(
            egr::EagerUtils::autograd_meta(&tensors));
        ctx->forward_output_tensor_is_duplicable.push_back(true);
      }
    }
  }
  if (outputs_tensor.empty()) [[unlikely]] {
    PADDLE_THROW(common::errors::InvalidArgument(
        "%s : At least one output of `PyLayer.forward` is a `Tensor`.",
        classname));
  }
  VLOG(6) << classname << ":"
          << "PyLayer forward function finish...";

#ifdef PADDLE_WITH_CUDA
  bool has_grad = false;
#endif
  if (require_any_grad && trace_backward) {
    auto non_differentiable = GetTensorsFromPyObject(ctx->non_differentiable);
    for (size_t i = 0; i < outputs_autograd_meta.size(); i++) {
      for (size_t j = 0; j < outputs_autograd_meta[i].size(); j++) {
        if (non_differentiable.find(outputs_tensor[i][j]) !=
            non_differentiable.end()) {
          outputs_autograd_meta[i][j]->SetStopGradient(true);
        } else {
          outputs_autograd_meta[i][j]->SetStopGradient(false);
        }
      }
    }

    for (auto inplace_tensor : inplace_tensors) {
      auto inplace_tensor_autograd_meta =
          egr::EagerUtils::autograd_meta(inplace_tensor);
      PADDLE_ENFORCE_EQ(!inplace_tensor_autograd_meta->StopGradient() &&
                            egr::EagerUtils::IsLeafTensor(*inplace_tensor),
                        false,
                        common::errors::InvalidArgument(
                            "%s : Leaf Var (%s) that doesn't stop gradient "
                            "can't use inplace strategy.",
                            classname,
                            inplace_tensor->name()));
      inplace_tensor->bump_inplace_version();
      VLOG(3) << "Tensor(" << inplace_tensor->name()
              << ") uses Inplace Strategy.";
    }

    auto grad_node =
        std::make_shared<egr::GradNodePyLayer>(reinterpret_cast<PyObject*>(ctx),
                                               outputs_autograd_meta.size(),
                                               inputs_autograd_meta.size());
    VLOG(3) << classname << ":"
            << "Create grad node " << grad_node->name() << " addr "
            << grad_node;
    // For dump call stack
    if (FLAGS_check_nan_inf || FLAGS_call_stack_level == 3) {
      grad_node->SetForwardTrace(forward_stack);
    }
    // Set for Record Subgraph
    if (egr::EagerBackwardSubGraphNodeRecorder::Instance()
            .NeedCaptureSubGraph()) {
      VLOG(3) << "Capture the grad node" << grad_node->name() << "("
              << grad_node.get() << ")"
              << "for subgraph.";
      egr::EagerBackwardSubGraphNodeRecorder::Instance().AddGradNode(
          grad_node.get());
    }

#ifdef PADDLE_WITH_CUDA
    has_grad = true;
#endif

    ctx->grad_node = grad_node;

    if (ctx->materialize_grads) {
      grad_node->SaveForwardOutputsMeta(outputs_tensor);
    }
    grad_node->SetGradInDtypeConsistent(ctx->grad_in_dtype_consistent);

    for (size_t i = 0; i < inputs_autograd_meta.size(); i++) {
      if (ctx->forward_input_tensor_is_duplicable[i]) {
        std::vector<const Tensor*> tmp;
        for (auto t : inputs_tensor[i]) {
          tmp.push_back(t);
        }
        grad_node->SetGradOutMeta(tmp, i);
      } else {
        grad_node->SetGradOutMeta(*inputs_tensor[i][0], i);
      }
    }

    for (size_t i = 0; i < outputs_autograd_meta.size(); i++) {
      if (ctx->forward_output_tensor_is_duplicable[i]) {
        egr::EagerUtils::SetOutRankWithSlot(&outputs_autograd_meta[i], i);
        egr::EagerUtils::SetHistory(&outputs_autograd_meta[i], grad_node);
        grad_node->SetGradInMeta(outputs_tensor[i], i);
      } else {
        egr::EagerUtils::SetOutRankWithSlot(outputs_autograd_meta[i][0], i);
        egr::EagerUtils::SetHistory(outputs_autograd_meta[i][0], grad_node);
        grad_node->SetGradInMeta(*outputs_tensor[i][0], i);
      }
    }
    VLOG(6) << classname << ":"
            << "PyLayer construct backward node finish...";
  }

  if (outputs_size == 1) {
    if (!PyTuple_Check(outputs) && !PyList_Check(outputs)) {
      Py_XDECREF(outputs);
      outputs = PyTuple_GetItem(outputs_tuple, 0);
      Py_INCREF(outputs);
      Py_XDECREF(outputs_tuple);
    }
  }
  VLOG(3) << classname << ":"
          << "PyLayer output size " << outputs_size;

  if (PyList_Check(outputs)) {
    Py_XDECREF(outputs_tuple);
  }

  Py_XDECREF(forward_args);
  Py_XDECREF(kwargs_value_list);
  Py_XDECREF(backward_function);
  Py_XDECREF(forward_fn);

#ifdef PADDLE_WITH_CUDA
  if (has_grad && FLAGS_offload_retry_times > 0) {
    auto grad_node = ctx->grad_node.lock();
    PADDLE_ENFORCE_NOT_NULL(
        grad_node,
        common::errors::InvalidArgument("%s : Cannot be null", classname));
    PyLayerAddOffloadActivation(ctx, grad_node->name());
  }
#endif
  Py_XDECREF(ctx);

  if (FLAGS_check_cuda_error) [[unlikely]] {
    egr::CUDAErrorCheck("pylayer_method_apply " +
                        std::string(Py_TYPE(ctx)->tp_name) + " finish");
  }
  VLOG(3) << classname << ":"
          << "Finish PyLayer Apply";
  if (VLOG_IS_ON(2)) egr::LogIndent::Instance().DecreaseIndentLevel();
  if (VLOG_IS_ON(3) || FLAGS_enable_unique_name) {
    const char* INPUT_PRINT_TEMPLATE =
        "\nForward Debug Info {\nAPI_Name: %s \nInput: [%s]  \nOutput: [%s] } ";
    std::string input_str = "";
    std::string output_str = "";
    int i = 0;
    for (auto& tensors : inputs_tensor) {
      const char* TENSOR_INPUT_TEMPLATE = " \n( input%d , %s), ";
      std::string input_x_str = paddle::string::Sprintf(
          TENSOR_INPUT_TEMPLATE, i, egr::EagerUtils::TensorStr(tensors));
      input_str += input_x_str;
      i++;
    }
    i = 0;
    for (auto& tensors : outputs_tensor) {
      egr::SetTensorName(unique_api_name, "out" + std::to_string(i), &tensors);
      const char* TENSOR_OUT_TEMPLATE = " \n( out%d , %s), ";
      std::string output_out_str = paddle::string::Sprintf(
          TENSOR_OUT_TEMPLATE, i, egr::EagerUtils::TensorStr(tensors));
      output_str += output_out_str;
      i++;
    }
    VLOG(3) << paddle::string::Sprintf(
        INPUT_PRINT_TEMPLATE, unique_api_name, input_str, output_str);
  }
  return outputs;
  EAGER_CATCH_AND_THROW_RETURN_NULL
}

// Deep-traverse a PyObject to collect shared_ptr<phi::DenseTensor> for all
// DenseTensors found (Tensor / Tuple / List / Dict, recursively).  Used by
// tensor_properties_set_container and ctx._hold_tensors to hold strong
// references so _clear_dataptr() cannot free the underlying allocation
// before backward.  DFS-walks obj and calls fn(tensor) for every Tensor
// leaf.  CollectDenseTensors and RestoreDenseTensors are built on top.
template <typename Fn>
static void WalkDenseTensors(PyObject* obj, Fn&& fn) {
  if (!obj || obj == Py_None) return;
  if (PyCheckTensor(obj)) {
    fn(reinterpret_cast<TensorObject*>(obj)->tensor);
    return;
  }
  if (PyTuple_Check(obj)) {
    Py_ssize_t n = PyTuple_GET_SIZE(obj);
    for (Py_ssize_t i = 0; i < n; ++i)
      WalkDenseTensors(PyTuple_GET_ITEM(obj, i), fn);
    return;
  }
  if (PyList_Check(obj)) {
    Py_ssize_t n = PyList_GET_SIZE(obj);
    for (Py_ssize_t i = 0; i < n; ++i)
      WalkDenseTensors(PyList_GET_ITEM(obj, i), fn);
    return;
  }
  if (PyDict_Check(obj)) {
    PyObject *k = nullptr, *v = nullptr;
    Py_ssize_t pos = 0;
    while (PyDict_Next(obj, &pos, &k, &v)) {
      WalkDenseTensors(v, fn);
    }
    return;
  }
}

static void CollectDenseTensors(
    PyObject* obj, std::vector<std::shared_ptr<phi::TensorBase>>* holder) {
  WalkDenseTensors(obj, [holder](const paddle::Tensor& tensor) {
    if (tensor.impl()) holder->push_back(tensor.impl());
  });
}

// Re-installs impl() for tensors cleared by _clear_dataptr(), using the
// shared_ptrs stored in holder (same DFS order as CollectDenseTensors).
static void RestoreDenseTensors(
    PyObject* obj,
    const std::vector<std::shared_ptr<phi::TensorBase>>& holder) {
  size_t idx = 0;
  WalkDenseTensors(obj, [&holder, &idx](paddle::Tensor& tensor) {
    if (idx < holder.size()) {
      if (!tensor.impl()) tensor.set_impl(holder[idx]);
      ++idx;
    }
  });
}

PyObject* call_unpack_hook(PyLayerObject* self) {
  auto unpack_hook = self->unpack_hook;
  auto packed_value = self->container;

  auto packed_value_size = PyTuple_GET_SIZE(packed_value);
  auto unpacked_value = PyTuple_New(packed_value_size);

  for (Py_ssize_t i = 0; i < packed_value_size; i++) {
    PyObject* obj = PyTuple_GET_ITEM(packed_value, i);
    if (PyList_Check(obj)) {
      Py_ssize_t len = PyList_Size(obj);
      auto tmp_list = PyList_New(len);
      for (Py_ssize_t j = 0; j < len; j++) {
        PyObject* o = PyList_GetItem(obj, j);
        PyTuple_SET_ITEM(tmp_list,
                         j,
                         reinterpret_cast<PyObject*>(((*unpack_hook)(
                             reinterpret_cast<void*>(o), nullptr))));
      }
      PyTuple_SET_ITEM(unpacked_value, i, tmp_list);
    } else if (PyTuple_Check(obj)) {
      Py_ssize_t len = PyTuple_Size(obj);
      auto tmp_tuple = PyTuple_New(len);
      for (Py_ssize_t j = 0; j < len; j++) {
        PyObject* o = PyTuple_GetItem(obj, j);
        PyTuple_SET_ITEM(tmp_tuple,
                         j,
                         reinterpret_cast<PyObject*>((*unpack_hook)(
                             reinterpret_cast<void*>(o), nullptr)));
      }
      PyTuple_SET_ITEM(unpacked_value, i, tmp_tuple);
    } else {
      PyTuple_SET_ITEM(unpacked_value,
                       i,
                       reinterpret_cast<PyObject*>((*unpack_hook)(
                           reinterpret_cast<void*>(obj), nullptr)));
    }
  }

  return unpacked_value;
}

PyObject* tensor_properties_get_container(PyLayerObject* self, void* closure) {
  EAGER_TRY
  if (self->container == nullptr) {
    RETURN_PY_NONE;
  }
  if (self->container_be_packed) {
    return call_unpack_hook(self);
  }
  // Re-attach any DenseTensor impls that were freed by _clear_dataptr().
  // tensor_hold_helper keeps the underlying allocations alive; walk the
  // container in the same DFS order as CollectDenseTensors and reinstall
  // impls for tensors whose impl() is currently null.
  if (!self->tensor_hold_helper.empty()) {
    RestoreDenseTensors(self->container, self->tensor_hold_helper);
  }
  Py_INCREF(self->container);
  return self->container;
  EAGER_CATCH_AND_THROW_RETURN_NULL
}

void call_pack_hook(PyLayerObject* self, PyObject* value) {
  PyObject* saved_value = nullptr;
  if (PyTuple_Check(value)) {
    saved_value = value;
  } else if (PyList_Check(value)) {
    saved_value = PyList_AsTuple(value);
  } else {
    saved_value = PyTuple_New(1);
    Py_INCREF(value);
    PyTuple_SET_ITEM(saved_value, 0, value);
  }

  auto pack_hook = egr::SavedTensorsHooks::GetInstance().GetPackHook();
  self->unpack_hook = egr::SavedTensorsHooks::GetInstance().GetUnPackHook();

  auto saved_value_size = PyTuple_GET_SIZE(saved_value);
  PyObject* packed_value = PyTuple_New(saved_value_size);

  for (Py_ssize_t i = 0; i < saved_value_size; i++) {
    PyObject* obj = PyTuple_GET_ITEM(saved_value, i);
    if (PyCheckTensor(obj)) {
      PyTuple_SET_ITEM(packed_value,
                       i,
                       reinterpret_cast<PyObject*>(
                           (*pack_hook)(reinterpret_cast<void*>(obj))));
    } else if (PyList_Check(obj)) {
      Py_ssize_t len = PyList_Size(obj);
      auto tmp_list = PyList_New(len);
      for (Py_ssize_t j = 0; j < len; j++) {
        PyObject* o = PyList_GetItem(obj, j);
        if (PyCheckTensor(o)) {
          PyTuple_SET_ITEM(tmp_list,
                           j,
                           reinterpret_cast<PyObject*>(
                               (*pack_hook)(reinterpret_cast<void*>(o))));
        } else if (o == Py_None) {
          PyTuple_SET_ITEM(tmp_list,
                           j,
                           reinterpret_cast<PyObject*>(
                               (*pack_hook)(reinterpret_cast<void*>(o))));
        } else {
          PADDLE_THROW(common::errors::InvalidArgument(
              "save_for_backward only support Tensor, list of Tensor, tuple of "
              "Tensor."));
        }
      }
      PyTuple_SET_ITEM(packed_value, i, tmp_list);
    } else if (PyTuple_Check(obj)) {
      Py_ssize_t len = PyTuple_Size(obj);
      auto tmp_tuple = PyTuple_New(len);
      for (Py_ssize_t j = 0; j < len; j++) {
        PyObject* o = PyTuple_GetItem(obj, j);
        if (PyCheckTensor(o)) {
          PyTuple_SET_ITEM(tmp_tuple,
                           j,
                           reinterpret_cast<PyObject*>(
                               (*pack_hook)(reinterpret_cast<void*>(o))));
        } else if (o == Py_None) {
          PyTuple_SET_ITEM(tmp_tuple,
                           j,
                           reinterpret_cast<PyObject*>(
                               (*pack_hook)(reinterpret_cast<void*>(o))));
        } else {
          PADDLE_THROW(common::errors::InvalidArgument(
              "save_for_backward only support Tensor, list of Tensor, tuple of "
              "Tensor."));
        }
      }
      PyTuple_SET_ITEM(packed_value, i, tmp_tuple);
    } else if (obj == Py_None) {
      PyTuple_SET_ITEM(packed_value,
                       i,
                       reinterpret_cast<PyObject*>(
                           (*pack_hook)(reinterpret_cast<void*>(obj))));
    } else {
      PADDLE_THROW(common::errors::InvalidArgument(
          "save_for_backward only support Tensor, list of Tensor, tuple of "
          "Tensor."));
    }
  }

  if (PyTuple_Check(value)) {
    Py_XDECREF(saved_value);
  }

  Py_XDECREF(self->container);
  self->container = packed_value;
  self->container_be_packed = true;
}

int tensor_properties_set_container(PyLayerObject* self,
                                    PyObject* value,
                                    void* closure) {
  EAGER_TRY
  if (egr::SavedTensorsHooks::GetInstance().IsEnable()) {
    // Note 1: when hooks are enabled the tensors are packed; do NOT populate
    // tensor_hold_helper (the hook system manages tensor lifetimes itself).
    call_pack_hook(self, value);
  } else {
    Py_XINCREF(value);
    Py_XDECREF(self->container);
    self->container = value;
    // Note 2: deep-traverse value (Tensor / Tuple / List / nested) to hold
    // strong references to every DenseTensor impl, preventing _clear_dataptr()
    // from freeing the underlying allocation before backward runs.
    self->tensor_hold_helper.clear();
    CollectDenseTensors(value, &self->tensor_hold_helper);
  }
  return 0;
  EAGER_CATCH_AND_THROW_RETURN_NEG
}

PyObject* tensor_properties_get_non_differentiable(PyLayerObject* self,
                                                   void* closure) {
  EAGER_TRY
  if (self->non_differentiable == nullptr) {
    RETURN_PY_NONE;
  }
  Py_INCREF(self->non_differentiable);
  return self->non_differentiable;
  EAGER_CATCH_AND_THROW_RETURN_NULL
}

int tensor_properties_set_non_differentiable(PyLayerObject* self,
                                             PyObject* value,
                                             void* closure) {
  EAGER_TRY
  Py_XINCREF(value);
  Py_XDECREF(self->non_differentiable);
  self->non_differentiable = value;
  return 0;
  EAGER_CATCH_AND_THROW_RETURN_NEG
}

PyObject* tensor_properties_get_not_inplace_tensors(PyLayerObject* self,
                                                    void* closure) {
  EAGER_TRY
  if (self->not_inplace_tensors == nullptr) {
    RETURN_PY_NONE;
  }
  Py_INCREF(self->not_inplace_tensors);
  return self->not_inplace_tensors;
  EAGER_CATCH_AND_THROW_RETURN_NULL
}

int tensor_properties_set_not_inplace_tensors(PyLayerObject* self,
                                              PyObject* value,
                                              void* closure) {
  EAGER_TRY
  Py_XINCREF(value);
  Py_XDECREF(self->not_inplace_tensors);
  self->not_inplace_tensors = value;
  return 0;
  EAGER_CATCH_AND_THROW_RETURN_NEG
}

int tensor_properties_set_materialize_grads(PyLayerObject* self,
                                            PyObject* value,
                                            void* closure) {
  EAGER_TRY
  self->materialize_grads = CastPyArg2AttrBoolean(value, 0);
  return 0;
  EAGER_CATCH_AND_THROW_RETURN_NEG
}
int tensor_properties_set_grad_in_dtype_consistent(PyLayerObject* self,
                                                   PyObject* value,
                                                   void* closure) {
  EAGER_TRY
  self->grad_in_dtype_consistent = CastPyArg2AttrBoolean(value, 0);
  return 0;
  EAGER_CATCH_AND_THROW_RETURN_NEG
}

// ctx._pop_saved_impl(tensor)
// Removes the strong reference held in tensor_hold_helper for the given
// tensor's underlying DenseTensor, allowing its memory to be freed early
// (e.g. inside backward when the tensor is no longer needed).
// The tensor must have a valid impl() — i.e. pass the recovered tensor
// returned by ctx.saved_tensor(), not the already-cleared one.
PyObject* pylayer_pop_saved_impl(PyObject* self_, PyObject* args) {
  EAGER_TRY
  auto* self = reinterpret_cast<PyLayerObject*>(self_);
  PyObject* tensor_obj = nullptr;
  if (!PyArg_ParseTuple(args, "O", &tensor_obj)) {
    RETURN_PY_NONE;
  }
  if (!tensor_obj || !PyCheckTensor(tensor_obj)) {
    RETURN_PY_NONE;
  }
  const auto& tensor = reinterpret_cast<TensorObject*>(tensor_obj)->tensor;
  if (!tensor.impl() || !tensor.is_dense_tensor()) {
    RETURN_PY_NONE;
  }
  auto* raw = static_cast<phi::DenseTensor*>(tensor.impl().get());
  for (auto it = self->tensor_hold_helper.begin();
       it != self->tensor_hold_helper.end();
       ++it) {
    if (it->get() == raw) {
      self->tensor_hold_helper.erase(it);
      break;
    }
  }
  RETURN_PY_NONE;
  EAGER_CATCH_AND_THROW_RETURN_NULL
}

// ctx._hold_tensors(obj)
// Keep strong refs to the owning container (Py_INCREF'd) and the impl() of
// every DenseTensor leaf found in obj (Tensor / Tuple / List / Dict).
// Covers tensors captured via Python closure of the forward function that
// bypass save_for_backward / container.  Skipped when saved_tensors_hooks is
// enabled (the hook system owns tensor lifetime in that case).
PyObject* pylayer_hold_tensors(PyObject* self_, PyObject* args) {
  EAGER_TRY
  auto* self = reinterpret_cast<PyLayerObject*>(self_);
  PyObject* obj = nullptr;
  if (!PyArg_ParseTuple(args, "O", &obj)) {
    RETURN_PY_NONE;
  }
  if (obj && obj != Py_None &&
      !egr::SavedTensorsHooks::GetInstance().IsEnable()) {
    Py_INCREF(obj);
    Py_XDECREF(self->closure_obj);
    self->closure_obj = obj;
    self->closure_tensor_hold_helper.clear();
    CollectDenseTensors(obj, &self->closure_tensor_hold_helper);
  }
  RETURN_PY_NONE;
  EAGER_CATCH_AND_THROW_RETURN_NULL
}

// ctx._restore_held_tensors()
// Re-install impl() on any Python Tensor previously registered via
// _hold_tensors whose impl_ has been nulled by _clear_dataptr().  Typically
// called at the start of backward before recompute re-runs forward.
PyObject* pylayer_restore_held_tensors(PyObject* self_, PyObject* /*unused*/) {
  EAGER_TRY
  auto* self = reinterpret_cast<PyLayerObject*>(self_);
  if (self->closure_obj && !self->closure_tensor_hold_helper.empty()) {
    RestoreDenseTensors(self->closure_obj, self->closure_tensor_hold_helper);
  }
  RETURN_PY_NONE;
  EAGER_CATCH_AND_THROW_RETURN_NULL
}

PyMethodDef pylayer_methods[] = {
    {"name",  // NOLINT
     (PyCFunction)(void (*)())pylayer_method_name,
     METH_NOARGS,
     nullptr},
    {"apply",
     (PyCFunction)(void (*)())pylayer_method_apply,
     METH_CLASS | METH_VARARGS | METH_KEYWORDS,
     nullptr},
    {"_pop_saved_impl",
     (PyCFunction)(void (*)())pylayer_pop_saved_impl,
     METH_VARARGS,
     "Release the strong reference held for a "
     "specific DenseTensor saved via "
     "save_for_backward, allowing its memory to "
     "be freed early if no other holder exists."},
    {"_hold_tensors",
     (PyCFunction)(void (*)())pylayer_hold_tensors,
     METH_VARARGS,
     "Deep-traverse the given object (Tensor / tuple / list / dict) and "
     "keep strong references to every DenseTensor impl found, plus the "
     "owning Python Tensor object.  Used to protect tensors captured in "
     "Python closures against _clear_dataptr() in pipeline parallel."},
    {"_restore_held_tensors",
     (PyCFunction)(void (*)())pylayer_restore_held_tensors,
     METH_NOARGS,
     "Reinstall impl() on Python Tensor objects previously registered via "
     "_hold_tensors, if their impl() has been nulled by _clear_dataptr()."},
    {nullptr, nullptr, 0, nullptr}};

struct PyGetSetDef pylayer_properties[] {  // NOLINT
  {"container",
   (getter)tensor_properties_get_container,
   (setter)tensor_properties_set_container,
   nullptr,
   nullptr},
      {"non_differentiable",
       (getter)tensor_properties_get_non_differentiable,
       (setter)tensor_properties_set_non_differentiable,
       nullptr,
       nullptr},
      {"not_inplace_tensors",
       (getter)tensor_properties_get_not_inplace_tensors,
       (setter)tensor_properties_set_not_inplace_tensors,
       nullptr,
       nullptr},
      {"materialize_grads",
       nullptr,
       (setter)tensor_properties_set_materialize_grads,
       nullptr,
       nullptr},
      {"grad_in_dtype_consistent",
       nullptr,
       (setter)tensor_properties_set_grad_in_dtype_consistent,
       nullptr,
       nullptr},
  {
    nullptr, nullptr, nullptr, nullptr, nullptr
  }
};

void BindEagerPyLayer(PyObject* module) {
  auto heap_type = reinterpret_cast<PyHeapTypeObject*>(
      PyType_Type.tp_alloc(&PyType_Type, 0));
  heap_type->ht_name = ToPyObject("PyLayer");
  heap_type->ht_qualname = ToPyObject("PyLayer");
  auto type = &heap_type->ht_type;
  type->tp_name = "PyLayer";
  type->tp_basicsize = sizeof(PyLayerObject);
  type->tp_dealloc = (destructor)PyLayerDealloc;
  type->tp_methods = pylayer_methods;
  type->tp_getset = pylayer_properties;
  type->tp_new = (newfunc)PyLayerNew;
  Py_INCREF(&PyBaseObject_Type);
  type->tp_base = reinterpret_cast<PyTypeObject*>(&PyBaseObject_Type);
  type->tp_flags |=
      Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HEAPTYPE;  // NOLINT
  type->tp_as_async = &heap_type->as_async;
  p_pylayer_type = type;

  if (PyType_Ready(type) < 0) {
    PADDLE_THROW(
        common::errors::Fatal("Init Paddle error in BindEager(PyType_Ready)."));
    return;
  }

  Py_INCREF(type);
  if (PyModule_AddObject(module, "PyLayer", reinterpret_cast<PyObject*>(type)) <
      0) {
    Py_DECREF(type);
    Py_DECREF(module);
    PADDLE_THROW(common::errors::Fatal(
        "Init Paddle error in BindEager(PyModule_AddObject)."));
    return;
  }
}

}  // namespace paddle::pybind