449 lines
21 KiB
C++
449 lines
21 KiB
C++
// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#include "paddle/fluid/inference/api/helper.h"
|
|
#include <cstdint>
|
|
|
|
#include "paddle/common/enforce.h"
|
|
#include "paddle/common/errors.h"
|
|
#include "paddle/common/flags.h"
|
|
#include "paddle/fluid/framework/custom_operator.h"
|
|
#include "paddle/fluid/framework/custom_operator_utils.h"
|
|
#include "paddle/fluid/framework/operator.h"
|
|
#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
|
|
#include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
|
|
#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
|
|
#include "paddle/fluid/pir/drr/src/ir_operation_factory.h"
|
|
#include "paddle/fluid/platform/init.h"
|
|
#include "paddle/phi/api/ext/op_meta_info.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
#include "paddle/pir/include/core/builtin_attribute.h"
|
|
#include "paddle/pir/include/core/builtin_type.h"
|
|
#include "paddle/pir/include/core/ir_context.h"
|
|
#include "paddle/pir/include/core/operation.h"
|
|
#include "paddle/pir/include/core/value.h"
|
|
|
|
namespace paddle::inference {
|
|
|
|
template <>
|
|
std::string to_string<std::vector<float>>(
|
|
const std::vector<std::vector<float>> &vec) {
|
|
std::stringstream ss;
|
|
for (const auto &piece : vec) {
|
|
ss << to_string(piece) << "\n";
|
|
}
|
|
return ss.str();
|
|
}
|
|
|
|
template <>
|
|
std::string to_string<std::vector<std::vector<float>>>(
|
|
const std::vector<std::vector<std::vector<float>>> &vec) {
|
|
std::stringstream ss;
|
|
for (const auto &line : vec) {
|
|
for (const auto &rcd : line) {
|
|
ss << to_string(rcd) << ";\t";
|
|
}
|
|
ss << '\n';
|
|
}
|
|
return ss.str();
|
|
}
|
|
|
|
void RegisterAllCustomOperator(bool use_pir) {
|
|
const auto &meta_info_map = OpMetaInfoMap::Instance().GetMap();
|
|
for (auto &pair : meta_info_map) {
|
|
if (use_pir) {
|
|
auto *custom_dialect =
|
|
::pir::IrContext::Instance()
|
|
->GetOrRegisterDialect<paddle::dialect::CustomOpDialect>();
|
|
if (custom_dialect->HasRegistered(pair.first)) {
|
|
VLOG(3) << "The operator `" << pair.first
|
|
<< "` has been registered. "
|
|
"Therefore, we will not repeat the registration here.";
|
|
continue;
|
|
}
|
|
for (const auto &meta_info : pair.second) {
|
|
VLOG(3) << "register pir custom op: " << pair.first;
|
|
custom_dialect->RegisterCustomOp(meta_info);
|
|
}
|
|
|
|
std::string pir_op_name =
|
|
paddle::framework::kCustomDialectPrefix + pair.first;
|
|
paddle::drr::OperationFactory::Instance().RegisterOperationCreator(
|
|
pir_op_name,
|
|
[pair, pir_op_name](
|
|
const std::vector<::pir::Value> &inputs,
|
|
const ::pir::AttributeMap &attrs,
|
|
::pir::PatternRewriter &rewriter) mutable -> ::pir::Operation * {
|
|
const auto &meta_inputs =
|
|
paddle::OpMetaInfoHelper::GetInputs(pair.second[0]);
|
|
const auto &meta_attrs =
|
|
paddle::OpMetaInfoHelper::GetAttrs(pair.second[0]);
|
|
const auto &meta_outputs =
|
|
paddle::OpMetaInfoHelper::GetOutputs(pair.second[0]);
|
|
const auto &inplace_map =
|
|
paddle::OpMetaInfoHelper::GetInplaceMap(pair.second[0]);
|
|
const auto &inplace_reverse_map =
|
|
paddle::OpMetaInfoHelper::GetInplaceReverseMap(pair.second[0]);
|
|
auto infershape_func =
|
|
OpMetaInfoHelper::GetInferShapeFn(pair.second[0]);
|
|
auto inferdtype_func =
|
|
OpMetaInfoHelper::GetInferDtypeFn(pair.second[0]);
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
meta_inputs.size(),
|
|
inputs.size(),
|
|
common::errors::InvalidArgument(
|
|
"The number of inputs for the custom operator [%s] given "
|
|
"in the Pattern needs to be consistent with the number at "
|
|
"implementation time.",
|
|
pir_op_name));
|
|
PADDLE_ENFORCE_EQ(
|
|
meta_attrs.size(),
|
|
attrs.size(),
|
|
common::errors::InvalidArgument(
|
|
"The number of attrs for the custom operator [%s] given "
|
|
"in the Pattern needs to be consistent with the number at "
|
|
"implementation time.",
|
|
pir_op_name));
|
|
|
|
if (!inplace_map.empty()) {
|
|
pir_op_name += "_";
|
|
}
|
|
::pir::OperationArgument argument(
|
|
rewriter.ir_context()->GetRegisteredOpInfo(pir_op_name));
|
|
argument.attributes = attrs;
|
|
argument.inputs = 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 < meta_inputs.size(); ++i) {
|
|
const auto &meta_input = meta_inputs.at(i);
|
|
if (!inputs[i]) {
|
|
VLOG(6) << "Add un-initialized tensor because the optional "
|
|
"input is None.";
|
|
if (paddle::framework::detail::IsDuplicableVar(meta_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[meta_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[meta_inputs[i]] = input_index;
|
|
input_index++;
|
|
}
|
|
continue;
|
|
}
|
|
if (paddle::framework::detail::IsDuplicableVar(meta_input)) {
|
|
PADDLE_ENFORCE_EQ(
|
|
inputs[i].type().isa<::pir::VectorType>(),
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"The [%d] input of the custom operator [%s] "
|
|
"should be a pir::VectorType.",
|
|
i,
|
|
pir_op_name));
|
|
std::vector<std::vector<int64_t>> tmp_input_shapes;
|
|
std::vector<phi::DataType> tmp_input_dtypes;
|
|
vec_input_name2id_map[meta_inputs[i]] = vec_input_index;
|
|
vec_input_index++;
|
|
auto input_value_types =
|
|
inputs[i].type().dyn_cast<::pir::VectorType>().data();
|
|
for (auto &input_value_type : input_value_types) {
|
|
auto 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);
|
|
} else {
|
|
input_name2id_map[meta_inputs[i]] = input_index;
|
|
input_index++;
|
|
auto input_tensor =
|
|
inputs[i]
|
|
.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()));
|
|
}
|
|
}
|
|
|
|
for (const auto &meta_attr : meta_attrs) {
|
|
auto attr_name_and_type = paddle::ParseAttrStr(meta_attr);
|
|
auto attr_name = attr_name_and_type[0];
|
|
auto attr_type = attr_name_and_type[1];
|
|
PADDLE_ENFORCE_EQ(attrs.count(attr_name),
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"The attr [%s] in the custom operator [%s] "
|
|
"specified in the Pattern needs to be "
|
|
"consistent with the implementation",
|
|
attr_name,
|
|
pir_op_name));
|
|
VLOG(6) << "Custom operator add attrs " << attr_name
|
|
<< " to CustomOpKernelContext. Attribute type = "
|
|
<< attr_type;
|
|
if (attr_type == "bool") {
|
|
auto bool_attr =
|
|
attrs.at(attr_name).dyn_cast<::pir::BoolAttribute>().data();
|
|
custom_attrs.emplace_back(bool_attr);
|
|
} else if (attr_type == "int") {
|
|
int int_attr = attrs.at(attr_name)
|
|
.dyn_cast<::pir::Int32Attribute>()
|
|
.data();
|
|
custom_attrs.emplace_back(int_attr);
|
|
} else if (attr_type == "float") {
|
|
float float_attr = attrs.at(attr_name)
|
|
.dyn_cast<::pir::FloatAttribute>()
|
|
.data();
|
|
custom_attrs.emplace_back(float_attr);
|
|
} else if (attr_type == "int64_t") {
|
|
int64_t long_attr = attrs.at(attr_name)
|
|
.dyn_cast<::pir::Int64Attribute>()
|
|
.data();
|
|
custom_attrs.emplace_back(long_attr);
|
|
} else if (attr_type == "std::string") {
|
|
std::string str_attr = attrs.at(attr_name)
|
|
.dyn_cast<::pir::StrAttribute>()
|
|
.AsString();
|
|
custom_attrs.emplace_back(str_attr);
|
|
} else if (attr_type == "std::vector<int>") {
|
|
auto vec_attr = attrs.at(attr_name)
|
|
.dyn_cast<::pir::ArrayAttribute>()
|
|
.AsVector();
|
|
std::vector<int> vec_int_attr;
|
|
for (const auto &int_attr : vec_attr) {
|
|
vec_int_attr.push_back(
|
|
int_attr.dyn_cast<::pir::Int32Attribute>().data());
|
|
}
|
|
custom_attrs.emplace_back(vec_int_attr);
|
|
} else if (attr_type == "std::vector<float>") {
|
|
auto vec_attr = attrs.at(attr_name)
|
|
.dyn_cast<::pir::ArrayAttribute>()
|
|
.AsVector();
|
|
std::vector<float> vec_float_attr;
|
|
for (const auto &float_attr : vec_attr) {
|
|
vec_float_attr.push_back(
|
|
float_attr.dyn_cast<::pir::FloatAttribute>().data());
|
|
}
|
|
custom_attrs.emplace_back(vec_float_attr);
|
|
} else if (attr_type == "std::vector<int64_t>") {
|
|
auto vec_attr = attrs.at(attr_name)
|
|
.dyn_cast<::pir::ArrayAttribute>()
|
|
.AsVector();
|
|
std::vector<int64_t> vec_long_attr;
|
|
for (const auto &long_attr : vec_attr) {
|
|
vec_long_attr.push_back(
|
|
long_attr.dyn_cast<::pir::Int64Attribute>().data());
|
|
}
|
|
custom_attrs.emplace_back(vec_long_attr);
|
|
} else if (attr_type == "std::vector<std::string>") {
|
|
auto vec_attr = attrs.at(attr_name)
|
|
.dyn_cast<::pir::ArrayAttribute>()
|
|
.AsVector();
|
|
std::vector<std::string> vec_string_attr;
|
|
for (const auto &string_attr : vec_attr) {
|
|
vec_string_attr.push_back(
|
|
string_attr.dyn_cast<::pir::StrAttribute>().AsString());
|
|
}
|
|
custom_attrs.emplace_back(vec_string_attr);
|
|
} else {
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"Unsupported `%s` type value as custom attribute now. "
|
|
"Supported data types include `bool`, `int`, `float`, "
|
|
"`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));
|
|
}
|
|
}
|
|
|
|
paddle::framework::CheckDefaultInferShapeDtype(
|
|
infershape_func, inferdtype_func, pair.second[0]);
|
|
std::vector<std::vector<int64_t>> output_shapes =
|
|
paddle::framework::RunInferShape(infershape_func,
|
|
pair.second[0],
|
|
input_shapes,
|
|
input_name2id_map,
|
|
vec_input_shapes,
|
|
vec_input_name2id_map,
|
|
custom_attrs);
|
|
std::vector<phi::DataType> output_dtypes =
|
|
paddle::framework::RunInferDtype(inferdtype_func,
|
|
pair.second[0],
|
|
input_dtypes,
|
|
input_name2id_map,
|
|
vec_input_dtypes,
|
|
vec_input_name2id_map,
|
|
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 (const auto &output : meta_outputs) {
|
|
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].empty() ? 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()));
|
|
|
|
size_t value_index = 0;
|
|
for (const auto &output : meta_outputs) {
|
|
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)) {
|
|
auto value_num = output_name2value_num[output];
|
|
std::vector<pir::Type> out_types;
|
|
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;
|
|
out_types.push_back(paddle::dialect::DenseTensorType::get(
|
|
pir::IrContext::Instance(),
|
|
paddle::dialect::TransToIrDataType(dtype),
|
|
ddims,
|
|
layout,
|
|
lod,
|
|
0));
|
|
value_index++;
|
|
}
|
|
pir::Type out_vector_type =
|
|
pir::VectorType::get(pir::IrContext::Instance(), out_types);
|
|
argument_outputs.push_back(out_vector_type);
|
|
} 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);
|
|
argument_outputs.push_back(out_type);
|
|
value_index++;
|
|
}
|
|
}
|
|
|
|
argument.AddOutputs(argument_outputs.begin(),
|
|
argument_outputs.end());
|
|
::pir::PassStopGradientsDefaultly(argument);
|
|
return rewriter.Build(std::move(argument));
|
|
});
|
|
}
|
|
const auto &all_op_kernels{framework::OperatorWithKernel::AllOpKernels()};
|
|
if (all_op_kernels.find(pair.first) == all_op_kernels.end()) {
|
|
framework::RegisterOperatorWithMetaInfo(pair.second);
|
|
} else {
|
|
VLOG(3) << "The operator `" << pair.first
|
|
<< "` has been registered. Therefore, we will not repeat the "
|
|
"registration here.";
|
|
}
|
|
}
|
|
}
|
|
|
|
void InitGflagsFromEnv() {
|
|
// support set gflags from environment.
|
|
std::vector<std::string> gflags;
|
|
const phi::ExportedFlagInfoMap &env_map = phi::GetExportedFlagInfoMap();
|
|
std::ostringstream os;
|
|
for (auto &pair : env_map) {
|
|
os << pair.second.name << ",";
|
|
}
|
|
std::string tryfromenv_str = os.str();
|
|
if (!tryfromenv_str.empty()) {
|
|
tryfromenv_str.pop_back();
|
|
tryfromenv_str = "--tryfromenv=" + tryfromenv_str;
|
|
gflags.push_back(tryfromenv_str);
|
|
}
|
|
framework::InitGflags(gflags);
|
|
}
|
|
|
|
} // namespace paddle::inference
|