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

#pragma once
#include <functional>

#include "paddle/fluid/eager/api/utils/global_utils.h"
#include "paddle/fluid/framework/custom_operator_utils.h"
#include "paddle/fluid/framework/new_executor/instruction/custom_kernel_instruction.h"
#include "paddle/fluid/framework/python_operator.h"
#include "paddle/fluid/pir/dialect/distributed/ir/dist_tools.h"
#include "paddle/fluid/pir/dialect/distributed/ir/dist_type.h"
#include "paddle/fluid/pir/dialect/operator/ir/api_builder.h"
#include "paddle/fluid/pir/dialect/operator/ir/manual_api.h"
#include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
#include "paddle/fluid/pir/dialect/operator/ir/pd_api.h"
#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
#include "paddle/fluid/pybind/eager_utils.h"
#include "paddle/fluid/pybind/exception.h"
#include "paddle/fluid/pybind/op_callstack_utils.h"
#include "paddle/fluid/pybind/op_function_common.h"
#include "paddle/fluid/pybind/static_op_function.h"
#include "paddle/phi/api/ext/native_meta_tensor.h"
#include "paddle/phi/common/int_array.h"
#include "paddle/phi/core/enforce.h"
#include "paddle/phi/infermeta/spmd_rules/rules.h"
#include "paddle/pir/include/core/attribute.h"
#include "paddle/pir/include/core/builtin_op.h"

