// Copyright (c) 2021 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.

#include "paddle/fluid/eager/grad_node_info.h"

#include "glog/logging.h"
#include "paddle/common/errors.h"
#include "paddle/fluid/eager/accumulation/accumulation_node.h"
#include "paddle/fluid/eager/autograd_meta.h"
#include "paddle/fluid/eager/utils.h"
#include "paddle/fluid/framework/convert_utils.h"
#include "paddle/fluid/framework/data_type.h"
#include "paddle/fluid/framework/data_type_transform.h"
#include "paddle/fluid/framework/var_type.h"
#include "paddle/fluid/platform/enforce.h"
#include "paddle/phi/common/data_type.h"
#include "paddle/phi/core/dense_tensor.h"

#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
#include "paddle/phi/core/sparse_coo_tensor.h"
#include "paddle/phi/core/sparse_csr_tensor.h"

COMMON_DECLARE_bool(enable_unique_name);

/**
 * Implementation of GradNodeBase and Edge.
 **/
namespace egr {

static void CheckTensor(const paddle::Tensor& pre, const paddle::Tensor& post) {
  if (!pre.has_allocation() && post.has_allocation()) {
    PADDLE_THROW(common::errors::PermissionDenied(
        "The tensor in before and after hook are not consistent"));
  }
  if (pre.has_allocation() && post.has_allocation()) {
    VLOG(7) << phi::DataTypeToString(pre.dtype()) << " "
            << phi::DataTypeToString(post.dtype());
    PADDLE_ENFORCE_EQ(
        pre.dtype(),
        post.dtype(),
        common::errors::PermissionDenied(
            "The dtype of tensor before(%s) and after(%s) hook are not "
            "consistent",
            phi::DataTypeToString(pre.dtype()),
            phi::DataTypeToString(post.dtype())));
    PADDLE_ENFORCE_EQ(pre.place(),
                      post.place(),
                      common::errors::PermissionDenied(
                          "The place of tensor before(%s) and after(%s) "
                          "hook are not consistent",
                          pre.place().DebugString(),
                          post.place().DebugString()));
  }
}

GradNodeBase::GradNodeBase(size_t bwd_in_slot_num, size_t bwd_out_slot_num)
    : bwd_out_meta_(), bwd_in_meta_(), gradient_hooks_() {
  VLOG(7) << "Construct GradNodeBase";
  bwd_in_meta_.resize(bwd_in_slot_num);
  bwd_out_meta_.resize(bwd_out_slot_num);
}

const paddle::small_vector<std::vector<GradSlotMeta>, kSlotSmallVectorSize>&
GradNodeBase::InputMeta() const {
  return bwd_in_meta_;
}

const paddle::small_vector<std::vector<GradSlotMeta>, kSlotSmallVectorSize>&
GradNodeBase::OutputMeta() const {
  return bwd_out_meta_;
}

paddle::small_vector<std::vector<GradSlotMeta>, kSlotSmallVectorSize>&
GradNodeBase::MutableOutputMeta() {
  return bwd_out_meta_;
}

paddle::small_vector<std::vector<GradSlotMeta>, kSlotSmallVectorSize>&
GradNodeBase::MutableInputMeta() {
  return bwd_in_meta_;
}
void GradNodeBase::SetGradInMeta(const paddle::Tensor& fwd_out,
                                 size_t slot_rank) {
  VLOG(7) << "Set GradSlotMeta for Grad Inputs";
  auto* fwd_out_meta = egr::EagerUtils::nullable_autograd_meta(fwd_out);
  PADDLE_ENFORCE_LE(
      slot_rank,
      (bwd_in_meta_.size() - 1),
      common::errors::InvalidArgument(
          "Slot Rank should less equal than bwd_in_meta_ size, since "
          "bwd_in_meta_ is designed to hold as same num as backward "
          "inputs."));
  auto& metas = bwd_in_meta_.at(slot_rank);
  if (metas.empty()) {
    metas.resize(1);
  }

  auto& meta = metas[0];
  if (fwd_out_meta && fwd_out_meta->StopGradient()) {
    meta.SetStopGradient(fwd_out_meta->StopGradient());
  }

  if (!fwd_out.has_allocation()) {
    if (fwd_out.defined() && fwd_out.is_dist_tensor() &&
        phi::distributed::NeedComputationClipForPP(fwd_out.impl())) {
      VLOG(5) << "Tensor " << fwd_out.name() << " is DistTensor,"
              << " and needs computation clip for pipeline parallel."
              << " Still SetGradInMeta for it.";
    } else {
      VLOG(6)
          << "Skip Configuring GradSlotMeta for uninitialized GradInput Tensor";
      return;
    }
  }

  const phi::DenseTensor* dense_tensor = nullptr;
  // Record TensorMeta
  if (phi::DenseTensor::classof(fwd_out.impl().get())) {
    // Only Copy Meta
    dense_tensor = static_cast<phi::DenseTensor*>(fwd_out.impl().get());
  } else if (phi::SparseCooTensor::classof(fwd_out.impl().get())) {
    phi::SparseCooTensor* coo_tensor =
        static_cast<phi::SparseCooTensor*>(fwd_out.impl().get());
    dense_tensor = coo_tensor->mutable_non_zero_elements();
  } else if (phi::SparseCsrTensor::classof(fwd_out.impl().get())) {
    phi::SparseCsrTensor* csr_tensor =
        static_cast<phi::SparseCsrTensor*>(fwd_out.impl().get());
    dense_tensor = csr_tensor->mutable_non_zero_elements();
  } else if (phi::distributed::DistTensor::classof(fwd_out.impl().get())) {
    dense_tensor =  // NOLINT
        &(static_cast<phi::distributed::DistTensor*>(fwd_out.impl().get())
              ->value());
    meta.SetDistAttr(
        static_cast<phi::distributed::DistTensor*>(fwd_out.impl().get())
            ->dist_attr());
    meta.SetDistTensorGlobalDims(
        static_cast<phi::distributed::DistTensor*>(fwd_out.impl().get())
            ->dims());
    SetIsRunAutoParallel(true);
  } else {
    VLOG(5) << "Unable to initialize the DenseTensorMeta of GradSlotMeta with "
               "non-DenseTensor argument.";
    return;
  }
  PADDLE_ENFORCE_NE(
      dense_tensor->meta().dtype,
      phi::DataType::UNDEFINED,
      common::errors::Fatal(
          "Attempting to copy DenseTensorMeta with phi::DataType::UNDEFINED,"
          "which is illegal."));

  meta.SetTensorMeta(dense_tensor->meta());
  meta.SetPlace(fwd_out.place());

  if (dense_tensor->type() == phi::DataType::COMPLEX64 ||
      dense_tensor->type() == phi::DataType::COMPLEX128) {
    need_complex_to_real_ = true;
  }
}

void GradNodeBase::SetGradInMeta(const std::vector<paddle::Tensor>& fwd_out,
                                 size_t slot_rank) {
  VLOG(7) << "Set GradSlotMeta for Grad Inputs";
  size_t slot_size = fwd_out.size();
  PADDLE_ENFORCE_LE(
      slot_rank,
      (bwd_in_meta_.size() - 1),
      common::errors::InvalidArgument(
          "Slot Rank should less equal than bwd_in_meta_ size, since "
          "bwd_in_meta_ is designed to hold as same num as backward "
          "inputs."));
  auto& metas = bwd_in_meta_.at(slot_rank);
  // Init stop gradient vector before use to avoid push back
  if (metas.size() < slot_size) {
    VLOG(7) << "Init bwd_in_meta_ with slot rank: " << slot_rank;
    metas.resize(slot_size);
  }
  for (size_t i = 0; i < slot_size; i++) {
    auto& meta = metas[i];
    const auto& fwd_out_tensor = fwd_out[i];
    auto* fwd_out_meta =
        egr::EagerUtils::nullable_autograd_meta(fwd_out_tensor);
    PADDLE_ENFORCE_NOT_NULL(fwd_out_meta,
                            common::errors::PreconditionNotMet(
                                "Bwd_in_meta should only be called while "
                                "autograd_meta is not null. If you got this "
                                "error, it indicates bugs in framework."));
    if (fwd_out_meta && fwd_out_meta->StopGradient()) {
      // Set Stop Gradient only when its true or non-initialized autograd_meta,
      // since all default value is false.
      meta.SetStopGradient(fwd_out_meta->StopGradient());
    }

    if (!fwd_out_tensor.has_allocation()) {
      if (fwd_out_tensor.defined() && fwd_out_tensor.is_dist_tensor() &&
          phi::distributed::NeedComputationClipForPP(fwd_out_tensor.impl())) {
        VLOG(3) << "Tensor " << fwd_out_tensor.name() << " is DistTensor,"
                << " and needs computation clip for pipeline parallel."
                << " Still SetGradInMeta for it.";
      } else {
        VLOG(6) << "Skip Configuring GradSlotMeta for uninitialized GradInput "
                   "Tensor";
        return;
      }
    }

    // Record TensorMeta
    if (phi::DenseTensor::classof(fwd_out_tensor.impl().get())) {
      // Only Copy Meta
      phi::DenseTensor* dense_tensor =
          static_cast<phi::DenseTensor*>(fwd_out_tensor.impl().get());

      PADDLE_ENFORCE_NE(
          dense_tensor->meta().dtype,
          phi::DataType::UNDEFINED,
          common::errors::Fatal("Attempting to copy DenseTensorMeta "
                                "with phi::DataType::UNDEFINED,"
                                "which is illegal."));
      meta.SetTensorMeta(dense_tensor->meta());
      meta.SetPlace(fwd_out_tensor.place());

      if (dense_tensor->type() == phi::DataType::COMPLEX64 ||
          dense_tensor->type() == phi::DataType::COMPLEX128) {
        need_complex_to_real_ = true;
      }
    } else if (phi::distributed::DistTensor::classof(
                   fwd_out_tensor.impl().get())) {
      // Only Copy Meta
      meta.SetDistAttr(static_cast<phi::distributed::DistTensor*>(
                           fwd_out_tensor.impl().get())
                           ->dist_attr());
      meta.SetDistTensorGlobalDims(static_cast<phi::distributed::DistTensor*>(
                                       fwd_out_tensor.impl().get())
                                       ->dims());
      SetIsRunAutoParallel(true);

      auto dense_tensor = static_cast<phi::distributed::DistTensor*>(
                              fwd_out_tensor.impl().get())
                              ->value();

      PADDLE_ENFORCE_NE(
          dense_tensor.meta().dtype,
          phi::DataType::UNDEFINED,
          common::errors::Fatal("Attempting to copy DenseTensorMeta "
                                "with phi::DataType::UNDEFINED,"
                                "which is illegal."));
      meta.SetTensorMeta(dense_tensor.meta());
      meta.SetPlace(fwd_out_tensor.place());

      if (dense_tensor.type() == phi::DataType::COMPLEX64 ||
          dense_tensor.type() == phi::DataType::COMPLEX128) {
        need_complex_to_real_ = true;
      }
    } else {
      VLOG(7) << "Unable to initialize the DenseTensorMeta of GradSlotMeta "
                 "with non-DenseTensor argument.";
    }
  }
}

void GradNodeBase::SetGradInMeta(const std::vector<paddle::Tensor*>& fwd_out,
                                 size_t slot_rank) {
  VLOG(7) << "Set GradSlotMeta for Grad Inputs";
  size_t slot_size = fwd_out.size();
  PADDLE_ENFORCE_LE(
      slot_rank,
      (bwd_in_meta_.size() - 1),
      common::errors::InvalidArgument(
          "Slot Rank should less equal than bwd_in_meta_ size, since "
          "bwd_in_meta_ is designed to hold as same num as backward "
          "inputs."));
  auto& metas = bwd_in_meta_.at(slot_rank);
  // Init stop gradient vector before use to avoid push back
  if (metas.size() < slot_size) {
    VLOG(7) << "Init bwd_in_meta_ with slot rank: " << slot_rank;
    metas.resize(slot_size);
  }
  for (size_t i = 0; i < slot_size; i++) {
    auto& meta = metas[i];
    const auto& fwd_out_tensor = *fwd_out[i];
    auto* fwd_out_meta =
        egr::EagerUtils::nullable_autograd_meta(fwd_out_tensor);
    PADDLE_ENFORCE_NOT_NULL(fwd_out_meta,
                            common::errors::PreconditionNotMet(
                                "Bwd_in_meta should only be called while "
                                "autograd_meta is not null. If you got this "
                                "error, it indicates bugs in framework."));
    if (fwd_out_meta && fwd_out_meta->StopGradient()) {
      // Set Stop Gradient only when its true or non-initialized autograd_meta,
      // since all default value is false.
      meta.SetStopGradient(fwd_out_meta->StopGradient());
    }

    if (!fwd_out_tensor.has_allocation()) {
      if (fwd_out_tensor.defined() && fwd_out_tensor.is_dist_tensor() &&
          phi::distributed::NeedComputationClipForPP(fwd_out_tensor.impl())) {
        VLOG(3) << "Tensor " << fwd_out_tensor.name() << " is DistTensor,"
                << " and needs computation clip for pipeline parallel."
                << " Still SetGradInMeta for it.";
      } else {
        VLOG(7) << "Skip Configuring GradSlotMeta for uninitialized GradInput "
                   "Tensor";
        return;
      }
    }

    // Record TensorMeta
    if (phi::DenseTensor::classof(fwd_out_tensor.impl().get())) {
      // Only Copy Meta
      phi::DenseTensor* dense_tensor =
          static_cast<phi::DenseTensor*>(fwd_out_tensor.impl().get());

      PADDLE_ENFORCE_NE(
          dense_tensor->meta().dtype,
          phi::DataType::UNDEFINED,
          common::errors::Fatal("Attempting to copy DenseTensorMeta "
                                "with phi::DataType::UNDEFINED,"
                                "which is illegal."));
      meta.SetTensorMeta(dense_tensor->meta());
      meta.SetPlace(fwd_out_tensor.place());

      if (dense_tensor->type() == phi::DataType::COMPLEX64 ||
          dense_tensor->type() == phi::DataType::COMPLEX128) {
        need_complex_to_real_ = true;
      }
    } else if (phi::distributed::DistTensor::classof(
                   fwd_out_tensor.impl().get())) {
      // Only Copy Meta
      meta.SetDistAttr(static_cast<phi::distributed::DistTensor*>(
                           fwd_out_tensor.impl().get())
                           ->dist_attr());
      meta.SetDistTensorGlobalDims(static_cast<phi::distributed::DistTensor*>(
                                       fwd_out_tensor.impl().get())
                                       ->dims());
      SetIsRunAutoParallel(true);

      auto dense_tensor = static_cast<phi::distributed::DistTensor*>(
                              fwd_out_tensor.impl().get())
                              ->value();

      PADDLE_ENFORCE_NE(
          dense_tensor.meta().dtype,
          phi::DataType::UNDEFINED,
          common::errors::Fatal("Attempting to copy DenseTensorMeta "
                                "with phi::DataType::UNDEFINED,"
                                "which is illegal."));
      meta.SetTensorMeta(dense_tensor.meta());
      meta.SetPlace(fwd_out_tensor.place());

      if (dense_tensor.type() == phi::DataType::COMPLEX64 ||
          dense_tensor.type() == phi::DataType::COMPLEX128) {
        need_complex_to_real_ = true;
      }
    } else {
      VLOG(5) << "Unable to initialize the DenseTensorMeta of GradSlotMeta "
                 "with non-DenseTensor argument.";
    }
  }
}

void GradNodeBase::SetGradOutMeta(const paddle::Tensor& fwd_in,
                                  size_t slot_rank) {
  auto* fwd_in_meta = egr::EagerUtils::nullable_autograd_meta(fwd_in);
  PADDLE_ENFORCE_LE(
      (slot_rank + 1),
      bwd_out_meta_.size(),
      common::errors::InvalidArgument(
          "Slot Rank should less equal than bwd_out_meta_ size, "
          "since bwd_out_meta_ is designed to hold as same num as "
          "backward outputs."));
  auto& metas = bwd_out_meta_.at(slot_rank);
  // Init stop gradient vector before use to avoid push back
  if (metas.empty()) {
    metas.resize(1);
  }
  auto& meta = metas[0];
  if (VLOG_IS_ON(6) || FLAGS_enable_unique_name) {
    // Record the forward input tensor name
    meta.SetForwardTensorName(fwd_in.name());
  }

  // Set Stop_gradient
  if (fwd_in_meta) {
    meta.SetStopGradient(fwd_in_meta->StopGradient());
  } else {
    meta.SetStopGradient(true);
  }
  // Set Adj Edges
  if (fwd_in_meta && !fwd_in_meta->StopGradient()) {
    auto node = fwd_in_meta->GetMutableGradNode();
    if (!node || !node.get()) {
      fwd_in_meta->SetGradNode(
          std::make_shared<egr::GradNodeAccumulation>(fwd_in));
    }
    VLOG(5) << "Add Edges for slot: " << slot_rank << ", the Edge is from "
            << this->name() << "(" << this << ")"
            << " to " << fwd_in_meta->GetMutableGradNode()->name() << "("
            << fwd_in_meta->GetMutableGradNode().get() << ")";

    meta.SetEdge(fwd_in_meta->GetMutableGradNode(), fwd_in_meta->OutRankInfo());
  }
  // Record TensorMeta
  if (fwd_in.impl() && fwd_in.impl().get()) {
    if (phi::DenseTensor::classof(fwd_in.impl().get())) {
      // Only Copy Meta
      phi::DenseTensor* dense_tensor =
          static_cast<phi::DenseTensor*>(fwd_in.impl().get());
      PADDLE_ENFORCE_NE(
          dense_tensor->meta().dtype,
          phi::DataType::UNDEFINED,
          common::errors::Fatal("Attempting to copy DenseTensorMeta "
                                "with phi::DataType::UNDEFINED,"
                                "which is illegal."));
      meta.SetTensorMeta(dense_tensor->meta());
      meta.SetPlace(fwd_in.place());
    } else if (phi::distributed::DistTensor::classof(fwd_in.impl().get())) {
      const phi::distributed::DistTensor* dist_tensor =
          static_cast<phi::distributed::DistTensor*>(fwd_in.impl().get());
      const phi::DenseTensor& dense_tensor = dist_tensor->value();
      PADDLE_ENFORCE_NE(
          dense_tensor.meta().dtype,
          phi::DataType::UNDEFINED,
          common::errors::Fatal("Attempting to copy DenseTensorMeta "
                                "with phi::DataType::UNDEFINED,"
                                "which is illegal."));
      meta.SetTensorMeta(dense_tensor.meta());
      meta.SetPlace(fwd_in.place());
      // Set DistAttr
      // Forward input DistTensor could be uninitialized.
      PADDLE_ENFORCE_NE(
          dist_tensor->dist_attr().empty(),
          true,
          common::errors::InvalidArgument(
              "The forward input DistTensor's dist attr is empty."));
      auto dist_attr = dist_tensor->dist_attr();
      dist_attr.clean_partial_status();
      meta.SetDistAttr(dist_attr);
      meta.SetDistTensorGlobalDims(dist_tensor->dims());
      SetIsRunAutoParallel(true);
    } else if (phi::SparseCsrTensor::classof(fwd_in.impl().get())) {
      phi::SparseCsrTensor* sparse_tensor =
          static_cast<phi::SparseCsrTensor*>(fwd_in.impl().get());
      const phi::DenseTensor dense_tensor =
          static_cast<const phi::DenseTensor>(sparse_tensor->values());
      PADDLE_ENFORCE_NE(
          dense_tensor.dtype(),
          phi::DataType::UNDEFINED,
          common::errors::Fatal("Attempting to copy DenseTensorMeta "
                                "with phi::DataType::UNDEFINED,"
                                "which is illegal."));
      meta.SetTensorMeta(dense_tensor.meta());
      meta.SetPlace(fwd_in.place());
    } else if (phi::SparseCooTensor::classof(fwd_in.impl().get())) {
      phi::SparseCooTensor* sparse_tensor =
          static_cast<phi::SparseCooTensor*>(fwd_in.impl().get());
      const phi::DenseTensor dense_tensor =
          static_cast<const phi::DenseTensor>(sparse_tensor->values());
      PADDLE_ENFORCE_NE(
          dense_tensor.dtype(),
          phi::DataType::UNDEFINED,
          common::errors::Fatal("Attempting to copy DenseTensorMeta "
                                "with phi::DataType::UNDEFINED,"
                                "which is illegal."));
      meta.SetTensorMeta(dense_tensor.meta());
      meta.SetPlace(fwd_in.place());
    } else {
      VLOG(7)
          << "Unable to initialize the DenseTensorMeta of GradSlotMeta with "
             "non-DenseTensor argument.";
    }
  } else {
    VLOG(5) << "Unable to initialize the DenseTensorMeta because the Tensor "
               "is not initialized.";
  }
}

/*
special func for matmul_double_grad etc. dx exists when x and y exists,
if stop_gradient of y is true, dy is None who is matmul_grad's out_put, ddy is
None, so dx = ddy * dout should be None.
*/
void GradNodeBase::SetGradOutMeta(const paddle::Tensor& fwd_in,
                                  const AutogradMeta* fwd_out_meta,
                                  size_t slot_rank) {
  auto* fwd_in_meta = egr::EagerUtils::nullable_autograd_meta(fwd_in);
  PADDLE_ENFORCE_LE(
      (slot_rank + 1),
      bwd_out_meta_.size(),
      common::errors::InvalidArgument(
          "Slot Rank should less equal than bwd_out_meta_ size, "
          "since bwd_out_meta_ is designed to hold as same num as "
          "backward outputs."));
  auto& metas = bwd_out_meta_.at(slot_rank);
  // Init stop gradient vector before use to avoid push back
  if (metas.empty()) {
    metas.resize(1);
  }
  auto& meta = metas[0];
  if (VLOG_IS_ON(6) || FLAGS_enable_unique_name) {
    // Record the forward input tensor name
    meta.SetForwardTensorName(fwd_in.name());
  }
  // Set Stop_gradient
  if (fwd_in_meta && !fwd_in_meta->StopGradient() && fwd_out_meta) {
    meta.SetStopGradient(false);
  } else {
    meta.SetStopGradient(true);
  }
  // Set Adj Edges
  if (fwd_in_meta && !fwd_in_meta->StopGradient() && fwd_out_meta) {
    auto node = fwd_in_meta->GetMutableGradNode();
    if (!node || !node.get()) {
      fwd_in_meta->SetGradNode(
          std::make_shared<egr::GradNodeAccumulation>(fwd_in));
    }
    VLOG(5) << "Add Edges for slot: " << slot_rank << ", the Edge is from "
            << this->name() << "(" << this << ")"
            << " to " << fwd_in_meta->GetMutableGradNode()->name() << "("
            << fwd_in_meta->GetMutableGradNode().get() << ")";

    meta.SetEdge(fwd_in_meta->GetMutableGradNode(), fwd_in_meta->OutRankInfo());
  }
  // Record TensorMeta
  if (fwd_in.impl() && fwd_in.impl().get()) {
    if (phi::DenseTensor::classof(fwd_in.impl().get())) {
      // Only Copy Meta
      phi::DenseTensor* dense_tensor =
          static_cast<phi::DenseTensor*>(fwd_in.impl().get());
      PADDLE_ENFORCE_NE(
          dense_tensor->meta().dtype,
          phi::DataType::UNDEFINED,
          common::errors::Fatal("Attempting to copy DenseTensorMeta "
                                "with phi::DataType::UNDEFINED,"
                                "which is illegal."));
      meta.SetTensorMeta(dense_tensor->meta());
      meta.SetPlace(fwd_in.place());
    } else if (phi::distributed::DistTensor::classof(fwd_in.impl().get())) {
      // Only Copy Meta
      meta.SetDistAttr(
          static_cast<phi::distributed::DistTensor*>(fwd_in.impl().get())
              ->dist_attr());
      meta.SetDistTensorGlobalDims(
          static_cast<phi::distributed::DistTensor*>(fwd_in.impl().get())
              ->dims());
      SetIsRunAutoParallel(true);
      auto dense_tensor =
          static_cast<phi::distributed::DistTensor*>(fwd_in.impl().get())
              ->value();
      PADDLE_ENFORCE_NE(
          dense_tensor.meta().dtype,
          phi::DataType::UNDEFINED,
          common::errors::Fatal("Attempting to copy DenseTensorMeta "
                                "with phi::DataType::UNDEFINED,"
                                "which is illegal."));
      meta.SetTensorMeta(dense_tensor.meta());
      meta.SetPlace(fwd_in.place());
    }
  } else {
    VLOG(5) << "Unable to initialize the DenseTensorMeta of GradSlotMeta with "
               "non-DenseTensor argument.";
  }
}

/*
Special func for inplace ops in auto parallel. For now, this func is only used
in reshape_.
*/
void GradNodeBase::SetGradOutMeta(
    const paddle::Tensor& fwd_in,
    size_t slot_rank,
    const phi::distributed::TensorDistAttr& fwd_in_dist_attr,
    const phi::DDim& fwd_in_dims) {
  auto* fwd_in_meta = egr::EagerUtils::nullable_autograd_meta(fwd_in);
  PADDLE_ENFORCE_LE(
      (slot_rank + 1),
      bwd_out_meta_.size(),
      common::errors::InvalidArgument(
          "Slot Rank should less equal than bwd_out_meta_ size, "
          "since bwd_out_meta_ is designed to hold as same num as "
          "backward outputs."));
  auto& metas = bwd_out_meta_.at(slot_rank);
  // Init stop gradient vector before use to avoid push back
  if (metas.empty()) {
    metas.resize(1);
  }
  auto& meta = metas[0];
  if (VLOG_IS_ON(6) || FLAGS_enable_unique_name) {
    meta.SetForwardTensorName(fwd_in.name());
  }
  // Set Stop_gradient
  if (fwd_in_meta) {
    meta.SetStopGradient(fwd_in_meta->StopGradient());
  } else {
    meta.SetStopGradient(true);
  }
  // Set Adj Edges
  if (fwd_in_meta && !fwd_in_meta->StopGradient()) {
    auto node = fwd_in_meta->GetMutableGradNode();
    if (!node || !node.get()) {
      fwd_in_meta->SetGradNode(
          std::make_shared<egr::GradNodeAccumulation>(fwd_in));
    }
    VLOG(5) << "Add Edges for slot: " << slot_rank << ", the Edge is from "
            << this->name() << "(" << this << ")"
            << " to " << fwd_in_meta->GetMutableGradNode()->name() << "("
            << fwd_in_meta->GetMutableGradNode().get() << ")";

    meta.SetEdge(fwd_in_meta->GetMutableGradNode(), fwd_in_meta->OutRankInfo());
  }
  // Record TensorMeta
  if (fwd_in.impl() && fwd_in.impl().get()) {
    if (phi::distributed::DistTensor::classof(fwd_in.impl().get())) {
      const phi::distributed::DistTensor* dist_tensor =
          static_cast<phi::distributed::DistTensor*>(fwd_in.impl().get());
      const phi::DenseTensor& dense_tensor = dist_tensor->value();
      PADDLE_ENFORCE_NE(
          dense_tensor.meta().dtype,
          phi::DataType::UNDEFINED,
          common::errors::Fatal("Attempting to copy DenseTensorMeta "
                                "with phi::DataType::UNDEFINED,"
                                "which is illegal."));
      meta.SetTensorMeta(dense_tensor.meta());
      meta.SetPlace(fwd_in.place());
      // Set DistAttr
      // Forward input DistTensor could be uninitialized.
      PADDLE_ENFORCE_NE(
          dist_tensor->dist_attr().empty(),
          true,
          common::errors::InvalidArgument(
              "The forward input DistTensor's dist attr is empty."));
      auto dist_attr = fwd_in_dist_attr;
      dist_attr.clean_partial_status();
      meta.SetDistAttr(dist_attr);
      meta.SetDistTensorGlobalDims(fwd_in_dims);
      SetIsRunAutoParallel(true);
    } else {
      VLOG(7)
          << "Unable to initialize the DenseTensorMeta of GradSlotMeta with "
             "non-DistTensor argument.";
    }
  } else {
    VLOG(7) << "Unable to initialize the DenseTensorMeta because the Tensor "
               "is not initialized.";
  }
}

void GradNodeBase::SetGradOutMeta(const std::vector<paddle::Tensor>& fwd_in,
                                  size_t slot_rank) {
  size_t slot_size = fwd_in.size();
  PADDLE_ENFORCE_LE(
      slot_rank,
      (bwd_out_meta_.size() - 1),
      common::errors::InvalidArgument(
          "Slot Rank should less equal than bwd_out_meta_ size, "
          "since bwd_out_meta_ is designed to hold as same num as "
          "backward outputs."));
  auto& metas = bwd_out_meta_.at(slot_rank);
  // Init stop gradient vector before use to avoid push back
  if (metas.size() < slot_size) {
    metas.resize(slot_size);
  }
  for (size_t i = 0; i < slot_size; i++) {
    const auto& fwd_in_tensor = fwd_in[i];
    auto& meta = metas[i];
    if (VLOG_IS_ON(6) || FLAGS_enable_unique_name) {
      meta.SetForwardTensorName(fwd_in_tensor.name());
    }
    auto* fwd_in_meta = egr::EagerUtils::nullable_autograd_meta(fwd_in_tensor);
    // Set Stop_gradient
    if (fwd_in_meta) {
      meta.SetStopGradient(fwd_in_meta->StopGradient());
    }
    // Set Adj Edges
    if (fwd_in_meta && !fwd_in_meta->StopGradient()) {
      auto node = fwd_in_meta->GetMutableGradNode();
      if (!node || !node.get()) {
        fwd_in_meta->SetGradNode(
            std::make_shared<egr::GradNodeAccumulation>(fwd_in_tensor));
      }
      VLOG(5) << "Add Edges for slot: " << slot_rank << ", the Edge is from "
              << this->name() << "(" << this << ")"
              << " to " << fwd_in_meta->GetMutableGradNode()->name() << "("
              << fwd_in_meta->GetMutableGradNode().get() << ")";

      meta.SetEdge(fwd_in_meta->GetMutableGradNode(),
                   fwd_in_meta->OutRankInfo());
    }
    // Record TensorMeta
    if (fwd_in_tensor.impl() && fwd_in_tensor.impl().get()) {
      if (phi::DenseTensor::classof(fwd_in_tensor.impl().get())) {
        phi::DenseTensor* dense_tensor =
            static_cast<phi::DenseTensor*>(fwd_in_tensor.impl().get());
        PADDLE_ENFORCE_NE(
            dense_tensor->dtype(),
            phi::DataType::UNDEFINED,
            common::errors::Fatal("Attempting to copy DenseTensorMeta "
                                  "with phi::DataType::UNDEFINED,"
                                  "which is illegal."));
        meta.SetTensorMeta(dense_tensor->meta());
        meta.SetPlace(fwd_in_tensor.place());
      } else if (phi::distributed::DistTensor::classof(
                     fwd_in_tensor.impl().get())) {
        meta.SetDistAttr(static_cast<phi::distributed::DistTensor*>(
                             fwd_in_tensor.impl().get())
                             ->dist_attr());
        meta.SetDistTensorGlobalDims(static_cast<phi::distributed::DistTensor*>(
                                         fwd_in_tensor.impl().get())
                                         ->dims());
        SetIsRunAutoParallel(true);
        auto dense_tensor = static_cast<phi::distributed::DistTensor*>(
                                fwd_in_tensor.impl().get())
                                ->value();
        PADDLE_ENFORCE_NE(
            dense_tensor.dtype(),
            phi::DataType::UNDEFINED,
            common::errors::Fatal("Attempting to copy DenseTensorMeta "
                                  "with phi::DataType::UNDEFINED,"
                                  "which is illegal."));
        meta.SetTensorMeta(dense_tensor.meta());
        meta.SetPlace(fwd_in_tensor.place());
      }
    } else {
      VLOG(7)
          << "Unable to initialize the DenseTensorMeta of GradSlotMeta with "
             "non-DenseTensor argument.";
    }
  }
}

void GradNodeBase::SetGradOutMeta(
    const std::vector<const paddle::Tensor*>& fwd_in, size_t slot_rank) {
  size_t slot_size = fwd_in.size();
  PADDLE_ENFORCE_LE(
      slot_rank,
      (bwd_out_meta_.size() - 1),
      common::errors::InvalidArgument(
          "Slot Rank should less equal than bwd_out_meta_ size, "
          "since bwd_out_meta_ is designed to hold as same num as "
          "backward outputs."));
  auto& metas = bwd_out_meta_.at(slot_rank);
  // Init stop gradient vector before use to avoid push back
  if (metas.size() < slot_size) {
    metas.resize(slot_size);
  }
  for (size_t i = 0; i < slot_size; i++) {
    const auto& fwd_in_tensor = (*fwd_in[i]);
    auto& meta = metas[i];
    if (VLOG_IS_ON(6) || FLAGS_enable_unique_name) {
      meta.SetForwardTensorName(fwd_in_tensor.name());
    }
    auto* fwd_in_meta = egr::EagerUtils::nullable_autograd_meta(fwd_in_tensor);
    // Set Stop_gradient
    if (fwd_in_meta) {
      meta.SetStopGradient(fwd_in_meta->StopGradient());
    }
    // Set Adj Edges
    if (fwd_in_meta && !fwd_in_meta->StopGradient()) {
      auto node = fwd_in_meta->GetMutableGradNode();
      if (!node || !node.get()) {
        fwd_in_meta->SetGradNode(
            std::make_shared<egr::GradNodeAccumulation>(fwd_in_tensor));
      }
      VLOG(5) << "Add Edges for slot: " << slot_rank << ", the Edge is from "
              << this->name() << "(" << this << ")"
              << " to " << fwd_in_meta->GetMutableGradNode()->name() << "("
              << fwd_in_meta->GetMutableGradNode().get() << ")";

      meta.SetEdge(fwd_in_meta->GetMutableGradNode(),
                   fwd_in_meta->OutRankInfo());
    }
    // Record TensorMeta
    if (fwd_in_tensor.impl() && fwd_in_tensor.impl().get()) {
      if (phi::DenseTensor::classof(fwd_in_tensor.impl().get())) {
        // Only Copy Meta
        phi::DenseTensor* dense_tensor =
            static_cast<phi::DenseTensor*>(fwd_in_tensor.impl().get());
        PADDLE_ENFORCE_NE(
            dense_tensor->dtype(),
            phi::DataType::UNDEFINED,
            common::errors::Fatal("Attempting to copy DenseTensorMeta "
                                  "with phi::DataType::UNDEFINED,"
                                  "which is illegal."));
        meta.SetTensorMeta(dense_tensor->meta());
        meta.SetPlace(fwd_in_tensor.place());
      } else if (phi::distributed::DistTensor::classof(
                     fwd_in_tensor.impl().get())) {
        // Only Copy Meta
        meta.SetDistAttr(static_cast<phi::distributed::DistTensor*>(
                             fwd_in_tensor.impl().get())
                             ->dist_attr());
        meta.SetDistTensorGlobalDims(static_cast<phi::distributed::DistTensor*>(
                                         fwd_in_tensor.impl().get())
                                         ->dims());
        SetIsRunAutoParallel(true);
        auto dense_tensor = static_cast<phi::distributed::DistTensor*>(
                                fwd_in_tensor.impl().get())
                                ->value();
        PADDLE_ENFORCE_NE(
            dense_tensor.dtype(),
            phi::DataType::UNDEFINED,
            common::errors::Fatal("Attempting to copy DenseTensorMeta "
                                  "with phi::DataType::UNDEFINED,"
                                  "which is illegal."));
        meta.SetTensorMeta(dense_tensor.meta());
        meta.SetPlace(fwd_in_tensor.place());
      }
    } else {
      VLOG(7)
          << "Unable to initialize the DenseTensorMeta of GradSlotMeta with "
             "non-DenseTensor argument.";
    }
  }
}

void GradNodeBase::SetDefaultGradInOutMeta() {
  PADDLE_ENFORCE((bwd_out_meta_.size() == 1) && (bwd_in_meta_.size() == 1),
                 common::errors::PreconditionNotMet(
                     "We can only support 1 input and 1 output in default grad "
                     "meta setter, other size of inputs and outputs should "
                     "create with Setter and Getters"));
  // Default stop_gradient is false and slot id is 0, slot size is 1;
  bwd_out_meta_[0].resize(1);
  bwd_in_meta_[0].resize(1);
}

int64_t GradNodeBase::RegisterGradientHook(
    size_t slot_id, size_t rank, std::shared_ptr<egr::TensorHook>&& hook) {
  gradient_hooks_.emplace(next_hook_id_,
                          std::make_tuple(slot_id, rank, std::move(hook)));
  return next_hook_id_++;
}

paddle::small_vector<std::vector<paddle::Tensor>, kSlotSmallVectorSize>
GradNodeBase::ApplyGradientHooks(
    const paddle::small_vector<std::vector<paddle::Tensor>,
                               kSlotSmallVectorSize>& tensors) {
  paddle::small_vector<std::vector<paddle::Tensor>, kSlotSmallVectorSize> outs(
      tensors.size());
  for (auto& hook_pair : gradient_hooks_) {
    size_t slot_id = std::get<0>(hook_pair.second);
    size_t rank = std::get<1>(hook_pair.second);

    auto hook = std::get<2>(hook_pair.second);

    PADDLE_ENFORCE(slot_id < tensors.size(),
                   common::errors::Fatal(
                       "Slot_id from registered hook should be smaller than "
                       "slot size of grad_tensors"));

    PADDLE_ENFORCE(rank < tensors[slot_id].size(),
                   common::errors::Fatal(
                       "rank of slot %d from registered hook should be smaller "
                       "than rank size of grad_tensors",
                       slot_id));

    std::vector<paddle::Tensor>& slot_out = outs[slot_id];
    slot_out.resize(tensors[slot_id].size());
    paddle::Tensor& out = slot_out[rank];
    if (!out.defined() || !out.has_allocation()) {
      out = (*hook)(tensors[slot_id][rank]);
    } else {
      // If more than one hook is registered, the input to the next hook func
      // should be the output of the previous hook
      out = (*hook)(out);
    }
  }

  for (size_t i = 0; i < outs.size(); i++) {
    if (outs[i].empty() && (!tensors[i].empty())) {
      outs[i].resize(tensors[i].size());
    }
    // TODO(Jiabin): Optimize this if we only add hook slot by slot
    for (size_t j = 0; j < outs[i].size(); j++) {
      if (!outs[i][j].defined() || !outs[i][j].initialized()) {
        outs[i][j] = tensors[i][j];
      }
      CheckTensor(tensors[i][j], outs[i][j]);
    }
  }

  return outs;
}

void GradNodeBase::HandleComplexGradToRealGrad(
    paddle::small_vector<std::vector<paddle::Tensor>, kSlotSmallVectorSize>*
        out_grads) {
  for (size_t slot_id = 0; slot_id < out_grads->size(); slot_id++) {
    const std::vector<paddle::Tensor>& slot_out_grads = (*out_grads)[slot_id];
    for (size_t rank_id = 0; rank_id < slot_out_grads.size(); rank_id++) {
      if (bwd_out_meta_[slot_id].size() == 0) continue;
      const GradSlotMeta& slot_meta = bwd_out_meta_[slot_id][rank_id];
      PADDLE_ENFORCE(
          slot_meta.HasTensorMeta() > 0,
          common::errors::Fatal(
              "We require TensorMeta in GradInputMeta() to obtain forward data "
              "types."
              "However, no TensorMeta is detected in bwd_out_meta_."));

      auto fwd_data_type = paddle::framework::TransToProtoVarType(
          slot_meta.GetTensorMeta().dtype);
      const paddle::Tensor& grad = slot_out_grads[rank_id];

      if (paddle::framework::IsComplexType(fwd_data_type)) continue;
      if (!grad.impl()) continue;

      // Only Handle Complex To Real for DenseTensor for now
      if (phi::DenseTensor::classof(grad.impl().get())) {
        phi::DenseTensor* grad_dense_tensor =
            static_cast<phi::DenseTensor*>(grad.impl().get());

        auto curr_data_type =
            paddle::framework::TransToProtoVarType(grad_dense_tensor->type());
        if (!paddle::framework::IsComplexType(curr_data_type)) continue;

        // Convert Complex GradOut to Real
        auto out = std::make_shared<phi::DenseTensor>();
        paddle::framework::TransComplexToReal(
            fwd_data_type, curr_data_type, *grad_dense_tensor, out.get());

        (*out_grads)[slot_id][rank_id].set_impl(out);
      } else if (phi::distributed::DistTensor::classof(grad.impl().get())) {
        auto grad_dense_tensor =
            static_cast<phi::distributed::DistTensor*>(grad.impl().get())
                ->value();

        auto curr_data_type =
            paddle::framework::TransToProtoVarType(grad_dense_tensor.type());
        if (!paddle::framework::IsComplexType(curr_data_type)) continue;
        if (grad_dense_tensor.dims().size() == -1) continue;

        // Convert Complex GradOut to Real
        auto out = std::make_shared<phi::DenseTensor>();
        paddle::framework::TransComplexToReal(
            fwd_data_type, curr_data_type, grad_dense_tensor, out.get());

        *(static_cast<phi::distributed::DistTensor*>(
              (*out_grads)[slot_id][rank_id].impl().get())
              ->unsafe_mutable_value()) = *(out.get());
      }
    }
  }
}

std::vector<std::shared_ptr<GradNodeBase>> GradNodeBase::NextFunctions() {
  std::vector<std::shared_ptr<GradNodeBase>> next_nodes;
  const paddle::small_vector<std::vector<GradSlotMeta>, kSlotSmallVectorSize>&
      metas = OutputMeta();

  for (const auto& meta_list : metas) {
    for (const GradSlotMeta& meta : meta_list) {
      const auto& edge = meta.GetEdge();
      std::shared_ptr<GradNodeBase> next_node = edge.GetMutableGradNode();
      next_nodes.push_back(next_node);
    }
  }

  return next_nodes;
}

uintptr_t GradNodeBase::GetPtr() const {
  return reinterpret_cast<uintptr_t>(this);
}

int64_t GradNodeBase::RegisterNodePostHook(
    std::shared_ptr<NodePostHookBase>&& hook) {
  post_hooks_.emplace(next_post_hook_id_, std::move(hook));
  return next_post_hook_id_++;
}

bool GradNodeBase::RemoveNodePostHook(int64_t hook_id) {
  auto remove_cnt = post_hooks_.erase(hook_id);
  if (remove_cnt == 0) {
    return false;
  }
  return true;
}

bool GradNodeBase::HasNodePostHook() { return !post_hooks_.empty(); }

paddle::small_vector<std::vector<paddle::Tensor>, egr::kSlotSmallVectorSize>
GradNodeBase::ApplyNodePostHooks(
    const paddle::small_vector<std::vector<paddle::Tensor>,
                               egr::kSlotSmallVectorSize>& grad_outputs,
    const paddle::small_vector<std::vector<paddle::Tensor>,
                               egr::kSlotSmallVectorSize>& grad_inputs) {
  paddle::small_vector<std::vector<paddle::Tensor>, kSlotSmallVectorSize> outs =
      grad_outputs;
  for (auto& iter : post_hooks_) {
    auto hook = iter.second;
    outs = (*hook)(outs, grad_inputs);
  }

  for (size_t i = 0; i < outs.size(); i++) {
    if (outs[i].empty() && (!grad_outputs[i].empty())) {
      outs[i].resize(grad_outputs[i].size());
    }

    for (size_t j = 0; j < outs[i].size(); j++) {
      if (!outs[i][j].defined() || !outs[i][j].initialized()) {
        outs[i][j] = grad_outputs[i][j];
      } else {
        CheckTensor(grad_outputs[i][j], outs[i][j]);
      }
    }
  }

  return outs;
}

void Edge::SetGradNode(const std::shared_ptr<GradNodeBase>& node) {
  VLOG(7) << "Resetting Edge(" << this << ")'s Grad Node"
          << " from " << (grad_node_ ? grad_node_->name() : "nullptr") << "("
          << grad_node_.get() << ") to " << (node ? node->name() : "nullptr")
          << "(" << node.get() << ")";
  grad_node_ = node;
}

}  // namespace egr