wehub-resource-sync
186 lines
7.2 KiB
Django/Jinja
186 lines
7.2 KiB
Django/Jinja
{% import "common.j2" as common %}
|
|
// Auto Generated by decomp_gen.py, DO NOT EDIT!
|
|
|
|
#include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
|
|
#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
|
|
#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
|
|
#include "paddle/fluid/primitive/decomp_rule/decomp_rule/composite.h"
|
|
#include "paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h"
|
|
#include "paddle/fluid/primitive/vjp_interface/generated/generated_vjp.h"
|
|
#include "paddle/fluid/primitive/base/lazy_tensor.h"
|
|
#include "paddle/phi/api/include/tensor.h"
|
|
#include "paddle/phi/common/int_array.h"
|
|
#include "paddle/pir/include/core/builtin_op.h"
|
|
#include "paddle/pir/include/core/op_base.h"
|
|
|
|
namespace paddle {
|
|
namespace dialect {
|
|
using IntArray = paddle::experimental::IntArray;
|
|
{% macro sig(fwd_name, class_name, inputs, attrs, outputs) %}
|
|
{% set input_names=[] %}
|
|
{% set attr_names=[] %}
|
|
{% set output_names=[] %}
|
|
{% set output_types=[] %}
|
|
|
|
std::vector<std::vector<pir::Value>> {{class_name}}::DecompVjp(pir::Operation* op) {
|
|
VLOG(4) << "Decomp call {{fwd_name}}'s decomp interface begin";
|
|
|
|
{{class_name}} op_obj = op->dyn_cast<{{class_name}}>();
|
|
(void)op_obj;
|
|
|
|
FLAGS_tensor_operants_mode = "static";
|
|
|
|
VLOG(6) << "Decomp Prepare inputs of {{fwd_name}}";
|
|
|
|
{% for item in inputs -%}
|
|
{% do input_names.append(item.name) %}
|
|
{% if item.typename == "Tensor" %} {#- Tensor or Tensor[] #}
|
|
{% if item.optional %}
|
|
paddle::optional<Tensor> {{item.name}};
|
|
if (!IsEmptyValue(op_obj.{{item.name}}())){
|
|
{{item.name}} = paddle::make_optional<Tensor>(Tensor(std::make_shared<primitive::LazyTensor>(op_obj.{{item.name}}())));
|
|
}
|
|
{% else %}
|
|
{{item.typename}} {{item.name}}(std::make_shared<primitive::LazyTensor>(op_obj.{{item.name}}()));
|
|
{% endif %}
|
|
{% elif item.typename == "Tensor[]" %}
|
|
{% if item.optional %}
|
|
|
|
paddle::optional<std::vector<Tensor>> {{item.name}};
|
|
if (!IsEmptyValue(op_obj.{{item.name}}())){
|
|
pir::CombineOp combine_op_obj =
|
|
op_obj.{{item.name}}().defining_op()->dyn_cast<pir::CombineOp>();
|
|
std::vector<Tensor> optional_{{item.name}};
|
|
for (size_t idx = 0; idx < combine_op_obj.inputs().size(); idx++) {
|
|
optional_{{item.name}}.emplace_back(
|
|
std::make_shared<primitive::LazyTensor>(combine_op_obj.inputs()[idx]));
|
|
}
|
|
{{item.name}} = paddle::make_optional<std::vector<Tensor>>(optional_{{item.name}});
|
|
}
|
|
|
|
{% else %}
|
|
pir::CombineOp combine_op_obj_{{item.name}} =
|
|
op_obj.{{item.name}}().defining_op()->dyn_cast<pir::CombineOp>();
|
|
std::vector<Tensor> {{item.name}};
|
|
for (size_t idx = 0; idx < combine_op_obj_{{item.name}}.inputs().size(); idx++) {
|
|
{{item.name}}.emplace_back(
|
|
std::make_shared<primitive::LazyTensor>(combine_op_obj_{{item.name}}.inputs()[idx]));
|
|
}
|
|
{% endif %}
|
|
{% endif %}
|
|
{% endfor %}
|
|
|
|
VLOG(6) << "Decomp prepare attributes of {{fwd_name}}";
|
|
{% if attrs %}
|
|
{% for item in attrs %}
|
|
{% do attr_names.append(item.name) %}
|
|
{% if item.typename == "Scalar" and item.support_tensor %}
|
|
|
|
Tensor {{item.name}}_(std::make_shared<primitive::LazyTensor>(op_obj.{{item.name}}()));
|
|
|
|
auto* {{item.name}}_define_op =
|
|
std::static_pointer_cast<primitive::LazyTensor>({{item.name}}_.impl())
|
|
->value()
|
|
.defining_op();
|
|
if ({{item.name}}_define_op->name() != "pd_op.full") {
|
|
PADDLE_THROW(
|
|
common::errors::Unimplemented("We don't support dynamic tensors "
|
|
"attribute {{item.name}} for {{fwd_name}} decomposition "
|
|
"for now. "));
|
|
}
|
|
Scalar {{item.name}} = {{item.name}}_define_op->attribute("value").dyn_cast<paddle::dialect::ScalarAttribute>().data();
|
|
|
|
{% elif item.typename == "IntArray" and item.support_tensor %}
|
|
|
|
Tensor {{item.name}}_(std::make_shared<primitive::LazyTensor>(op_obj.{{item.name}}()));
|
|
|
|
auto* {{item.name}}_define_op =
|
|
std::static_pointer_cast<primitive::LazyTensor>({{item.name}}_.impl())
|
|
->value()
|
|
.defining_op();
|
|
if ({{item.name}}_define_op->name() != "pd_op.full_int_array") {
|
|
PADDLE_THROW(
|
|
common::errors::Unimplemented("We don't support dynamic tensors "
|
|
"attribute {{item.name}} for {{fwd_name}} decomposition "
|
|
"for now. "));
|
|
}
|
|
IntArray {{item.name}} = phi::IntArray(
|
|
paddle::dialect::GetInt64Vector({{item.name}}_define_op->attribute("value")));
|
|
|
|
{% else %}
|
|
{% if item.mapped_type[0] == "pir::StrAttribute" %}
|
|
{{item.mapped_type[1]}} {{item.name}} = op->attribute("{{item.name}}").dyn_cast<{{item.mapped_type[0]}}>().AsString();
|
|
{% elif "[]" in item.typename %}
|
|
auto array_list = op->attribute("{{item.name}}").dyn_cast<pir::ArrayAttribute>().AsVector();
|
|
{% set temp_type= item.mapped_type[0]|replace('pir::ArrayAttribute<', '')|replace('>', '')%}
|
|
{{item.mapped_type[1]|replace('const ', '')|replace('&', '')}} {{item.name}};
|
|
if (array_list.size() > 0) {
|
|
if (array_list[0].isa<{{temp_type}}>()) {
|
|
for (size_t i = 0; i < array_list.size(); ++i) {
|
|
{{item.name}}.push_back(
|
|
array_list[i].dyn_cast<{{temp_type}}>().data());
|
|
}
|
|
|
|
} else {
|
|
PADDLE_THROW(common::errors::Unimplemented("attr is not vector of {{temp_type}} "));
|
|
}
|
|
}
|
|
{% else %}
|
|
{{item.mapped_type[1]}} {{item.name}} = op->attribute("{{item.name}}").dyn_cast<{{item.mapped_type[0]}}>().data();
|
|
{% endif %}
|
|
{% endif %}
|
|
{% endfor %}
|
|
{% endif %}
|
|
|
|
VLOG(6) << "Decomp call {{fwd_name}}'s backward composite rule prepare";
|
|
|
|
std::vector<std::vector<bool>> stop_gradients = ConstructStopGradient(op);
|
|
|
|
std::vector<std::vector<paddle::Tensor>> tensor_res;
|
|
for (auto arg : stop_gradients) {
|
|
tensor_res.push_back(std::vector<paddle::Tensor>(arg.size()));
|
|
}
|
|
std::string op_name = "{{fwd_name}}";
|
|
FLAGS_tensor_operants_mode = "static";
|
|
VLOG(4) << "Call Pir Decomposed backward op {{fwd_name}}";
|
|
|
|
|
|
{% for k in range(outputs|length) %}
|
|
paddle::Tensor* {{outputs[k].name}} = !stop_gradients[{{k}}][0] ? &tensor_res[{{k}}][0] : nullptr;
|
|
{% endfor %}
|
|
|
|
{% for item in outputs %}
|
|
{% do output_names.append(item.name) %}
|
|
{% do output_types.append(item.mapped_type) %}
|
|
{% endfor %}
|
|
|
|
paddle::primitive::details::{{fwd_name}}<primitive::LazyTensor>(
|
|
{{common.args(input_names, attr_names)}}, {{common.sequence('', '', ', ', output_names)}});
|
|
std::vector<std::vector<pir::Value>> res(tensor_res.size());
|
|
|
|
for (size_t i = 0; i < tensor_res.size(); ++i) {
|
|
res[i].resize(tensor_res[i].size());
|
|
for (size_t j = 0; j < tensor_res[i].size(); ++j) {
|
|
if (tensor_res[i][j].defined()) {
|
|
res[i][j] = std::static_pointer_cast<primitive::LazyTensor>(
|
|
tensor_res[i][j].impl())
|
|
->value();
|
|
}
|
|
}
|
|
}
|
|
|
|
VLOG(4) << "Decomp call {{fwd_name}}'s decomp interface end";
|
|
return res;
|
|
}
|
|
{% endmacro %}
|
|
|
|
{% for api in apis -%}
|
|
{% if api.name in decomp_vjp_white_list %}
|
|
{{sig(api.name, api.class_name, api.inputs, api.attrs, api.outputs)}}
|
|
{% else %} {# render nothing #}
|
|
{% endif %}
|
|
{% endfor %}
|
|
|
|
} // namespace dialect
|
|
} // namespace paddle
|