namespace paddle {

namespace pybind {
static PyObject *static_api_parameter(PyObject *self,
                                      PyObject *args,
                                      PyObject *kwargs) {
  try {
    VLOG(6) << "Add parameter op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);

    // Parse Attributes
    PyObject *name_obj = PyTuple_GET_ITEM(args, 0);
    std::string name = CastPyArg2String(name_obj, "name", 0);
    // Call ir static api
    CallStackRecorder callstack_recorder("parameter");
    callstack_recorder.Record();
    auto static_api_out = paddle::dialect::parameter(name);
    callstack_recorder.AttachToOps();
    return ToPyObject(static_api_out);
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

static PyObject *static_api_set_parameter(PyObject *self,
                                          PyObject *args,
                                          PyObject *kwargs) {
  try {
    VLOG(6) << "Add set_parameter op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);

    // Get Value from args
    PyObject *parameter_obj = PyTuple_GET_ITEM(args, 0);
    auto parameter = CastPyArg2Value(parameter_obj, "parameter", 0, false);

    // Parse Attributes
    PyObject *name_obj = PyTuple_GET_ITEM(args, 1);
    std::string name = CastPyArg2String(name_obj, "name", 1);
    // Call ir static api
    CallStackRecorder callstack_recorder("set_parameter");
    callstack_recorder.Record();
    paddle::dialect::set_parameter(parameter, name);
    callstack_recorder.AttachToOps();
    Py_RETURN_NONE;
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

static PyObject *static_api_update_parameter(PyObject *self,
                                             PyObject *args,
                                             PyObject *kwargs) {
  try {
    VLOG(6) << "Add uodata_parameter op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);

    // Get Value from args
    PyObject *parameter_obj = PyTuple_GET_ITEM(args, 0);
    auto parameter = CastPyArg2Value(parameter_obj, "parameter", 0, false);

    // Parse Attributes
    PyObject *name_obj = PyTuple_GET_ITEM(args, 1);
    std::string name = CastPyArg2String(name_obj, "name", 1);
    // Call ir static api
    CallStackRecorder callstack_recorder("uodata_parameter");
    callstack_recorder.Record();
    paddle::dialect::update_parameter(parameter, name);
    callstack_recorder.AttachToOps();
    Py_RETURN_NONE;
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

static PyObject *static_api_set_persistable_value(PyObject *self,
                                                  PyObject *args,
                                                  PyObject *kwargs) {
  try {
    VLOG(6) << "Add shadow_output op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);

    // Get OpResult from args
    PyObject *persist_value_obj = PyTuple_GET_ITEM(args, 0);
    auto persist_value =
        CastPyArg2Value(persist_value_obj, "persist_value", 0, false);

    // Parse Attributes
    PyObject *name_obj = PyTuple_GET_ITEM(args, 1);
    std::string name = CastPyArg2String(name_obj, "name", 1);
    // Call ir static api
    CallStackRecorder callstack_recorder("shadow_output");
    callstack_recorder.Record();
    paddle::dialect::shadow_output(persist_value, name);
    callstack_recorder.AttachToOps();
    Py_RETURN_NONE;
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

PyObject *static_api_full(PyObject *self, PyObject *args, PyObject *kwargs) {
  try {
    VLOG(6) << "Add full op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);

    // Parse Attributes
    PyObject *shape_obj = PyTuple_GET_ITEM(args, 0);
    PyObject *value_obj = PyTuple_GET_ITEM(args, 1);
    PyObject *dtype_obj = PyTuple_GET_ITEM(args, 2);
    PyObject *place_obj = PyTuple_GET_ITEM(args, 3);

    DataType dtype = CastPyArg2DataTypeDirectly(dtype_obj, "full", 2);
    Place place = CastPyArg2Place(place_obj, "full", 3);

    if (!PyObject_CheckIRValue(shape_obj) &&
        !PyObject_CheckIRVectorOfValue(shape_obj) &&
        !PyObject_CheckIRValue(value_obj)) {
      std::vector<int64_t> shape = CastPyArg2Longs(shape_obj, "full", 0);
      if (PyComplex_Check(value_obj)) {
        phi::dtype::complex<float> complex_value =
            CastPyArg2Complex(value_obj, "full", 1);
        CallStackRecorder callstack_recorder("full");
        callstack_recorder.Record();
        auto static_api_out = paddle::dialect::full(
            shape, complex_value.real, complex_value.imag, dtype, place);
        callstack_recorder.AttachToOps();
        return ToPyObject(static_api_out);
      } else {
        double value = CastPyArg2Double(value_obj, "full", 1);
        CallStackRecorder callstack_recorder("full");
        callstack_recorder.Record();
        auto static_api_out = paddle::dialect::full(shape, value, dtype, place);
        callstack_recorder.AttachToOps();
        return ToPyObject(static_api_out);
      }
    } else {
      pir::Value shape, value;

      if (PyObject_CheckIRValue(shape_obj)) {
        shape = CastPyArg2Value(shape_obj, "full", 0, false);
      } else if (PyObject_CheckIRVectorOfValue(shape_obj)) {
        std::vector<pir::Value> shape_tmp =
            CastPyArg2VectorOfValue(shape_obj, "full", 0, false);
        shape = paddle::dialect::stack(shape_tmp, 0);
      } else {
        std::vector<int64_t> shape_tmp = CastPyArg2Longs(shape_obj, "full", 0);
        shape = paddle::dialect::full_int_array(
            shape_tmp, DataType::INT64, CPUPlace());
      }

      if (PyObject_CheckIRValue(value_obj)) {
        value = CastPyArg2Value(value_obj, "full", 1, false);
      } else {
        if (PyComplex_Check(value_obj)) {
          phi::dtype::complex<float> complex_value_tmp =
              CastPyArg2Complex(value_obj, "full", 1);
          value = paddle::dialect::full(std::vector<int64_t>{1},
                                        complex_value_tmp.real,
                                        complex_value_tmp.imag,
                                        dtype,
                                        place);
        } else {
          double value_tmp = CastPyArg2Double(value_obj, "full", 1);
          value = paddle::dialect::full(std::vector<int64_t>{1},
                                        value_tmp,
                                        DataType::FLOAT32,
                                        CPUPlace());
        }
      }

      CallStackRecorder callstack_recorder("full_with_tensor");
      callstack_recorder.Record();
      auto static_api_out =
          paddle::dialect::full_with_tensor(value, shape, dtype);
      callstack_recorder.AttachToOps();

      return ToPyObject(static_api_out);
    }
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

static PyObject *static_api_create_array(PyObject *self,
                                         PyObject *args,
                                         PyObject *kwargs) {
  try {
    VLOG(6) << "Add create_array op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);

    // Get dtype from args
    PyObject *dtype_obj = PyTuple_GET_ITEM(args, 0);
    DataType dtype = CastPyArg2DataTypeDirectly(dtype_obj, "create_array", 0);

    // Call ir static api
    CallStackRecorder callstack_recorder("create_array");
    callstack_recorder.Record();
    auto static_api_out = paddle::dialect::create_array(dtype);
    callstack_recorder.AttachToOps();

    return ToPyObject(static_api_out);
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

static PyObject *static_api_create_array_like(PyObject *self,
                                              PyObject *args,
                                              PyObject *kwargs) {
  try {
    VLOG(6) << "Add create_array_like op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);

    // Get Value from args
    PyObject *input_obj = PyTuple_GET_ITEM(args, 0);
    auto input = CastPyArg2Value(input_obj, "create_array_like", 0, false);

    // Parse Attributes
    PyObject *value_obj = PyTuple_GET_ITEM(args, 1);
    float value = CastPyArg2Float(value_obj, "create_array_like", 1);

    // Call ir static api
    CallStackRecorder callstack_recorder("create_array_like");
    callstack_recorder.Record();
    auto static_api_out = paddle::dialect::create_array_like(input, value);
    callstack_recorder.AttachToOps();
    return ToPyObject(static_api_out);
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

static PyObject *static_api_array_length(PyObject *self,
                                         PyObject *args,
                                         PyObject *kwargs) {
  try {
    VLOG(6) << "Add array_length op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);

    // Get Value from args
    PyObject *x_obj = PyTuple_GET_ITEM(args, 0);
    auto x = CastPyArg2Value(x_obj, "array_length", 0, false);

    // Call ir static api
    CallStackRecorder callstack_recorder("array_length");
    callstack_recorder.Record();
    auto static_api_out = paddle::dialect::array_length(x);
    callstack_recorder.AttachToOps();

    return ToPyObject(static_api_out);
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}
static PyObject *static_api_array_read(PyObject *self,
                                       PyObject *args,
                                       PyObject *kwargs) {
  try {
    VLOG(6) << "Add array_read op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);

    // Get Value from args
    PyObject *array_obj = PyTuple_GET_ITEM(args, 0);
    auto array = CastPyArg2Value(array_obj, "array_read", 0, false);

    PyObject *i_obj = PyTuple_GET_ITEM(args, 1);
    pir::Value i;
    if (PyObject_CheckIRValue(i_obj)) {
      i = CastPyArg2Value(i_obj, "array_read", 1, false);
    } else {
      int64_t i_tmp = CastPyArg2Int(i_obj, "array_read", 1);
      i = paddle::dialect::full(
          std::vector<int64_t>{1}, i_tmp, DataType::INT64, CPUPlace());
    }

    // Call ir static api
    CallStackRecorder callstack_recorder("array_read");
    callstack_recorder.Record();
    auto static_api_out = paddle::dialect::array_read(array, i);
    callstack_recorder.AttachToOps();

    return ToPyObject(static_api_out);
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

static PyObject *static_api_fetch(PyObject *self,
                                  PyObject *args,
                                  PyObject *kwargs) {
  try {
    VLOG(6) << "Add fetch op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);

    // Get Value from args
    PyObject *value_obj = PyTuple_GET_ITEM(args, 0);
    auto value = CastPyArg2Value(value_obj, "fetch", 0, false);

    std::string name = CastPyArg2AttrString(PyTuple_GET_ITEM(args, 1), 1);
    int col = CastPyArg2Int(PyTuple_GET_ITEM(args, 2), "array_read", 2);

    // Call ir static api
    CallStackRecorder callstack_recorder("fetch");
    callstack_recorder.Record();
    auto static_api_out = paddle::dialect::fetch(value, name, col);
    callstack_recorder.AttachToOps();

    return ToPyObject(static_api_out);
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

static PyObject *static_api_array_write_(PyObject *self,
                                         PyObject *args,
                                         PyObject *kwargs) {
  try {
    VLOG(6) << "Add array_write_ op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);

    // Get Value from args
    PyObject *array_obj = PyTuple_GET_ITEM(args, 0);
    auto array = CastPyArg2Value(array_obj, "array_write_", 0, false);
    PyObject *x_obj = PyTuple_GET_ITEM(args, 1);
    auto x = CastPyArg2Value(x_obj, "array_write_", 1, false);
    PyObject *i_obj = PyTuple_GET_ITEM(args, 2);
    pir::Value i;
    if (PyObject_CheckIRValue(i_obj)) {
      i = CastPyArg2Value(i_obj, "array_write_", 2, false);
    } else {
      int64_t i_tmp = CastPyArg2Int(i_obj, "array_write_", 2);
      i = paddle::dialect::full(
          std::vector<int64_t>{1}, i_tmp, DataType::INT64, CPUPlace());
    }

    // Call ir static api
    CallStackRecorder callstack_recorder("array_write_");
    callstack_recorder.Record();
    auto static_api_out = paddle::dialect::array_write_(array, x, i);
    callstack_recorder.AttachToOps();

    return ToPyObject(static_api_out);
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

static PyObject *static_api_array_to_tensor(PyObject *self,
                                            PyObject *args,
                                            PyObject *kwargs) {
  try {
    VLOG(6) << "Add array_to_tensor op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);

    // Get Value from args
    PyObject *x_obj = PyTuple_GET_ITEM(args, 0);
    pir::Value x;
    if (PyObject_CheckIRValue(x_obj)) {
      x = CastPyArg2Value(x_obj, "array_to_tensor", 0, false);
    } else if (PyObject_CheckIRVectorOfValue(x_obj)) {
      std::vector<pir::Value> x_tmp =
          CastPyArg2VectorOfValue(x_obj, "array_to_tensor", 0, false);
      if (x_tmp.size() != 1) {
        PADDLE_THROW(common::errors::InvalidArgument(
            "Input x expects only one input, but %d are given.",
            x_tmp.size()));  // NOLINT
      }
      x = x_tmp[0];
    }

    PyObject *axis_obj = PyTuple_GET_ITEM(args, 1);
    auto axis = CastPyArg2Int(axis_obj, "array_to_tensor", 1);

    PyObject *use_stack_obj = PyTuple_GET_ITEM(args, 2);
    auto use_stack = CastPyArg2Boolean(use_stack_obj, "array_to_tensor", 2);

    // Call ir static api
    CallStackRecorder callstack_recorder("array_to_tensor");
    callstack_recorder.Record();
    auto static_api_out = paddle::dialect::array_to_tensor(x, axis, use_stack);
    callstack_recorder.AttachToOps();

    return ToPyObject(static_api_out);
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

PyObject *static_api_add_n_array(PyObject *self,
                                 PyObject *args,
                                 PyObject *kwargs) {
  try {
    VLOG(6) << "Add add_n_array op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);

    // Get Value from args
    PyObject *inputs_obj = PyTuple_GET_ITEM(args, 0);
    auto inputs = CastPyArg2VectorOfValue(inputs_obj, "add_n", 0, false);

    CallStackRecorder callstack_recorder("add_n_array");
    callstack_recorder.Record();
    auto static_api_out = paddle::dialect::add_n_array(inputs);
    callstack_recorder.AttachToOps();

    return ToPyObject(static_api_out);
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

static PyObject *static_api_slice_array(PyObject *self,
                                        PyObject *args,
                                        PyObject *kwargs) {
  try {
    VLOG(6) << "Add slice_array op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);

    // Get Value from args
    PyObject *input_obj = PyTuple_GET_ITEM(args, 0);
    auto input = CastPyArg2Value(input_obj, "slice_array", 0, false);

    PyObject *starts_obj = PyTuple_GET_ITEM(args, 1);
    pir::Value starts;
    if (PyObject_CheckIRValue(starts_obj)) {
      starts = CastPyArg2Value(starts_obj, "slice_array", 1, false);
    } else if (PyObject_CheckIRVectorOfValue(starts_obj)) {
      std::vector<pir::Value> starts_tmp =
          CastPyArg2VectorOfValue(starts_obj, "slice_array", 1, false);
      starts = paddle::dialect::stack(starts_tmp, /*axis*/ 0);
    } else {
      std::vector<int64_t> starts_tmp =
          CastPyArg2Longs(starts_obj, "slice_array", 1);
      starts = paddle::dialect::full_int_array(
          starts_tmp, DataType::INT64, CPUPlace());
    }

    PyObject *ends_obj = PyTuple_GET_ITEM(args, 2);
    pir::Value ends;
    if (PyObject_CheckIRValue(ends_obj)) {
      ends = CastPyArg2Value(ends_obj, "slice_array", 2, false);
    } else if (PyObject_CheckIRVectorOfValue(ends_obj)) {
      std::vector<pir::Value> ends_tmp =
          CastPyArg2VectorOfValue(ends_obj, "slice_array", 2, false);
      ends = paddle::dialect::stack(ends_tmp, /*axis*/ 0);
    } else {
      std::vector<int64_t> ends_tmp =
          CastPyArg2Longs(ends_obj, "slice_array", 2);
      ends = paddle::dialect::full_int_array(
          ends_tmp, DataType::INT64, CPUPlace());
    }

    // Call ir static api
    CallStackRecorder callstack_recorder("slice_array");
    callstack_recorder.Record();
    auto static_api_out = paddle::dialect::slice_array(input, starts, ends);
    callstack_recorder.AttachToOps();

    return ToPyObject(static_api_out);
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

static PyObject *static_api_slice_array_dense(PyObject *self,
                                              PyObject *args,
                                              PyObject *kwargs) {
  try {
    VLOG(6) << "Add slice_array_dense op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);

    // Get Value from args
    PyObject *input_obj = PyTuple_GET_ITEM(args, 0);
    auto input = CastPyArg2Value(input_obj, "slice_array_dense", 0, false);

    PyObject *starts_obj = PyTuple_GET_ITEM(args, 1);
    pir::Value starts;
    if (PyObject_CheckIRValue(starts_obj)) {
      starts = CastPyArg2Value(starts_obj, "slice_array_dense", 1, false);
    } else if (PyObject_CheckIRVectorOfValue(starts_obj)) {
      std::vector<pir::Value> starts_tmp =
          CastPyArg2VectorOfValue(starts_obj, "slice_array_dense", 1, false);
      starts = paddle::dialect::stack(starts_tmp, /*axis*/ 0);

    } else {
      std::vector<int64_t> starts_tmp =
          CastPyArg2Longs(starts_obj, "slice_array_dense", 1);
      starts = paddle::dialect::full_int_array(
          starts_tmp, DataType::INT64, CPUPlace());
    }
    // Call ir static api
    CallStackRecorder callstack_recorder("slice_array_dense");
    callstack_recorder.Record();
    auto static_api_out = paddle::dialect::slice_array_dense(input, starts);
    callstack_recorder.AttachToOps();

    return ToPyObject(static_api_out);
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

extern PyObject *eager_api_run_custom_op(PyObject *self,
                                         PyObject *args,
                                         PyObject *kwargs);

static PyObject *static_api_run_custom_op(PyObject *self,
                                          PyObject *args,
                                          PyObject *kwargs) {
  std::string op_type = CastPyArg2AttrString(PyTuple_GET_ITEM(args, 0), 0);
  VLOG(7) << "Get things from python for Custom Op: " << op_type;
  const auto &meta_info_map = OpMetaInfoMap::Instance().GetMap();
  PADDLE_ENFORCE_NE(meta_info_map.find(op_type),
                    meta_info_map.end(),
                    common::errors::NotFound(
                        "Can't find %s in Eager OpMetaInfoMap which should be "
                        "created by LoadOpMetaInfoAndRegisterOp, please make "
                        "sure you registered your op first and try again. ",
                        op_type));

  const auto &vec_map = meta_info_map.at(op_type);
  const auto &inputs = paddle::OpMetaInfoHelper::GetInputs(vec_map[0]);
  const auto &attrs = paddle::OpMetaInfoHelper::GetAttrs(vec_map[0]);
  const auto &outputs = paddle::OpMetaInfoHelper::GetOutputs(vec_map[0]);
  const auto &inplace_map = paddle::OpMetaInfoHelper::GetInplaceMap(vec_map[0]);
  const auto &inplace_reverse_map =
      paddle::OpMetaInfoHelper::GetInplaceReverseMap(vec_map[0]);
  auto infershape_func = OpMetaInfoHelper::GetInferShapeFn(vec_map[0]);
  auto inferdtype_func = OpMetaInfoHelper::GetInferDtypeFn(vec_map[0]);

  std::string pir_op_name = paddle::framework::kCustomDialectPrefix + op_type;
  if (!inplace_map.empty()) {
    pir_op_name += "_";
  }
  pir::IrContext *ctx = pir::IrContext::Instance();
  pir::OpInfo pir_info = ctx->GetRegisteredOpInfo(pir_op_name);
  pir::OperationArgument argument(pir_info);
  std::vector<pir::Value> argument_inputs;
  std::vector<pir::Type> argument_outputs;

  std::vector<std::vector<int64_t>> input_shapes;
  std::vector<DataType> input_dtypes;
  std::unordered_map<std::string, int> input_name2id_map;
  std::vector<std::vector<std::vector<int64_t>>> vec_input_shapes;
  std::vector<std::vector<DataType>> vec_input_dtypes;
  std::unordered_map<std::string, int> vec_input_name2id_map;
  std::vector<paddle::any> custom_attrs;
  int input_index = 0;
  int vec_input_index = 0;

  for (size_t i = 0; i < inputs.size(); ++i) {
    const auto &input = inputs.at(i);
    // Parse op_type first, so that use i + 1
    PyObject *obj = PyTuple_GET_ITEM(args, i + 1);
    // Emplace Py_None from python, this means optional inputs passed to C++,
    // use one un-initialized tensor to indicate both Tensor and
    // vector<Tensor> inputs.
    if (obj == Py_None) {
      VLOG(7) << "Add un-initialized tensor "
                 "because the optional input is None";
      if (paddle::framework::detail::IsDuplicableVar(input)) {
        std::vector<std::vector<int64_t>> vec_input_shape;
        std::vector<DataType> vec_input_dtype;
        vec_input_shapes.emplace_back(vec_input_shape);
        vec_input_dtypes.emplace_back(vec_input_dtype);
        vec_input_name2id_map[inputs[i]] = vec_input_index;
        vec_input_index++;
      } else {
        std::vector<int64_t> input_shape;
        DataType input_dtype = DataType::UNDEFINED;
        input_shapes.emplace_back(input_shape);
        input_dtypes.emplace_back(input_dtype);
        input_name2id_map[inputs[i]] = input_index;
        input_index++;
      }
      argument_inputs.emplace_back();
      continue;
    }
    if (paddle::framework::detail::IsDuplicableVar(input)) {
      std::vector<std::vector<int64_t>> tmp_input_shapes;
      std::vector<DataType> tmp_input_dtypes;
      vec_input_name2id_map[inputs[i]] = vec_input_index;
      vec_input_index++;
      std::vector<pir::Value> input_values =
          CastPyArg2VectorOfValue(obj, op_type, i + 1, false);
      for (auto &input_value : input_values) {
        paddle::dialect::DenseTensorType input_tensor =
            input_value.type().dyn_cast<paddle::dialect::DenseTensorType>();
        tmp_input_shapes.push_back(phi::vectorize(input_tensor.dims()));
        tmp_input_dtypes.push_back(
            paddle::dialect::TransToPhiDataType(input_tensor.dtype()));
      }
      vec_input_shapes.push_back(tmp_input_shapes);
      vec_input_dtypes.push_back(tmp_input_dtypes);
      auto combine_op = paddle::dialect::ApiBuilder::Instance()
                            .GetBuilder()
                            ->Build<pir::CombineOp>(input_values);
      argument_inputs.push_back(combine_op.out());
    } else {
      input_name2id_map[inputs[i]] = input_index;
      input_index++;
      pir::Value input_value =
          CastPyArg2Value(obj, op_type, i + 1, false);  // NOLINT
      paddle::dialect::DenseTensorType input_tensor =
          input_value.type().dyn_cast<paddle::dialect::DenseTensorType>();
      input_shapes.push_back(phi::vectorize(input_tensor.dims()));
      input_dtypes.push_back(
          paddle::dialect::TransToPhiDataType(input_tensor.dtype()));
      argument_inputs.push_back(input_value);
    }
  }
  argument.AddInputs(argument_inputs);

  // Parse op_type and inputs first, so that use 1 + inputs.size() + i
  int attr_start_idx = static_cast<int>(1 + inputs.size());
  for (size_t i = 0; i < attrs.size(); ++i) {
    const auto &attr = attrs.at(i);
    std::vector<std::string> attr_name_and_type = paddle::ParseAttrStr(attr);
    auto attr_type_str = attr_name_and_type[1];
    VLOG(7) << "Custom operator add attrs " << attr_name_and_type[0]
            << " to CustomOpKernelContext. Attribute type = " << attr_type_str;
    PyObject *obj = PyTuple_GET_ITEM(args, attr_start_idx + i);
    if (attr_type_str == "bool") {
      bool bool_attr = CastPyArg2AttrBoolean(obj, attr_start_idx + i);
      custom_attrs.push_back(bool_attr);  // NOLINT
      argument.AddAttribute(
          attr_name_and_type[0],
          pir::BoolAttribute::get(pir::IrContext::Instance(), bool_attr));
    } else if (attr_type_str == "int") {
      int int_attr = CastPyArg2AttrInt(obj, attr_start_idx + i);
      custom_attrs.push_back(int_attr);  // NOLINT
      argument.AddAttribute(
          attr_name_and_type[0],
          pir::Int32Attribute::get(pir::IrContext::Instance(), int_attr));
    } else if (attr_type_str == "float") {
      float float_attr = CastPyArg2AttrFloat(obj, attr_start_idx + i);
      custom_attrs.push_back(float_attr);  // NOLINT
      argument.AddAttribute(
          attr_name_and_type[0],
          pir::FloatAttribute::get(pir::IrContext::Instance(), float_attr));
    } else if (attr_type_str == "double") {
      double double_attr = CastPyArg2AttrDouble(obj, attr_start_idx + i);
      custom_attrs.push_back(double_attr);  // NOLINT
      argument.AddAttribute(
          attr_name_and_type[0],
          pir::DoubleAttribute::get(pir::IrContext::Instance(), double_attr));
    } else if (attr_type_str == "int64_t") {
      int64_t long_attr = CastPyArg2AttrLong(obj, attr_start_idx + i);
      custom_attrs.push_back(long_attr);  // NOLINT
      argument.AddAttribute(
          attr_name_and_type[0],
          pir::Int64Attribute::get(pir::IrContext::Instance(), long_attr));
    } else if (attr_type_str == "std::string") {
      std::string str_attr = CastPyArg2AttrString(obj, attr_start_idx + i);
      custom_attrs.push_back(str_attr);  // NOLINT
      argument.AddAttribute(
          attr_name_and_type[0],
          pir::StrAttribute::get(pir::IrContext::Instance(), str_attr));
    } else if (attr_type_str == "std::vector<int>") {
      std::vector<int> vec_int_attr =
          CastPyArg2VectorOfInt(obj, attr_start_idx + i);
      custom_attrs.push_back(vec_int_attr);
      std::vector<pir::Attribute> array_attr;
      for (size_t i = 0; i < static_cast<size_t>(vec_int_attr.size()); i++) {
        pir::Attribute attr = pir::Int32Attribute::get(
            pir::IrContext::Instance(), vec_int_attr[i]);
        array_attr.push_back(attr);
      }
      argument.AddAttribute(
          attr_name_and_type[0],
          pir::ArrayAttribute::get(pir::IrContext::Instance(), array_attr));
    } else if (attr_type_str == "std::vector<float>") {
      std::vector<float> vec_float_attr =
          CastPyArg2VectorOfFloat(obj, attr_start_idx + i);
      custom_attrs.push_back(vec_float_attr);
      std::vector<pir::Attribute> array_attr;
      for (size_t i = 0; i < static_cast<size_t>(vec_float_attr.size()); i++) {
        pir::Attribute attr = pir::FloatAttribute::get(
            pir::IrContext::Instance(), vec_float_attr[i]);
        array_attr.push_back(attr);
      }
      argument.AddAttribute(
          attr_name_and_type[0],
          pir::ArrayAttribute::get(pir::IrContext::Instance(), array_attr));
    } else if (attr_type_str == "std::vector<int64_t>") {
      std::vector<int64_t> vec_long_attr =
          CastPyArg2VectorOfInt64(obj, attr_start_idx + i);
      custom_attrs.push_back(vec_long_attr);  // NOLINT
      std::vector<pir::Attribute> array_attr;
      for (size_t i = 0; i < static_cast<size_t>(vec_long_attr.size()); i++) {
        pir::Attribute attr = pir::Int64Attribute::get(
            pir::IrContext::Instance(), vec_long_attr[i]);
        array_attr.push_back(attr);
      }
      argument.AddAttribute(
          attr_name_and_type[0],
          pir::ArrayAttribute::get(pir::IrContext::Instance(), array_attr));
    } else if (attr_type_str == "std::vector<std::string>") {
      std::vector<std::string> vec_str_attr =
          CastPyArg2VectorOfString(obj, attr_start_idx + i);
      custom_attrs.push_back(vec_str_attr);  // NOLINT
      std::vector<pir::Attribute> array_attr;
      for (size_t i = 0; i < static_cast<size_t>(vec_str_attr.size()); i++) {
        pir::Attribute attr =
            pir::StrAttribute::get(pir::IrContext::Instance(), vec_str_attr[i]);
        array_attr.push_back(attr);
      }
      argument.AddAttribute(
          attr_name_and_type[0],
          pir::ArrayAttribute::get(pir::IrContext::Instance(), array_attr));
    } else {
      PADDLE_THROW(common::errors::Unimplemented(
          "Unsupported `%s` type value as custom attribute now. "
          "Supported data types include `bool`, `int`, `float`, `double`, "
          "`int64_t`, `std::string`, `std::vector<int>`, "
          "`std::vector<float>`, `std::vector<int64_t>`, "
          "`std::vector<std::string>`, Please check whether "
          "the attribute data type and data type string are matched.",
          attr_type_str));
    }
  }

  paddle::framework::CheckDefaultInferShapeDtype(
      infershape_func, inferdtype_func, vec_map[0]);
  std::vector<std::vector<int64_t>> output_shapes =
      paddle::framework::RunInferShape(infershape_func,
                                       vec_map[0],
                                       input_shapes,
                                       input_name2id_map,
                                       vec_input_shapes,
                                       vec_input_name2id_map,
                                       custom_attrs);
  std::vector<DataType> output_dtypes =
      paddle::framework::RunInferDtype(inferdtype_func,
                                       vec_map[0],
                                       input_dtypes,
                                       input_name2id_map,
                                       vec_input_dtypes,
                                       vec_input_name2id_map,
                                       custom_attrs);
  dialect::ProcessMeshAttribute op_mesh;
  bool run_auto_parallel = false;
  std::vector<pir::Attribute> dist_result_attrs;
  phi::distributed::SpmdInfo spmd_info;
  if (dialect::HasDistInput(argument_inputs, &op_mesh)) {
    VLOG(7) << "Custom Op: " << op_type << " InferSPMD";
    run_auto_parallel = true;
    spmd_info = paddle::framework::RunInferSpmd(
        vec_map[0], op_type, op_mesh, argument_inputs, custom_attrs);
  }

  size_t all_values_num = 0;
  // output name -> value num (that output should hold)
  std::unordered_map<std::string, size_t> output_name2value_num;
  for (size_t i = 0; i < outputs.size(); ++i) {
    const auto &output = outputs.at(i);
    if (paddle::framework::detail::IsDuplicableVar(output)) {
      PADDLE_ENFORCE_NE(
          inplace_reverse_map.find(output),
          inplace_reverse_map.end(),
          common::errors::InvalidArgument(
              "Only support vector output that is set for inplace, Please use "
              "`SetInplaceMap` in your output when registry custom operator."));
      const auto &input = inplace_reverse_map.at(output);
      auto index = vec_input_name2id_map[input];
      auto &vec_input_shape = vec_input_shapes[index];
      output_name2value_num[output] = vec_input_shape.size();
    } else {
      if (inplace_reverse_map.find(output) != inplace_reverse_map.end()) {
        const auto &input = inplace_reverse_map.at(output);
        auto index = input_name2id_map[input];
        // input_shapes[index] is dim of tensor, if the dim doesn't have
        // element, it must be a optional tensor that is None in custom operator
        output_name2value_num[output] = input_shapes[index].size() == 0 ? 0 : 1;
      } else {
        output_name2value_num[output]++;
      }
    }
    all_values_num += output_name2value_num[output];
  }

  PADDLE_ENFORCE_EQ(
      output_shapes.size(),
      all_values_num,
      common::errors::InvalidArgument(
          "The number of output shapes after running custom operator's "
          "InferShapeFunc is wrong, "
          "expected contains %d Tensors' shape, but actually contains %d "
          "Tensors' shape",
          all_values_num,
          output_shapes.size()));

  PADDLE_ENFORCE_EQ(
      output_dtypes.size(),
      all_values_num,
      common::errors::InvalidArgument(
          "The number of output dtypes after running custom operator's "
          "InferDtypeFunc is wrong, "
          "expected contains %d Tensors' dtype, but actually contains %d "
          "Tensors' dtype",
          all_values_num,
          output_dtypes.size()));
  if (run_auto_parallel) {
    PADDLE_ENFORCE_EQ(
        spmd_info.second.size(),
        all_values_num,
        common::errors::InvalidArgument(
            "The number of output dist_attr after running custom operator's "
            "InferSPMD is wrong, "
            "expected contains %d Tensors' dist_attr, but actually contains %d "
            "Tensors' dist_attr",
            all_values_num,
            spmd_info.second.size()));
  }
  size_t value_index = 0;
  for (size_t i = 0; i < outputs.size(); ++i) {
    const auto &output = outputs.at(i);
    auto value_num = output_name2value_num[output];
    if (value_num == 0) {
      // Optional value condition
      pir::Type out_type;
      argument_outputs.push_back(out_type);
      continue;
    }
    if (paddle::framework::detail::IsDuplicableVar(output)) {
      std::vector<pir::Type> out_types;
      std::vector<pir::Attribute> dist_attrs;
      for (size_t j = 0; j < value_num; ++j) {
        auto ddims = phi::make_ddim(output_shapes[value_index]);
        auto dtype = output_dtypes[value_index];
        phi::DataLayout layout{DataLayout::NCHW};
        phi::LegacyLoD lod;
        auto type = paddle::dialect::DenseTensorType::get(
            pir::IrContext::Instance(),
            paddle::dialect::TransToIrDataType(dtype),
            ddims,
            layout,
            lod,
            0);
        if (run_auto_parallel) {
          auto dist_attr = dialect::CvtToPirAttr(spmd_info.second[value_index]);
          out_types.push_back(dialect::CvtToPirDistType(type, dist_attr));
          dist_attrs.push_back(dist_attr);
        } else {
          out_types.push_back(std::move(type));
        }
        value_index++;
      }
      pir::Type out_vector_type =
          pir::VectorType::get(pir::IrContext::Instance(), out_types);
      argument_outputs.push_back(out_vector_type);
      if (run_auto_parallel) {
        dist_result_attrs.push_back(
            pir::ArrayAttribute::get(pir::IrContext::Instance(), dist_attrs));
      }
    } else {
      auto ddims = phi::make_ddim(output_shapes[value_index]);
      auto dtype = output_dtypes[value_index];
      phi::DataLayout layout{DataLayout::NCHW};
      phi::LegacyLoD lod;
      auto out_type = paddle::dialect::DenseTensorType::get(
          pir::IrContext::Instance(),
          paddle::dialect::TransToIrDataType(dtype),
          ddims,
          layout,
          lod,
          0);
      if (run_auto_parallel) {
        auto dist_attr = dialect::CvtToPirAttr(spmd_info.second[value_index]);
        argument_outputs.push_back(
            dialect::CvtToPirDistType(out_type, dist_attr));
        dist_result_attrs.push_back(dist_attr);
      } else {
        argument_outputs.push_back(out_type);
      }
      value_index++;
    }
  }

  // construct operator_dist_attr
  if (run_auto_parallel) {
    std::vector<pir::Attribute> dist_operand_attrs;
    for (auto &arg_dist : spmd_info.first) {
      dist_operand_attrs.push_back(dialect::CvtToPirAttr(arg_dist));
    }
    auto op_dist_attr = dialect::OperationDistAttribute::get(
        ctx, op_mesh, dist_operand_attrs, dist_result_attrs);
    std::ostringstream print_stream;
    print_stream << op_dist_attr;
    VLOG(7) << "Custom Op: " << op_type << " InferSPMD Operator dist attr"
            << print_stream.str();
    argument.AddAttribute(
        kAttrOpDistAttr,
        dialect::OperationDistAttribute::get(
            ctx, op_mesh, dist_operand_attrs, dist_result_attrs));
  }

  argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
  pir::PassStopGradientsDefaultly(argument);
  CallStackRecorder callstack_recorder("run_custom_op");
  callstack_recorder.Record();
  std::vector<pir::Value> op_results;
  pir::Operation *op =
      paddle::dialect::ApiBuilder::Instance().GetBuilder()->Build(
          std::move(argument));
  for (size_t i = 0; i < outputs.size(); ++i) {
    const auto &output = outputs.at(i);
    if (paddle::framework::detail::IsDuplicableVar(output)) {
      if (op->result(i).type().dyn_cast<pir::VectorType>()) {
        auto split_op = paddle::dialect::ApiBuilder::Instance()
                            .GetBuilder()
                            ->Build<pir::SplitOp>(op->result(i));
        auto split_outputs = split_op.outputs();
        op_results.insert(
            op_results.end(), split_outputs.begin(), split_outputs.end());
      }
    } else {
      op_results.push_back(op->result(i));
    }
  }
  callstack_recorder.AttachToOps();
  return ToPyObject(op_results);
}

static PyObject *run_custom_op(PyObject *self,
                               PyObject *args,
                               PyObject *kwargs) {
  if (egr::Controller::Instance().GetCurrentTracer() == nullptr) {
    VLOG(6) << "Call static_api_abs";
    return static_api_run_custom_op(self, args, kwargs);
  } else {
    VLOG(6) << "Call eager_api_abs";
    return eager_api_run_custom_op(self, args, kwargs);
  }
}

using IrTensor = paddle::dialect::IrTensor;

template <typename T>
auto CreatePyFuncRunner(void *py_func_ptr, const std::string &op_name) {
  static_assert(
      std::is_same_v<T, Tensor> || std::is_same_v<T, phi::NativeMetaTensor>,
      "T must be either Tensor or phi::NativeMetaTensor");

  using FuncInputType =
      std::conditional_t<std::is_same_v<T, phi::NativeMetaTensor>,
                         const std::vector<phi::NativeMetaTensor>,
                         std::vector<Tensor>>;

  using FuncOutputType = std::vector<T>;

  return [=](FuncInputType &inputs) -> FuncOutputType {
    py::gil_scoped_acquire acquire;
    PyObject *py_func = reinterpret_cast<PyObject *>(py_func_ptr);

    py::tuple py_args(inputs.size());
    size_t index = 0;
    for (auto &tensor : inputs) {
      py_args[index++] = py::cast(tensor);
    }
    Py_INCREF(py_func);
    PyObject *raw_result = PyObject_CallObject(py_func, py_args.ptr());
    Py_DECREF(py_func);

    if (raw_result == nullptr) {
      PyErr_Print();
      PADDLE_THROW(
          common::errors::Fatal("Execution of the Python OP (%s) failed.\n"
                                "Please review your code, and you may use "
                                "breakpoint() for debugging.",
                                op_name));
    }

    py::object result = py::reinterpret_steal<py::object>(raw_result);
    std::vector<T> outputs;

    if (py::isinstance<py::tuple>(result)) {
      py::tuple tuple_result = py::cast<py::tuple>(result);
      for (const auto &item : tuple_result) {
        outputs.push_back(py::cast<T>(item));
      }
    } else {
      outputs.push_back(py::cast<T>(result));
    }
    return outputs;
  };
}

static PyObject *run_python_op(PyObject *self,
                               PyObject *args,
                               PyObject *kwargs) {
  VLOG(6) << "Call run_python_op";

  if (kwargs == NULL) {
    PyErr_SetString(
        PyExc_TypeError,
        "kwargs cannot be NULL. Please add inputs/outputs/attr/inplace_map!");
    return NULL;
  }

  PyObject *py_op_name = PyDict_GetItemString(kwargs, "name");
  PyObject *py_input_names = PyDict_GetItemString(kwargs, "input_names");
  PyObject *py_output_names = PyDict_GetItemString(kwargs, "output_names");
  PyObject *py_attrs_dict = PyDict_GetItemString(kwargs, "attrs");
  PyObject *py_inplace_dict = PyDict_GetItemString(kwargs, "inplace_map");

  if (!py_op_name || !py_input_names || !py_output_names || !py_attrs_dict ||
      !py_inplace_dict) {
    PyErr_SetString(
        PyExc_KeyError,
        "Required key (inputs/outputs/attr/inplace_map) missing from kwargs.");
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }

  std::string op_name = CastPyArg2String(py_op_name, "run_python_op", 0);
  std::vector<std::string> inputs_vec =
      CastPyArg2Strings(py_input_names, "run_python_op", 0);
  std::vector<std::string> outputs_vec =
      CastPyArg2Strings(py_output_names, "run_python_op", 0);
  std::unordered_map<std::string, void *> attrs_map =
      ParsePythonOpAttrs(py_attrs_dict);
  std::unordered_map<std::string, std::string> op_inplace_map =
      ParseStringDict(py_inplace_dict);

  VLOG(6) << "Building Python OP [" << op_name << "] with attrs:" << std::endl
          << "    op_name: " << op_name << std::endl
          << "    inputs: " << paddle::string::join_strings(inputs_vec, ", ")
          << std::endl
          << "    outputs: " << paddle::string::join_strings(outputs_vec, ", ")
          << std::endl
          << "    attrs[infer_meta_fn_ptr]: "
          << reinterpret_cast<uintptr_t>(attrs_map["infer_meta_fn_ptr"])
          << std::endl
          << "    attrs[fn_ptr]: "
          << reinterpret_cast<uintptr_t>(attrs_map["fn_ptr"]);

  const auto &meta_info_map = OpMetaInfoMap::Instance().GetMap();

  auto py_func = CreatePyFuncRunner<Tensor>(attrs_map["fn_ptr"], op_name);
  auto infer_meta_py_func = CreatePyFuncRunner<phi::NativeMetaTensor>(
      attrs_map["infer_meta_fn_ptr"], op_name);

  if (meta_info_map.find(op_name) == meta_info_map.end()) {
    VLOG(6) << "Python OP " << op_name << " does not exist, registering...";
    paddle::framework::RegisterPythonOperator(
        op_name,
        std::move(inputs_vec),
        std::move(outputs_vec),
        {"infer_meta_fn_ptr: void*", "fn_ptr: void*"},
        std::move(op_inplace_map),
        std::move(py_func),
        std::move(infer_meta_py_func));
  }

  PADDLE_ENFORCE_NE(meta_info_map.find(op_name),
                    meta_info_map.end(),
                    common::errors::NotFound(
                        "Can't find %s in Eager OpMetaInfoMap which should be "
                        "created by LoadOpMetaInfoAndRegisterOp, please make "
                        "sure you registered your op first and try again. ",
                        op_name));

  const auto &vec_map = meta_info_map.at(op_name);
  const auto &inputs = paddle::OpMetaInfoHelper::GetInputs(vec_map[0]);
  const auto &outputs = paddle::OpMetaInfoHelper::GetOutputs(vec_map[0]);
  const auto &inplace_map = paddle::OpMetaInfoHelper::GetInplaceMap(vec_map[0]);
  const auto &inplace_reverse_map =
      paddle::OpMetaInfoHelper::GetInplaceReverseMap(vec_map[0]);

  std::string pir_op_name =
      paddle::framework::kPythonOperatorDialectPrefix + op_name;
  if (!inplace_map.empty()) {
    pir_op_name += "_";
  }
  pir::IrContext *ctx = pir::IrContext::Instance();
  pir::OpInfo pir_info = ctx->GetRegisteredOpInfo(pir_op_name);
  pir::OperationArgument argument(pir_info);
  std::vector<pir::Value> argument_inputs;
  std::vector<pir::Type> argument_outputs;

  std::vector<std::vector<int64_t>> input_shapes;
  std::vector<DataType> input_dtypes;
  std::unordered_map<std::string, int> input_name2id_map;
  std::vector<std::vector<std::vector<int64_t>>> vec_input_shapes;
  std::vector<std::vector<DataType>> vec_input_dtypes;
  std::unordered_map<std::string, int> vec_input_name2id_map;
  std::vector<paddle::any> custom_attrs;
  int input_index = 0;

  std::vector<phi::NativeMetaTensor> inputs_meta;
  inputs_meta.reserve(inputs.size());

  for (size_t i = 0; i < inputs.size(); ++i) {
    const auto &input = inputs.at(i);
    PyObject *obj = PyTuple_GET_ITEM(args, i);
    // Emplace Py_None from python, this means optional inputs passed to C++,
    // use one un-initialized tensor to indicate both Tensor and
    // vector<Tensor> inputs.
    if (obj == Py_None) {
      PADDLE_THROW(common::errors::Unimplemented(
          "Currently, optional Tensor input is not supported in "
          "Python operator."));
    }
    if (paddle::framework::detail::IsDuplicableVar(input)) {
      PADDLE_THROW(common::errors::Unimplemented(
          "Currently, optional vector<Tensor> input is not supported in "
          "Python operator."));
    } else {
      input_name2id_map[inputs[i]] = input_index;
      input_index++;
      pir::Value input_value =
          CastPyArg2Value(obj, op_name, i, false);  // NOLINT
      paddle::dialect::DenseTensorType input_tensor =
          input_value.type().dyn_cast<paddle::dialect::DenseTensorType>();
      argument_inputs.push_back(input_value);

      inputs_meta.push_back(phi::NativeMetaTensor(
          paddle::dialect::TransToPhiDataType(input_tensor.dtype()),
          input_tensor.dims()));
    }
  }
  argument.AddInputs(argument_inputs);

  custom_attrs.push_back(attrs_map["infer_meta_fn_ptr"]);
  custom_attrs.push_back(attrs_map["fn_ptr"]);
  argument.AddAttribute(
      "infer_meta_fn_ptr",
      pir::PointerAttribute::get(pir::IrContext::Instance(),
                                 attrs_map["infer_meta_fn_ptr"]));
  argument.AddAttribute("fn_ptr",
                        pir::PointerAttribute::get(pir::IrContext::Instance(),
                                                   attrs_map["fn_ptr"]));

  // Run infer meta
  VLOG(4) << "Start to run infer meta for " << op_name;
  std::vector<phi::NativeMetaTensor> outputs_meta =
      infer_meta_py_func(inputs_meta);
  VLOG(4) << "End to run infer meta for " << op_name;
  std::vector<IrTensor> process_result;
  process_result.reserve(outputs.size());
  for (auto &out_meta : outputs_meta) {
    process_result.push_back(
        IrTensor(out_meta.dtype(), out_meta.dims(), phi::DataLayout::NCHW, {}));
  }
  PADDLE_ENFORCE_EQ(
      process_result.size(),
      outputs.size(),
      common::errors::InvalidArgument("Expected output size %d, but got %d.",
                                      static_cast<int>(process_result.size()),
                                      static_cast<int>(outputs.size())));

  dialect::ProcessMeshAttribute op_mesh;
  bool run_auto_parallel = false;
  std::vector<pir::Attribute> dist_result_attrs;
  phi::distributed::SpmdInfo spmd_info;
  if (dialect::HasDistInput(argument_inputs, &op_mesh)) {
    VLOG(7) << "Custom Op: " << op_name << " InferSPMD";
    run_auto_parallel = true;
    spmd_info = paddle::framework::RunInferSpmd(
        vec_map[0], op_name, op_mesh, argument_inputs, custom_attrs);
  }

  size_t all_values_num = 0;
  // output name -> value num (that output should hold)
  std::unordered_map<std::string, size_t> output_name2value_num;
  for (size_t i = 0; i < outputs.size(); ++i) {
    const auto &output = outputs.at(i);
    if (paddle::framework::detail::IsDuplicableVar(output)) {
      PADDLE_ENFORCE_NE(
          inplace_reverse_map.find(output),
          inplace_reverse_map.end(),
          common::errors::InvalidArgument(
              "Only support vector output that is set for inplace, Please use "
              "`SetInplaceMap` in your output when registry custom operator."));
      const auto &input = inplace_reverse_map.at(output);
      auto index = vec_input_name2id_map[input];
      auto &vec_input_shape = vec_input_shapes[index];
      output_name2value_num[output] = vec_input_shape.size();
    } else {
      if (inplace_reverse_map.find(output) != inplace_reverse_map.end()) {
        const auto &input = inplace_reverse_map.at(output);
        auto index = input_name2id_map[input];
        // input_shapes[index] is dim of tensor, if the dim doesn't have
        // element, it must be a optional tensor that is None in custom operator
        output_name2value_num[output] = input_shapes[index].size() == 0 ? 0 : 1;
      } else {
        ++(output_name2value_num[output]);
      }
    }
    all_values_num += output_name2value_num[output];
  }

  if (run_auto_parallel) {
    PADDLE_ENFORCE_EQ(
        spmd_info.second.size(),
        all_values_num,
        common::errors::InvalidArgument(
            "The number of output dist_attr after running custom operator's "
            "InferSPMD is wrong, "
            "expected contains %d Tensors' dist_attr, but actually contains %d "
            "Tensors' dist_attr",
            all_values_num,
            spmd_info.second.size()));
  }

  size_t value_index = 0;
  for (size_t i = 0; i < outputs.size(); ++i) {
    const auto &output = outputs.at(i);
    auto value_num = output_name2value_num[output];
    if (value_num == 0) {
      // Optional value condition
      pir::Type out_type;
      argument_outputs.push_back(out_type);
      continue;
    }
    if (paddle::framework::detail::IsDuplicableVar(output)) {
      PADDLE_THROW(common::errors::Unimplemented(
          "Currently, vector<Tensor> output is not supported in Python "
          "operator."));
    } else {
      auto dense_out = process_result[value_index];
      auto out_type = paddle::dialect::DenseTensorType::get(
          pir::IrContext::Instance(),
          paddle::dialect::TransToIrDataType(dense_out.dtype()),
          dense_out.dims(),
          dense_out.layout(),
          dense_out.lod(),
          dense_out.offset());
      if (run_auto_parallel) {
        auto dist_attr = dialect::CvtToPirAttr(spmd_info.second[value_index]);
        argument_outputs.push_back(
            dialect::CvtToPirDistType(out_type, dist_attr));
        dist_result_attrs.push_back(dist_attr);
      } else {
        argument_outputs.push_back(out_type);
      }
      value_index++;
    }
  }

  // construct operator_dist_attr
  if (run_auto_parallel) {
    std::vector<pir::Attribute> dist_operand_attrs;
    for (auto &arg_dist : spmd_info.first) {
      dist_operand_attrs.push_back(dialect::CvtToPirAttr(arg_dist));
    }
    auto op_dist_attr = dialect::OperationDistAttribute::get(
        ctx, op_mesh, dist_operand_attrs, dist_result_attrs);
    std::ostringstream print_stream;
    print_stream << op_dist_attr;
    VLOG(7) << "Custom Op: " << op_name << " InferSPMD Operator dist attr"
            << print_stream.str();
    argument.AddAttribute(
        kAttrOpDistAttr,
        dialect::OperationDistAttribute::get(
            ctx, op_mesh, dist_operand_attrs, dist_result_attrs));
  }

  argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
  pir::PassStopGradientsDefaultly(argument);
  CallStackRecorder callstack_recorder("run_python_op");
  callstack_recorder.Record();
  std::vector<pir::Value> op_results;
  pir::Operation *op =
      paddle::dialect::ApiBuilder::Instance().GetBuilder()->Build(
          std::move(argument));
  for (size_t i = 0; i < outputs.size(); ++i) {
    const auto &output = outputs.at(i);
    if (paddle::framework::detail::IsDuplicableVar(output)) {
      if (op->result(i).type().dyn_cast<pir::VectorType>()) {
        auto split_op = paddle::dialect::ApiBuilder::Instance()
                            .GetBuilder()
                            ->Build<pir::SplitOp>(op->result(i));
        auto split_outputs = split_op.outputs();
        op_results.insert(
            op_results.end(), split_outputs.begin(), split_outputs.end());
      }
    } else {
      op_results.push_back(op->result(i));
    }
  }
  callstack_recorder.AttachToOps();
  return ToPyObject(op_results);
}

static PyObject *builtin_combine_op(PyObject *self,
                                    PyObject *args,
                                    PyObject *kwargs) {
  try {
    VLOG(6) << "Add builtin_combine op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);
    // Get Value from args
    PyObject *x_obj = PyTuple_GET_ITEM(args, 0);
    auto x = CastPyArg2VectorOfValue(x_obj, "builtin_combine", 0, false);
    CallStackRecorder callstack_recorder("builtin_combine_op");
    callstack_recorder.Record();
    auto static_api_out = paddle::dialect::builtin_combine(x);
    callstack_recorder.AttachToOps();
    return ToPyObject(static_api_out);
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

static PyObject *builtin_split_op(PyObject *self,
                                  PyObject *args,
                                  PyObject *kwargs) {
  try {
    VLOG(6) << "Add builtin_split op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);
    // Get Value from args
    PyObject *x_obj = PyTuple_GET_ITEM(args, 0);
    auto x = CastPyArg2Value(x_obj, "builtin_split", 0, false);
    CallStackRecorder callstack_recorder("builtin_builtin_split");
    callstack_recorder.Record();
    auto static_api_out = paddle::dialect::builtin_split(x);
    callstack_recorder.AttachToOps();
    return ToPyObject(static_api_out);
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

static PyObject *static_api_fused_gemm_epilogue(PyObject *self,
                                                PyObject *args,
                                                PyObject *kwargs) {
  try {
    VLOG(6) << "Running Static API: fused_gemm_epilogue";

    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);
    // Get OpResult from args
    PyObject *x_obj = PyTuple_GET_ITEM(args, 0);
    auto x = CastPyArg2Value(x_obj, "fused_gemm_epilogue", 0, false);
    PyObject *y_obj = PyTuple_GET_ITEM(args, 1);
    auto y = CastPyArg2Value(y_obj, "fused_gemm_epilogue", 1, false);
    PyObject *bias_obj = PyTuple_GET_ITEM(args, 2);
    auto bias = CastPyArg2Value(bias_obj, "fused_gemm_epilogue", 2, false);

    // Parse Attributes if needed
    PyObject *trans_x_obj = PyTuple_GET_ITEM(args, 3);
    bool trans_x = CastPyArg2Boolean(trans_x_obj, "fused_gemm_epilogue", 3);
    PyObject *trans_y_obj = PyTuple_GET_ITEM(args, 4);
    bool trans_y = CastPyArg2Boolean(trans_y_obj, "fused_gemm_epilogue", 4);
    PyObject *activation_obj = PyTuple_GET_ITEM(args, 5);
    std::string activation =
        CastPyArg2String(activation_obj, "fused_gemm_epilogue", 5);
    // Call ir static api
    CallStackRecorder callstack_recorder("fused_gemm_epilogue");
    callstack_recorder.Record();
    auto out = paddle::dialect::fused_gemm_epilogue(
        x, y, bias, trans_x, trans_y, activation);
    callstack_recorder.AttachToOps();

    return ToPyObject(out);
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}
static PyObject *static_api_array_pop(PyObject *self,
                                      PyObject *args,
                                      PyObject *kwargs) {
  try {
    VLOG(6) << "Add array_pop op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);

    // Get Value from args
    PyObject *input_obj = PyTuple_GET_ITEM(args, 0);
    auto input = CastPyArg2Value(input_obj, "array_pop", 0, false);

    PyObject *index_obj = PyTuple_GET_ITEM(args, 1);
    auto index = CastPyArg2Int(index_obj, "array_pop", 1);

    // Call ir static api
    CallStackRecorder callstack_recorder("array_pop");
    callstack_recorder.Record();
    auto static_api_out = paddle::dialect::array_pop(input, index);
    callstack_recorder.AttachToOps();
    return ToPyObject(static_api_out);
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

extern PyTypeObject *g_tensorrt_engine_params_pytype;

static PyObject *static_api_tensorrt_engine(PyObject *self,
                                            PyObject *args,
                                            PyObject *kwargs) {
  try {
    VLOG(6) << "Add tensorrt_engine op into program";

    // Get Value from args
    PyObject *x_obj = PyTuple_GET_ITEM(args, 0);
    auto x = CastPyArg2VectorOfValue(x_obj, "tensorrt_engine", 0, true);

    PyObject *param_obj = PyTuple_GET_ITEM(args, 1);
    if (!PyObject_TypeCheck(param_obj, g_tensorrt_engine_params_pytype)) {
      PADDLE_THROW(common::errors::InvalidType(
          "tensorrt_engine(): argument (position %d) must be "
          "EngineParams, but got %s",
          2,
          ((PyTypeObject *)param_obj->ob_type)->tp_name));  // NOLINT
    }
    auto trt_param =
        ::pybind11::handle(param_obj).cast<paddle::platform::EngineParams>();

    PyObject *input_names_obj = PyTuple_GET_ITEM(args, 2);
    auto input_names = CastPyArg2VectorOfString(input_names_obj, 2);

    PyObject *output_names_obj = PyTuple_GET_ITEM(args, 3);
    auto output_names = CastPyArg2VectorOfString(output_names_obj, 3);

    PyObject *outputs_shape_obj = PyTuple_GET_ITEM(args, 4);
    std::vector<std::vector<int64_t>> outputs_shape;
    if (PyList_Check(outputs_shape_obj)) {
      Py_ssize_t len = PyList_Size(outputs_shape_obj);
      PyObject *item = nullptr;
      for (Py_ssize_t i = 0; i < len; i++) {
        item = PyList_GetItem(outputs_shape_obj, i);
        outputs_shape.emplace_back(CastPyArg2VectorOfInt64(item, 4));
      }
    } else {
      PADDLE_THROW(common::errors::InvalidType(
          "argument (position %d) must be "
          "list but got %s",
          5,
          reinterpret_cast<PyTypeObject *>(outputs_shape_obj->ob_type)
              ->tp_name));
    }

    PyObject *outputs_dtype_obj = PyTuple_GET_ITEM(args, 5);
    std::vector<DataType> outputs_dtype;
    if (PyList_Check(outputs_dtype_obj)) {
      Py_ssize_t len = PyList_Size(outputs_dtype_obj);
      PyObject *item = nullptr;
      for (Py_ssize_t i = 0; i < len; i++) {
        item = PyList_GetItem(outputs_dtype_obj, i);
        outputs_dtype.emplace_back(
            CastPyArg2DataTypeDirectly(item, "tensorrt_engine", 5));
      }
    } else {
      PADDLE_THROW(common::errors::InvalidType(
          "argument (position %d) must be "
          "list but got %s",
          6,
          reinterpret_cast<PyTypeObject *>(outputs_dtype_obj->ob_type)
              ->tp_name));
    }
    PyObject *converter_debug_info_obj = PyTuple_GET_ITEM(args, 6);
    std::string converter_debug_info =
        CastPyArg2String(converter_debug_info_obj, "converter_debug_info", 6);
    // Call ir static api
    CallStackRecorder callstack_recorder("tensorrt_engine");
    callstack_recorder.Record();
    auto static_api_out =
        paddle::dialect::tensorrt_engine(x,
                                         trt_param,
                                         input_names,
                                         output_names,
                                         outputs_shape,
                                         outputs_dtype,
                                         converter_debug_info);
    callstack_recorder.AttachToOps();
    return ToPyObject(static_api_out);
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

extern PyObject *eager_api_fused_gemm_epilogue(PyObject *self,
                                               PyObject *args,
                                               PyObject *kwargs);

static PyObject *fused_gemm_epilogue(PyObject *self,
                                     PyObject *args,
                                     PyObject *kwargs) {
  if (egr::Controller::Instance().GetCurrentTracer() == nullptr) {
    VLOG(6) << "Call static_api_fused_gemm_epilogue";
    return static_api_fused_gemm_epilogue(self, args, kwargs);
  } else {
    VLOG(6) << "Call eager_api_fused_gemm_epilogue";
    return eager_api_fused_gemm_epilogue(self, args, kwargs);
  }
}

static PyObject *anchor_generator(PyObject *self,
                                  PyObject *args,
                                  PyObject *kwargs) {
  if (egr::Controller::Instance().GetCurrentTracer() == nullptr) {
    VLOG(6) << "Call static_api_anchor_generator";
    return static_api_anchor_generator(self, args, kwargs);
  } else {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

static PyObject *share_var(PyObject *self, PyObject *args, PyObject *kwargs) {
  try {
    VLOG(6) << "Add share_var op into program";
    VLOG(8) << "args count: " << (PyTuple_Size(args) / 2);
    // Get Value from args
    PyObject *input_obj = PyTuple_GET_ITEM(args, 0);
    auto inputs = CastPyArg2VectorOfValue(input_obj, "share_var", 0, false);
    CallStackRecorder callstack_recorder("share_var_op");
    callstack_recorder.Record();
    auto share_var_op = paddle::dialect::share_var(inputs);
    callstack_recorder.AttachToOps();
    return ToPyObject(share_var_op);
  } catch (...) {
    ThrowExceptionToPython(std::current_exception());
    return nullptr;
  }
}

static PyMethodDef ManualOpsAPI[] = {
    {"set_parameter",
     (PyCFunction)(void (*)(void))static_api_set_parameter,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for set_parameter."},
    {"update_parameter",
     (PyCFunction)(void (*)(void))static_api_update_parameter,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for update_parameter."},
    {"set_persistable_value",
     (PyCFunction)(void (*)(void))static_api_set_persistable_value,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for set_persistable_value."},
    {"parameter",
     (PyCFunction)(void (*)(void))static_api_parameter,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for parameter."},
    {"create_array",
     (PyCFunction)(void (*)(void))static_api_create_array,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for create_array."},
    {"create_array_like",
     (PyCFunction)(void (*)(void))static_api_create_array_like,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for create_array_like."},
    {"array_length",
     (PyCFunction)(void (*)(void))static_api_array_length,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for array_length."},
    {"array_read",
     (PyCFunction)(void (*)(void))static_api_array_read,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for array_read."},
    {"fetch",
     (PyCFunction)(void (*)(void))static_api_fetch,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for fetch."},
    {"array_write_",
     (PyCFunction)(void (*)(void))static_api_array_write_,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for array_write_."},
    {"array_to_tensor",
     (PyCFunction)(void (*)(void))static_api_array_to_tensor,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for array_to_tensor."},
    {"add_n_array",
     (PyCFunction)(void (*)(void))static_api_add_n_array,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for add_n_array."},
    {"slice_array",
     (PyCFunction)(void (*)(void))static_api_slice_array,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for slice_array."},
    {"slice_array_dense",
     (PyCFunction)(void (*)(void))static_api_slice_array_dense,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for slice_array_dense."},
    {"fused_gemm_epilogue",
     (PyCFunction)(void (*)(void))fused_gemm_epilogue,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for fused_gemm_epilogue."},
    {"anchor_generator",
     (PyCFunction)(void (*)(void))anchor_generator,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for anchor_generator."},
    {"_run_custom_op",
     (PyCFunction)(void (*)(void))run_custom_op,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for run_custom_op."},
    {"_run_python_op",
     (PyCFunction)(void (*)(void))run_python_op,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for run_python_op."},
    {"builtin_combine",
     (PyCFunction)(void (*)(void))builtin_combine_op,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for builtin_combine_op."},
    {"builtin_split",
     (PyCFunction)(void (*)(void))builtin_split_op,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for builtin_split_op."},
    {"tensorrt_engine",
     (PyCFunction)(void (*)(void))static_api_tensorrt_engine,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for tensorrt_engine."},
    {"array_pop",
     (PyCFunction)(void (*)(void))static_api_array_pop,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for array_pop."},
    {"share_var",
     (PyCFunction)(void (*)(void))share_var,
     METH_VARARGS | METH_KEYWORDS,
     "C++ interface function for share_var_op."},
    {nullptr, nullptr, 0, nullptr}};

}  // namespace pybind
}  // namespace paddle