wehub/paddlepaddle--paddle
1
0
Fork 0
Code Issues Pull Requests Actions 2 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
ff97fb6869c2eeecb7b4c09af41ce69f43f2c574
paddlepaddle--paddle/paddle/fluid/pybind/manual_static_op_function.h
T
wehub-resource-sync e25996e7db chore: import upstream snapshot with attribution
2026-07-13 12:40:42 +08:00

1673 lines
65 KiB
C++
Raw Blame History

// 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
Reference in New Issue View Git Blame Copy Permalink