Skip to content

Commit

Permalink
[PHI CAPI] Add support for registering a new operator, PART1 (PaddleP…
Browse files Browse the repository at this point in the history 
…addle#55532)
  • Loading branch information
@ronny1996 @wz1qqx
ronny1996 authored and wz1qqx committed Jul 31, 2023
1 parent de1cc2d commit df648e5
Show file tree
Hide file tree
Showing 5 changed files with 513 additions and 311 deletions.
286 changes: 0 additions & 286 deletions paddle/fluid/framework/custom_operator.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,9 +28,6 @@ limitations under the License. */
#include "paddle/fluid/eager/api/utils/global_utils.h"
#include "paddle/fluid/framework/attribute.h"
#include "paddle/fluid/framework/convert_utils.h"
#include "paddle/fluid/framework/custom_operator_utils.h"
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/framework/operator.h"
#include "paddle/fluid/framework/phi_utils.h"
#include "paddle/fluid/framework/tensor.h"
#include "paddle/fluid/platform/device/gpu/gpu_info.h"
Expand Down Expand Up @@ -843,289 +840,6 @@ class CustomOperator : public OperatorWithKernel {
}
};

class CustomOpMaker : public OpProtoAndCheckerMaker {
public:
explicit CustomOpMaker(const std::vector<std::string>& inputs,
const std::vector<std::string>& outputs,
const std::vector<std::string>& attrs)
: inputs_(inputs), outputs_(outputs), attrs_(attrs) {}

void Make() override {
for (auto& in_name : inputs_) {
auto input_var_builder =
AddInput(in_name, "The input " + in_name + "of Custom operator.");
if (detail::IsDuplicableVar(in_name)) {
input_var_builder.AsDuplicable();
}
if (detail::IsOptionalVar(in_name)) {
input_var_builder.AsDispensable();
}
}
for (auto& out_name : outputs_) {
auto output_var_builder =
AddOutput(out_name, "The output " + out_name + "of Custom Operator.");
if (detail::IsDuplicableVar(out_name)) {
output_var_builder.AsDuplicable();
}
if (detail::IsOptionalVar(out_name)) {
output_var_builder.AsDispensable();
}
}
for (auto& attr : attrs_) {
auto attr_name_and_type = paddle::ParseAttrStr(attr);
auto attr_name = attr_name_and_type[0];
auto attr_type_str = attr_name_and_type[1];
if (attr_type_str == "bool") {
AddAttr<bool>(attr_name, "custom operator bool attribute.")
.SetDefault(false);
} else if (attr_type_str == "int") {
AddAttr<int>(attr_name, "custom operator int attribute.").SetDefault(1);
} else if (attr_type_str == "float") {
AddAttr<float>(attr_name, "custom operator float attribute.")
.SetDefault(1.0f);
} else if (attr_type_str == "int64_t") {
AddAttr<int64_t>(attr_name, "custom operator int64_t attribute.")
.SetDefault(1);
} else if (attr_type_str == "std::string") {
AddAttr<std::string>(attr_name, "custom operator int attribute.")
.SetDefault("");
} else if (attr_type_str == "std::vector<int>") {
AddAttr<std::vector<int>>(attr_name,
"custom operator std::vector<int> attribute.")
.SetDefault({});
} else if (attr_type_str == "std::vector<float>") {
AddAttr<std::vector<float>>(
attr_name, "custom operator std::vector<float> attribute.")
.SetDefault({});
} else if (attr_type_str == "std::vector<int64_t>") {
AddAttr<std::vector<int64_t>>(
attr_name, "custom operator std::vector<int64_t> attribute.")
.SetDefault({});
} else if (attr_type_str == "std::vector<std::string>") {
AddAttr<std::vector<std::string>>(
attr_name, "custom operator std::vector<std::string> attribute.")
.SetDefault({});
} else {
PADDLE_THROW(platform::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_str));
}
}
AddComment(R"DOC(
Custom Operator.
According to the phi::DenseTensor operation function implemented by the user
independently of the framework, it is encapsulated into a framework
operator to adapt to various execution scenarios such as dynamic graph
mode, static graph mode, and inference mode.
)DOC");
}

private:
std::vector<std::string> inputs_;
std::vector<std::string> outputs_;
std::vector<std::string> attrs_;
};

template <typename T>
class CustomGradOpMaker;

template <>
class CustomGradOpMaker<OpDesc> : public SingleGradOpMaker<OpDesc> {
public:
explicit CustomGradOpMaker(
const OpDesc& fwd_op,
const std::unordered_set<std::string>& no_grad_set,
std::unordered_map<std::string, std::string>* grad_to_var,
const std::vector<BlockDesc*>& grad_block,
const std::string& name,
const std::vector<std::string>& inputs,
const std::vector<std::string>& outputs,
bool is_double_grad)
: SingleGradOpMaker<OpDesc>(fwd_op, no_grad_set, grad_to_var, grad_block),
name_(name),
inputs_(inputs),
outputs_(outputs),
is_double_grad_(is_double_grad) {}

protected:
void Apply(GradOpPtr<OpDesc> grad_op) const override {
grad_op->SetType(name_);

auto fwd_op_inputs = this->InputNames();
auto fwd_op_outputs = this->OutputNames();

for (auto& in_name : inputs_) {
VLOG(3) << "Custom Operator: GradOpDescMaker - input: " << in_name;
if (!detail::IsGradVar(in_name, is_double_grad_)) {
if (detail::IsMemberOf(fwd_op_inputs, in_name)) {
grad_op->SetInput(in_name, this->Input(in_name));
} else if (detail::IsMemberOf(fwd_op_outputs, in_name)) {
grad_op->SetInput(in_name, this->Output(in_name));
} else {
PADDLE_THROW(platform::errors::InvalidArgument(
"The input tensor name `%s` is invalid, expected it is the input "
"or output of forward operator.",
in_name));
}
} else {
if (this->HasOutput(detail::NoGrad(in_name))) {
grad_op->SetInput(in_name, this->OutputGrad(detail::NoGrad(in_name)));
} else {
// Maybe visit here! handle inplace optional case
PADDLE_ENFORCE(
in_name.find(paddle::kOptionalSuffix) != std::string::npos,
phi::errors::InvalidArgument(
"Custom operator couldn't find grad operator input name for "
"%s. If you are using inplace optional inputs & outputs, "
"please check your InplaceMap and `Outputs` again and make "
"sure %s is wrapped by `paddle::Optional`",
in_name,
in_name));
VLOG(3) << "Custom Operator: GradOpDescMaker - handle unfound input: "
<< in_name;
}
}
}
for (auto& out_name : outputs_) {
// Handle inplace optional case
if (!this->HasInput(detail::NoGrad(out_name, is_double_grad_))) {
PADDLE_ENFORCE(
out_name.find(paddle::kOptionalSuffix) != std::string::npos,
phi::errors::InvalidArgument(
"Custom operator couldn't find grad operator output name for "
"%s. If you are using inplace optional inputs & outputs, "
"please check your InplaceMap and `Outputs` again and make "
"sure %s is wrapped by `paddle::Optional`",
out_name,
out_name));
VLOG(3) << "Custom Operator: GradOpDescMaker - handle unfound output: "
<< out_name;
continue;
}
VLOG(3) << "Custom Operator: GradOpDescMaker - output: " << out_name;
if (detail::IsDuplicableVar(out_name)) {
grad_op->SetOutput(
out_name,
this->InputGrad(detail::NoGrad(out_name, is_double_grad_),
/*drop_empty_grad=*/false));
} else {
grad_op->SetOutput(
out_name,
this->InputGrad(detail::NoGrad(out_name, is_double_grad_)));
}
}
grad_op->SetAttrMap(this->Attrs());
}

private:
std::string name_;
std::vector<std::string> inputs_;
std::vector<std::string> outputs_;
bool is_double_grad_{false};
};

template <>
class CustomGradOpMaker<imperative::OpBase>
: public SingleGradOpMaker<imperative::OpBase> {
public:
explicit CustomGradOpMaker(
const std::string& type,
const imperative::NameVarBaseMap& var_base_map_in,
const imperative::NameVarBaseMap& var_base_map_out,
const AttributeMap& attrs,
const std::map<std::string, std::string>& inplace_map,
const std::string& name,
const std::vector<std::string>& inputs,
const std::vector<std::string>& outputs,
bool is_double_grad)
: SingleGradOpMaker<imperative::OpBase>(
type, var_base_map_in, var_base_map_out, attrs, inplace_map),
name_(name),
inputs_(inputs),
outputs_(outputs),
is_double_grad_(is_double_grad) {}

protected:
// TODO(chenweihang): The code is duplicated with the previous one, because
// ere OpMaker's Input, Output and other methods are protected. Putting the
// function implementation outside the class will cause the method to be
// uncallable,
// so it is still implemented in the class for the time being.
void Apply(GradOpPtr<imperative::OpBase> grad_op) const override {
grad_op->SetType(name_);

auto fwd_op_inputs = this->InputNames();
auto fwd_op_outputs = this->OutputNames();

for (auto& in_name : inputs_) {
VLOG(3) << "Custom Operator: GradOpBaseMaker - input: " << in_name;
if (!detail::IsGradVar(in_name, is_double_grad_)) {
if (detail::IsMemberOf(fwd_op_inputs, in_name)) {
grad_op->SetInput(in_name, this->Input(in_name));
} else if (detail::IsMemberOf(fwd_op_outputs, in_name)) {
grad_op->SetInput(in_name, this->Output(in_name));
} else {
PADDLE_THROW(platform::errors::InvalidArgument(
"The input tensor name `%s` is invalid, expected it is the input "
"or output of forward operator.",
in_name));
}
} else {
// Handle inplace optional case
if (this->HasOutput(detail::NoGrad(in_name))) {
grad_op->SetInput(in_name, this->OutputGrad(detail::NoGrad(in_name)));
} else {
PADDLE_ENFORCE(
in_name.find(paddle::kOptionalSuffix) != std::string::npos,
phi::errors::InvalidArgument(
"Custom operator couldn't find grad operator input name for "
"%s. If you are using inplace optional inputs & outputs, "
"please check your InplaceMap and `Outputs` again and make "
"sure %s is wrapped by `paddle::Optional`",
in_name,
in_name));
VLOG(3) << "Custom Operator: GradOpBaseMaker - handle unfound input: "
<< in_name;
}
}
}
for (auto& out_name : outputs_) {
// Handle inplace optional case
if (!this->HasInput(detail::NoGrad(out_name, is_double_grad_))) {
PADDLE_ENFORCE(
out_name.find(paddle::kOptionalSuffix) != std::string::npos,
phi::errors::InvalidArgument(
"Custom operator couldn't find grad operator output name for "
"%s. If you are using inplace optional inputs & outputs, "
"please check your InplaceMap and `Outputs` again and make "
"sure %s is wrapped by `paddle::Optional`",
out_name,
out_name));
VLOG(3) << "Custom Operator: GradOpBaseMaker - handle unfound output: "
<< out_name;
continue;
}
VLOG(3) << "Custom Operator: GradOpBaseMaker - output: " << out_name;
grad_op->SetOutput(
out_name, this->InputGrad(detail::NoGrad(out_name, is_double_grad_)));
}
grad_op->SetAttrMap(this->Attrs());
}

private:
std::string name_;
std::vector<std::string> inputs_;
std::vector<std::string> outputs_;
bool is_double_grad_{false};
};

//////////// Operator and Kernel Register //////////////

static void RegisterOperatorKernelWithPlace(
Expand Down
Loading

0 comments on commit df648e5

Please sign in to comment.