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/pir.cc
T
wehub-resource-sync e25996e7db chore: import upstream snapshot with attribution
2026-07-13 12:40:42 +08:00

3625 lines
142 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.
#include "paddle/fluid/pybind/pir.h"
#include <Python.h>
#include <algorithm>
#include <iterator>
#include <memory>
#include <sstream>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include "paddle/common/enforce.h"
#include "paddle/common/errors.h"
#include "paddle/common/flags.h"
#include "paddle/fluid/framework/executor.h"
#include "paddle/fluid/framework/ir/pass.h"
#include "paddle/fluid/framework/program_desc.h"
#include "paddle/fluid/framework/tensor_util.h"
#include "paddle/fluid/ir_adaptor/translator/program_translator.h"
#include "paddle/fluid/ir_adaptor/translator/translate.h"
#include "paddle/fluid/ir_adaptor/translator/utils.h"
#include "paddle/fluid/pir/dialect/distributed/ir/dist_attribute.h"
#include "paddle/fluid/pir/dialect/distributed/ir/dist_dialect.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/kernel/ir/kernel_type.h"
#include "paddle/fluid/pir/dialect/operator/interface/op_yaml_info.h"
#include "paddle/fluid/pir/dialect/operator/ir/api_builder.h"
#include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.h"
#include "paddle/fluid/pir/dialect/operator/ir/manual_pylayer_op.h"
#include "paddle/fluid/pir/dialect/operator/ir/op_attribute.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/ir/pd_api.h"
#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
#include "paddle/fluid/pir/dialect/operator/trait/inplace.h"
#include "paddle/fluid/pir/dialect/operator/utils/op_yaml_info_parser.h"
#include "paddle/fluid/pir/dialect/operator/utils/shape_analysis_utils.h"
#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
#include "paddle/fluid/pir/drr/include/drr_pattern_base.h"
#include "paddle/fluid/pir/serialize_deserialize/include/ir_serialize.h"
#include "paddle/fluid/pir/transforms/general/common_subexpression_elimination_pass.h"
#include "paddle/fluid/pir/transforms/general/dead_code_elimination_pass.h"
#include "paddle/fluid/pir/transforms/gpu/fused_bn_add_act_pass.h"
#include "paddle/fluid/pir/transforms/passes.h"
#include "paddle/fluid/pir/utils/general_functions.h"
#include "paddle/fluid/pir/utils/name_analysis.h"
#include "paddle/fluid/pybind/control_flow_api.h"
#include "paddle/fluid/pybind/eager_utils.h"
#include "paddle/fluid/pybind/pir_utils.h"
#include "paddle/fluid/pybind/pybind_variant_caster.h"
#include "paddle/fluid/pybind/size.h"
#include "paddle/phi/common/data_type.h"
#include "paddle/phi/common/place.h"
#include "paddle/phi/core/distributed/auto_parallel/process_mesh.h"
#include "paddle/phi/core/enforce.h"
#include "paddle/pir/include/core/attribute.h"
#include "paddle/pir/include/core/block.h"
#include "paddle/pir/include/core/builtin_attribute.h"
#include "paddle/pir/include/core/builtin_op.h"
#include "paddle/pir/include/core/ir_mapping.h"
#include "paddle/pir/include/core/ir_printer.h"
#include "paddle/pir/include/core/operation.h"
#include "paddle/pir/include/core/program.h"
#include "paddle/pir/include/core/type.h"
#include "paddle/pir/include/core/value.h"
#include "paddle/pir/include/core/visitors.h"
#include "paddle/pir/include/dialect/control_flow/ir/cf_dialect.h"
#include "paddle/pir/include/dialect/control_flow/ir/cf_op.h"
#include "paddle/pir/include/dialect/shape/ir/shape_attribute.h"
#include "paddle/pir/include/dialect/shape/ir/shape_dialect.h"
#include "paddle/pir/include/dialect/shape/transforms/shape_optimization_pass.h"
#include "paddle/pir/include/dialect/shape/utils/original_attributes_filter.h"
#include "paddle/pir/include/pass/pass.h"
#include "paddle/pir/include/pass/pass_manager.h"
#include "paddle/pir/include/pass/pass_registry.h"
#include "pybind11/stl.h"
#ifdef PADDLE_WITH_CINN
#include "paddle/ap/include/paddle/hlir/op_dialect.h"
#include "paddle/ap/include/paddle/pass/add_pcc_pass.h"
#include "paddle/cinn/hlir/dialect/operator/ir/op_dialect.h"
#include "paddle/cinn/hlir/dialect/operator/transforms/add_cinn_pass.h"
#include "paddle/cinn/hlir/dialect/operator/transforms/check_infer_symbolic_util.h"
#include "paddle/cinn/hlir/dialect/operator/transforms/pir_to_py_code_converter.h"
#include "paddle/cinn/hlir/dialect/operator/transforms/reduce_as_to_sum_pass.h"
#include "paddle/cinn/hlir/dialect/operator/transforms/specify_input_dynamic_dim_util.h"
#include "paddle/cinn/hlir/framework/pir_compiler.h"
#include "paddle/pir/include/dialect/shape/utils/shape_analysis.h"
#endif
using paddle::dialect::ApiBuilder;
using paddle::dialect::DenseTensorArrayType;
using paddle::dialect::DenseTensorType;
using paddle::dialect::DistDenseTensorType;
using paddle::dialect::DistTypeInterface;
using paddle::dialect::IfOp;
using paddle::dialect::PrintOp;
using paddle::dialect::PyLayerOp;
using paddle::dialect::SelectedRowsType;
using paddle::dialect::SparseCooTensorType;
using paddle::dialect::SparseCsrTensorType;
using paddle::dialect::WhileOp;
using pir::TuplePopOp;
using paddle::dialect::IntArrayAttribute;
using paddle::dialect::OperationDistAttribute;
using paddle::dialect::PlaceAttribute;
using paddle::dialect::TensorDistAttribute;
using pir::ArrayAttribute;
using pir::Attribute;
using pir::Block;
using pir::BlockArgument;
using pir::BoolAttribute;
using pir::CloneOptions;
using pir::Int32Attribute;
using pir::Int64Attribute;
using pir::IrContext;
using pir::IrMapping;
using pir::Operation;
using pir::OpOperand;
using pir::OpResult;
using pir::Pass;
using pir::PassManager;
using pir::Program;
using pir::StrAttribute;
using pir::Type;
using pir::Value;
using pir::VectorType;
using pybind11::return_value_policy;
namespace name_analysis = pir::utils::name_analysis;
COMMON_DECLARE_bool(print_ir);
namespace paddle {
namespace pybind {
PyTypeObject *g_ir_value_pytype = nullptr;
namespace py = pybind11;
void BindOpsAPI(pybind11::module *module);
Value FakeValue() {
// create a fake value to simplify `ForwardBackwardSplit`.
return Value(nullptr);
}
bool IsFakeValue(const Value &value) {
// create a fake value to simplify `ForwardBackwardSplit`.
return value.impl() == nullptr || !value.type();
}
inline int64_t GetProgramInt64Attr(const std::shared_ptr<Program> &program,
const std::string &attr_name,
int64_t default_value = 0) {
auto op = program->module_op();
if (op->HasAttribute(attr_name)) {
auto val = op->attribute(attr_name).dyn_cast<pir::Int64Attribute>().data();
return val;
} else {
return default_value;
}
}
inline void SetProgramInt64Attr(std::shared_ptr<Program> program,
const std::string &attr_name,
int64_t value) {
auto op = program->module_op();
op->set_attribute(attr_name,
pir::Int64Attribute::get(IrContext::Instance(), value));
}
std::string GetValueInfo(Value v) {
if (v.impl() == nullptr) {
return "nullptr value";
}
std::stringstream ss;
if (auto op_result = v.dyn_cast<OpResult>()) {
ss << "define_op_name=" << op_result.owner()->name();
ss << ", index=" << op_result.index();
} else if (auto arg = v.dyn_cast<BlockArgument>()) {
if (arg.is_kwarg()) {
ss << "keyword block_arg, keyword = " << arg.keyword();
} else {
ss << "position block_arg, index = " << arg.index();
}
}
if (!v.type()) {
ss << ", dtype=<<NULL TYPE>>";
} else {
ss << ", dtype=" << v.type();
if (v.type().isa<paddle::dialect::AllocatedDenseTensorType>()) {
ss << ", place="
<< v.type()
.dyn_cast<paddle::dialect::AllocatedDenseTensorType>()
.place();
}
}
auto stop_gradient = v.attribute<BoolAttribute>(kAttrStopGradients);
if (stop_gradient && !stop_gradient.data()) {
ss << ", stop_gradient=False";
} else {
ss << ", stop_gradient=True";
}
return ss.str();
}
py::object Clone(const Program &self, IrMapping *p_mapper = nullptr) {
IrMapping mapper;
if (p_mapper == nullptr) {
p_mapper = &mapper;
}
auto src_obj = py::cast(self);
auto new_obj = py::cast(self.Clone(*p_mapper));
for (auto item : src_obj.attr("__dict__").cast<py::dict>()) {
new_obj.attr(item.first.cast<std::string>().c_str()) = item.second;
}
return new_obj;
}
bool SomeInSet(const std::vector<Value> &vec, const std::set<Value> &set) {
for (auto &v : vec) {
if (set.find(v) != set.end()) {
return true;
}
}
return false;
}
Value AppendDataOp(pir::Block *block,
const Value &value,
const std::string &name,
const Operation &origin_op) {
IrContext *ctx = IrContext::Instance();
auto op_info = ctx->GetRegisteredOpInfo(paddle::dialect::DataOp::name());
pir::AttributeMap attribute_map = {
{"name", StrAttribute::get(ctx, name)},
{"shape",
paddle::dialect::IntArrayAttribute::get(
ctx, phi::IntArray(phi::vectorize(GetValueDims(value))))},
{"dtype",
paddle::dialect::DataTypeAttribute::get(ctx, pir::GetValueDtype(value))},
{"place", PlaceAttribute::get(ctx, Place())}};
std::vector<pir::Type> output_types{value.type()};
Operation *operation =
Operation::Create({}, attribute_map, output_types, op_info);
block->insert(origin_op, operation);
return operation->result(0);
}
std::vector<Value> GetRealOpInputs(Operation *op) {
if (op->isa<paddle::dialect::IfOp>() ||
op->isa<paddle::dialect::PyLayerOp>()) {
return pir::GetUsedExternalValue(*op);
} else if (op->isa<paddle::dialect::WhileOp>()) {
paddle::dialect::WhileOp whileop = op->dyn_cast<paddle::dialect::WhileOp>();
auto value_vector = op->operands_source();
auto value_vector2 = pir::GetUsedExternalValue(whileop.body());
value_vector.insert(
value_vector.end(), value_vector2.begin(), value_vector2.end());
return value_vector;
} else {
return op->operands_source();
}
}
/*
Variables in input_vars will be the pruned program's inputs,
and variables in output_vars will be the pruned program's outputs.
Therefore, the pruning logic includes replacing the input of
input_vars with the data op, and then preserving all connected
ops starting from output_vars.
Note: The returned program is the original program.
If you do not want the original program to be modified,
please pass in a cloned result.
*/
void PruneWithInput(const std::vector<Value> &input_vars,
const std::vector<Value> &output_vars,
Program *prog) {
auto global_block = prog->block();
std::vector<Value> new_input_vars;
if (!input_vars.empty()) {
std::vector<Value> new_input_vars;
for (uint64_t idx = 0; idx < input_vars.size(); idx++) {
auto input = input_vars[idx];
auto origin_op = input.defining_op();
std::string name = name_analysis::TryGetValueFirstName(input).value_or(
"input_" + std::to_string(idx));
auto new_input = AppendDataOp(global_block, input, name, *origin_op);
input.ReplaceAllUsesWith(new_input);
new_input_vars.push_back(new_input);
}
}
VLOG(6) << "program after add new feed op = " << *prog;
auto total_ops_list = global_block->ops();
std::vector<Operation *> total_ops(total_ops_list.begin(),
total_ops_list.end());
std::vector<bool> intersection_op_flags(total_ops.size(), true);
std::set<Value> output_vars_set(output_vars.begin(), output_vars.end());
for (uint32_t index = total_ops.size() - 1; index != (uint32_t)(-1);
--index) {
auto op = total_ops[index];
auto op_results = op->results();
if (SomeInSet(op_results, output_vars_set)) {
for (auto &operand : GetRealOpInputs(op)) {
output_vars_set.insert(operand);
}
} else {
VLOG(6) << "delete op " << index << ", name is " << op->name();
intersection_op_flags[index] = false;
}
}
std::set<Value> input_vars_set(new_input_vars.begin(), new_input_vars.end());
std::vector<Operation *> remove_ops;
for (uint32_t index = total_ops.size() - 1; index != (uint32_t)(-1);
--index) {
auto op = total_ops[index];
if (!intersection_op_flags[index]) {
auto op_results = op->results();
if (!input_vars_set.empty() && SomeInSet(op_results, input_vars_set)) {
PADDLE_THROW(common::errors::InvalidArgument(
"The input_var create by: '{%s}' is not involved in the "
"output_vars calculation. "
"Please remove it from input_vars.",
op->name()));
}
global_block->erase(*op);
}
}
}
void SetIsTestAttr(const std::shared_ptr<Program> &prog) {
for (auto &op : prog->block()->ops()) {
if (op->HasAttribute("is_test")) {
op->set_attribute("is_test",
pir::BoolAttribute::get(IrContext::Instance(), true));
}
}
}
using ComputeReturnType = std::variant<float,
double,
int32_t,
int64_t,
bool,
std::string,
std::vector<int32_t>,
std::vector<int64_t>,
std::vector<float>,
DataType,
Place>;
ComputeReturnType CastPyObjectToAny(const pybind11::object &obj,
const std::string &type_name) {
static const std::unordered_map<
std::string,
std::function<ComputeReturnType(const pybind11::object &)>>
type_casters = {
{"float",
[](const pybind11::object &obj) { return obj.cast<float>(); }},
{"double",
[](const pybind11::object &obj) { return obj.cast<double>(); }},
{"int32",
[](const pybind11::object &obj) { return obj.cast<int32_t>(); }},
{"int64",
[](const pybind11::object &obj) { return obj.cast<int64_t>(); }},
{"bool",
[](const pybind11::object &obj) { return obj.cast<bool>(); }},
{"string",
[](const pybind11::object &obj) { return obj.cast<std::string>(); }},
{"vector<int32>",
[](const pybind11::object &obj) {
return obj.cast<std::vector<int32_t>>();
}},
{"vector<int64>",
[](const pybind11::object &obj) {
return obj.cast<std::vector<int64_t>>();
}},
{"vector<float>",
[](const pybind11::object &obj) {
return obj.cast<std::vector<float>>();
}},
{"datatype",
[](const pybind11::object &obj) { return obj.cast<DataType>(); }},
{"place",
[](const pybind11::object &obj) { return obj.cast<Place>(); }}};
auto it = type_casters.find(type_name);
if (it == type_casters.end()) {
throw std::runtime_error("Unsupported type: " + type_name);
}
return it->second(obj);
}
void BindProgram(py::module *m) {
static int64_t global_prog_seed = 0;
py::class_<Program, std::shared_ptr<Program>> program(
*m, "Program", py::dynamic_attr(), R"DOC(
Create Python Program. Program is an abstraction of model structure, divided into
computational graphs and weights. The Program has a main block that stores the computational
graphs.
A set of Program usually contains startup program and main program.
A startup program is set to contain some initial work, eg. initialize the ``Parameter``, and the main
program will contain the network structure and vars for train.
A set of Program can be used for test or train, in train program ,
Paddle will contain all content to build a train network, in test
program Paddle will prune some content which is irrelevant to test, eg.
backward ops and vars.
**Notes**:
**we have** :ref:`api_paddle_static_default_startup_program` **and** :ref:`api_paddle_static_default_main_program`
**by default, a pair of them will shared the parameters. The** :ref:`api_paddle_static_default_startup_program` **only run once to initialize parameters,**
:ref:`api_paddle_static_default_main_program` **run in every mini batch and adjust the weights.**
Returns:
Program: An empty Program.
Examples:
.. code-block:: pycon
>>> import paddle
>>> import paddle.static as static
>>> paddle.enable_static()
>>> main_program = static.Program()
>>> startup_program = static.Program()
>>> with static.program_guard(main_program=main_program, startup_program=startup_program):
... x = static.data(name="x", shape=[-1, 784], dtype='float32')
... y = static.data(name="y", shape=[-1, 1], dtype='int32')
... z = static.nn.fc(name="fc", x=x, size=10, activation="relu")
>>> print("main program is: {}".format(main_program))
>>> print("start up program is: {}".format(startup_program))
)DOC");
program
.def(py::init([]() {
auto prog = std::make_shared<Program>(IrContext::Instance());
SetProgramInt64Attr(prog, "random_seed", global_prog_seed);
return prog;
}))
.def("__str__",
[](const std::shared_ptr<Program> &self) {
std::ostringstream print_stream;
self->Print(print_stream);
return print_stream.str();
})
.def("__repr__",
[](const std::shared_ptr<Program> &self) {
std::ostringstream print_stream;
self->Print(print_stream);
return print_stream.str();
})
.def("parameters_num",
[](const std::shared_ptr<Program> &self) {
return self->parameters_num();
})
.def("set_is_test_attr",
[](const std::shared_ptr<Program> &self) { SetIsTestAttr(self); })
.def("set_parameters_from",
[](const std::shared_ptr<Program> &self,
const std::shared_ptr<Program> &other) {
self->set_parameters(other->parameters());
})
.def(
"global_block",
[](std::shared_ptr<Program> self) { return self->block(); },
return_value_policy::reference)
.def("clone", [](Program &self) { return Clone(self); })
.def("clone",
[](Program &self, IrMapping &ir_mapper) {
return Clone(self, &ir_mapper);
})
.def(
"copy_to_block",
[](std::shared_ptr<Program> self,
pir::IrMapping &mapper,
Block *block) { return self->CopyToBlock(mapper, block); },
return_value_policy::reference)
.def("list_vars",
[](std::shared_ptr<Program> self) {
py::list vars;
for (auto op : self->block()->ops()) {
for (auto var : op->results()) {
vars.append(var);
}
}
return vars;
})
.def("_list_named_vars",
[](std::shared_ptr<Program> self) {
return name_analysis::GetAllNamedValues(*self);
})
.def(
"global_block",
[](const std::shared_ptr<Program> &self) { return self->block(); },
return_value_policy::reference)
.def_property(
"random_seed",
[](const std::shared_ptr<Program> &self) {
return GetProgramInt64Attr(self, "random_seed", 0);
},
[](std::shared_ptr<Program> self, int64_t random_seed) {
SetProgramInt64Attr(self, "random_seed", random_seed);
})
.def_property(
"_seed",
[](const std::shared_ptr<Program> &self) {
return GetProgramInt64Attr(self, "random_seed", 0);
},
[](std::shared_ptr<Program> self, int64_t random_seed) {
SetProgramInt64Attr(self, "random_seed", random_seed);
})
.def("global_seed",
[](std::shared_ptr<Program> self, int64_t random_seed) {
global_prog_seed = random_seed;
SetProgramInt64Attr(self, "random_seed", random_seed);
})
.def_property_readonly(
"num_blocks",
[](const std::shared_ptr<Program> &self) {
size_t num_blocks = 0;
auto top_level_op = self->module_op();
for (size_t i = 0; i < top_level_op->num_regions(); ++i) {
auto &region = top_level_op->region(i);
num_blocks += region.size();
}
return num_blocks;
})
.def_property_readonly(
"blocks",
[](const std::shared_ptr<Program> &self) {
// Note: We only return global block currently.
py::list op_list;
op_list.append(self->block());
return op_list;
},
return_value_policy::reference)
.def(
"get_value_by_op_id",
[](Program &self, py::object op_ids) {
std::vector<int> op_ids_list;
if (py::isinstance<py::int_>(op_ids)) {
op_ids_list.push_back(op_ids.cast<int>());
} else if (py::isinstance<py::list>(op_ids)) {
for (auto item : op_ids) {
op_ids_list.push_back(item.cast<int>());
}
} else {
PADDLE_THROW(
"Invalid op_ids format. Please provide either a single "
"integer or a list of integers.");
}
std::list<Operation *> all_ops = self.block()->get_recursive_ops();
std::vector<Value> value_list;
for (auto op : all_ops) {
if (std::find(op_ids_list.begin(), op_ids_list.end(), op->id()) !=
op_ids_list.end()) {
for (auto value : op->results()) {
value_list.push_back(value);
}
}
}
if (value_list.empty()) {
PADDLE_THROW(
"Can't find the corresponding opresult from the op ids");
}
return value_list;
})
.def("get_output_value_by_name",
[](Program &self, const std::string &name) {
return name_analysis::GetOutputValueByName(self, name);
})
.def("get_parameter_value_by_name",
[](Program &self, const std::string &name) {
return name_analysis::GetParameterValueByName(self, name);
})
.def("get_all_parameter_values",
[](Program &self) {
return name_analysis::GetAllParameterValues(self);
})
.def("num_ops", [](Program &self) { return self.num_ops(); })
.def("_state_dict",
[](std::shared_ptr<Program> self,
const std::string &mode = "all",
const framework::Scope &scope = framework::Scope()) {
std::unordered_map<std::string, DenseTensor> state_dict_all;
std::unordered_map<std::string, DenseTensor> state_dict_param;
std::unordered_map<std::string, DenseTensor> state_dict_opt;
for (auto op : self->block()->ops()) {
for (auto var : op->results()) {
auto is_persistable =
var.attribute<BoolAttribute>(kAttrIsPersistable);
if (is_persistable && is_persistable.data()) {
if (var.defining_op()->isa<pir::ParameterOp>()) {
std::string var_name =
name_analysis::GetValueFirstName(var);
auto tensor =
scope.FindVar(var_name)->GetMutable<DenseTensor>();
state_dict_param[var_name] = *tensor;
state_dict_all[var_name] = *tensor;
} else if (var.defining_op()
->isa<paddle::dialect::DataOp>()) {
std::string var_name =
name_analysis::GetValueFirstName(var);
auto tensor =
scope.FindVar(var_name)->GetMutable<DenseTensor>();
state_dict_opt[var_name] = *tensor;
state_dict_all[var_name] = *tensor;
}
}
}
}
if (mode == "all") {
return state_dict_all;
} else if (mode == "param") {
return state_dict_param;
} else if (mode == "opt") {
return state_dict_opt;
} else {
PADDLE_THROW(common::errors::InvalidArgument(
"The mode is not supported."));
}
})
.def(
"set_state_dict",
[](std::shared_ptr<Program> self,
const std::unordered_map<std::string, DenseTensor> &state_dict,
const framework::Scope &scope = framework::Scope(),
bool copy_tensor = false) {
for (auto item : state_dict) {
auto var = scope.FindVar(item.first);
if (var == nullptr) {
PADDLE_THROW(common::errors::NotFound(
"The variable %s is not found.", item.first));
} else {
if (copy_tensor) {
auto *mutable_tensor = var->GetMutable<DenseTensor>();
paddle::framework::TensorCopy(
item.second, item.second.place(), mutable_tensor);
} else {
*var->GetMutable<DenseTensor>() = item.second;
}
}
}
},
py::arg("state_dict"),
py::arg("scope"),
py::arg("copy_tensor") = false)
.def(
"_prune",
[](Program &self, std::vector<Value> output_vars) {
std::vector<Value> input_vars;
PruneWithInput(input_vars, output_vars, &self);
return &self;
},
py::arg("targets"),
"A description for the _prune method")
.def(
"_prune_with_input",
[](Program &self,
std::vector<Value> input_vars,
std::vector<Value> output_vars) {
PruneWithInput(input_vars, output_vars, &self);
return &self;
},
py::arg("feeded_vars"),
py::arg("targets"))
.def("_sync_with_cpp", [](const std::shared_ptr<Program> &self) {
// It's not need _sync_with_cpp in pir, but it's necessary in old static
// graph. Add empty function to avoid python call error.
});
}
void RefreshOpStopgradients(Operation *op) {
if (op->num_operands() == 0 || op->isa<pir::ParameterOp>() ||
op->isa<paddle::dialect::UniformOp>()) {
return;
} else if (op->isa<pir::SliceOp>()) {
op->dyn_cast<pir::SliceOp>().RefreshStopGradients();
} else if (op->isa<pir::SplitOp>()) {
op->dyn_cast<pir::SplitOp>().RefreshStopGradients();
} else {
RefreshStopGradientsDefaultly(op);
}
}
void BindBlock(py::module *m) {
py::class_<Block> block(*m, "Block", R"DOC(
In IR, a Block has a list of Operation and can represent a sub computational graph.
Notes:
The constructor of Block should not be invoked directly. You can
use `Program.block()` to get a block.
)DOC");
block.def("empty", &Block::empty)
.def(
"__str__",
[](Block &self) {
std::ostringstream print_stream;
pir::IrPrinter printer(print_stream);
printer.PrintBlock(self);
return print_stream.str();
},
return_value_policy::reference)
.def(
"front",
[](Block &self) { return &self.front(); },
return_value_policy::reference)
.def(
"back",
[](Block &self) { return &self.back(); },
return_value_policy::reference)
.def_property_readonly(
"parent_op",
[](Block &self) { return self.GetParentOp(); },
return_value_policy::reference)
.def_property_readonly(
"program",
[](Block &self) { return self.GetParentOp()->GetParentProgram(); },
return_value_policy::reference)
.def_property_readonly(
"parent_block",
[](Block &self) { return self.GetParentOp()->GetParent(); },
return_value_policy::reference)
.def_property_readonly("ops",
[](Block &self) -> py::list {
py::list op_list;
for (auto &op : self) {
op_list.append(&op);
}
return op_list;
})
.def("num_ops", [](Block &self) { return self.num_ops(); })
.def(
"__enter__",
[](Block &self) -> Block & {
ApiBuilder::Instance().PushInsertionPoint();
ApiBuilder::Instance().SetInsertionPointToBlockEnd(&self);
return self;
},
return_value_policy::reference)
.def("__exit__",
[](Block &self, py::object, py::object, py::object) {
ApiBuilder::Instance().LoadInsertionPoint();
})
.def("__len__", [](Block &self) { return self.size(); })
.def("args", &Block::args, return_value_policy::reference)
.def("kwargs", &Block::kwargs, return_value_policy::reference)
.def("add_arg", &Block::AddArg)
.def("add_kwarg", &Block::AddKwarg)
.def("erase_kwarg", &Block::EraseKwarg)
.def("get_values_by_op_idx",
[](Block &self, const py::list &op_idxs) -> py::list {
py::list value_list;
auto it = self.begin();
std::set<int> idxs_set;
for (py::handle item : op_idxs) {
idxs_set.insert(item.cast<int>());
}
for (int i = 0; it != self.end(); ++i, ++it) {
if (idxs_set.find(i) != idxs_set.end()) {
for (uint32_t j = 0; j < it->num_results(); ++j) {
value_list.append(static_cast<Value>(it->result(j)));
}
}
}
return value_list;
})
.def("remove_op",
[](Block &self, const Operation &op) { self.erase(op); })
.def(
"move_op",
[](Block &self, Operation *op, uint32_t offset) {
Block::Iterator position = self.begin();
std::advance(position, offset);
op->MoveTo(&self, position);
},
R"DOC(
Move an op to a specific position (block.begin() + offset).
Args:
op (pir.Operation): the operator to be moved.
offset (uint32_t) : offset relative to the begin of the block
Returns:
None
)DOC")
.def(
"move_op_to_block_end",
[](Block &self, Operation *op) { op->MoveTo(&self, self.end()); },
R"DOC(
Move an op to the end of the block.
Args:
op (pir.Operation): The operator to be moved.
Returns:
None
)DOC")
.def("all_parameters",
[](Block &self) -> py::list {
py::list param_list;
for (auto &op : self) {
if (op.name() == "builtin.parameter" &&
op.HasAttribute(kAttrIsPersistable)) {
auto attrs = op.attribute(kAttrIsPersistable)
.dyn_cast<pir::ArrayAttribute>()
.AsVector();
for (uint32_t i = 0; i < attrs.size(); i++) {
bool is_persistable =
attrs[i].dyn_cast<pir::BoolAttribute>().data();
if (is_persistable) {
param_list.append(static_cast<Value>(op.result(i)));
}
}
}
}
return param_list;
})
.def("refresh_stopgradient",
[](Block &self) {
for (auto &op : self) {
RefreshOpStopgradients(&op);
}
})
.def("_sync_with_cpp", [](const Block &self) {
// It's not need _sync_with_cpp in pir, but it's necessary in old static
// graph. Add empty function to avoid python call error.
});
}
void BindIrMapping(py::module *m) {
py::class_<IrMapping> ir_mapping(*m, "IrMapping");
ir_mapping.def(py::init<>())
.def("look_up",
[](IrMapping &self, Value from) { return self.Lookup(from); })
.def("has", [](IrMapping &self, Value from) { return self.Has(from); })
.def("add",
[](IrMapping &self, Value from, Value to) {
self.Add<Value>(from, to);
})
.def("size",
[](IrMapping &self) { return self.GetMutableMap<Value>().size(); });
}
void BindCloneOptions(py::module *m) {
py::class_<CloneOptions> clone_options(*m, "CloneOptions");
clone_options.def(
"__init__",
[](CloneOptions &self,
bool clone_regions,
bool clone_operands,
bool clone_successors) {
new (&self)
CloneOptions(clone_regions, clone_operands, clone_successors);
},
return_value_policy::reference);
}
void BindOperation(py::module *m) {
py::class_<Operation> op(*m, "Operation", R"DOC(
In IR, all the operation are represented by Operation, and Operation
is regarded as a build in an instruction of a Block. Users can call
python api to describe their neural network.
Notes:
The constructor of operator should not be invoked directly. Use
python api, for example: paddle.mean for building mean operation.
)DOC");
op.def("name", &Operation::name)
.def("get_parent_block",
&Operation::GetParent,
return_value_policy::reference)
.def("num_operands", &Operation::num_operands)
.def("num_results", &Operation::num_results)
.def("num_regions", &Operation::num_regions)
.def("id", &Operation::id)
.def("operand", &Operation::operand)
.def("result",
[](Operation &self, uint32_t index) {
return static_cast<Value>(self.result(index));
})
.def("operand_source", &Operation::operand_source)
.def("operands", &Operation::operands)
.def("results",
[](Operation &self) -> py::list {
py::list value_list;
for (uint32_t i = 0; i < self.num_results(); i++) {
value_list.append(static_cast<Value>(self.result(i)));
}
return value_list;
})
.def(
"blocks",
[](Operation &self) { return &self.blocks(); },
return_value_policy::reference)
.def("has_attr", &Operation::HasAttribute)
.def("str_attr",
[](Operation &self, const std::string &attr_name) -> py::object {
auto str_attr = self.attribute<StrAttribute>(attr_name);
if (str_attr) {
return py::cast(str_attr.AsString());
} else {
return py::cast<py::none>(Py_None);
}
})
.def("int_attr",
[](Operation &self, const std::string &attr_name) -> py::object {
auto int_attr = self.attribute<Int32Attribute>(attr_name);
if (int_attr) {
return py::cast(int_attr.data());
} else {
return py::cast<py::none>(Py_None);
}
})
.def("set_bool_attr",
[](Operation &self, std::string &attr_name, bool flag) {
self.set_attribute(
attr_name,
pir::BoolAttribute::get(IrContext::Instance(), flag));
})
.def("set_int_array_attr",
[](Operation &self,
std::string &attr_name,
const std::vector<int64_t> &val) {
auto attr = IntArrayAttribute::get(IrContext::Instance(),
phi::IntArray(val));
self.set_attribute(attr_name, attr);
})
.def("set_str_array_attr",
[](Operation &self,
std::string &attr_name,
const std::vector<std::string> &val) {
std::vector<Attribute> val_attr;
for (auto &str : val) {
val_attr.emplace_back(
StrAttribute::get(IrContext::Instance(), str));
}
auto attr =
pir::ArrayAttribute::get(IrContext::Instance(), val_attr);
self.set_attribute(attr_name, attr);
})
.def("set_str_attr",
[](Operation &self, std::string &attr_name, std::string &val) {
self.set_attribute(attr_name,
StrAttribute::get(IrContext::Instance(), val));
})
.def("set_int_attr",
[](Operation &self, std::string &attr_name, const int &val) {
self.set_attribute(
attr_name,
pir::Int32Attribute::get(IrContext::Instance(), val));
})
.def("erase_attr",
[](Operation &self, std::string &attr_name) {
self.erase_attribute(attr_name);
})
.def("attrs",
[](Operation &self) -> py::dict {
py::dict attrs_dict;
for (auto &pair : self.attributes()) {
// SymbolAttribute is only used in PIR, no need to pass to Python
if (pair.second.isa<pir::shape::SymbolAttribute>()) continue;
if (pair.first == kAttrOpDistAttr) {
attrs_dict[pair.first.c_str()] =
pair.second.dyn_cast<OperationDistAttribute>();
} else {
if (pair.second.isa<pir::FloatAttribute>()) {
VLOG(2) << "The value is stored with float32 precision, "
"which may cause precision issues for higher "
"precision requirements.";
}
attrs_dict[pair.first.c_str()] =
paddle::dialect::GetAttributeData(pair.second);
}
}
return attrs_dict;
})
.def("copy_attrs_from",
[](Operation &self, Operation &other) {
for (auto &pair : other.attributes()) {
self.set_attribute(pair.first, pair.second);
}
})
.def("set_execution_stream",
[](Operation &self, const std::string &exe_stream) {
self.set_attribute(
"execution_stream",
StrAttribute::get(IrContext::Instance(), exe_stream));
})
.def("set_scheduling_priority",
[](Operation &self, int64_t priority) {
self.set_attribute(
"scheduling_priority",
pir::Int64Attribute::get(IrContext::Instance(), priority));
})
.def("operands_source",
[](Operation &self) -> py::list {
py::list op_list;
for (uint32_t i = 0; i < self.num_operands(); i++) {
op_list.append(self.operand_source(i));
}
return op_list;
})
.def("get_input_names",
[](Operation &self) -> py::list {
if (self.HasInterface<paddle::dialect::OpYamlInfoInterface>() ==
false) {
PADDLE_THROW(common::errors::InvalidArgument(
"Currently, we can only get input names of Operation that "
"has OpYamlInfoInterface"));
}
py::list op_list;
paddle::dialect::OpYamlInfoInterface yaml_interface =
self.dyn_cast<paddle::dialect::OpYamlInfoInterface>();
auto inputs_info = std::get<0>(yaml_interface.GetOpInfo());
for (auto &input_info : inputs_info) {
op_list.append(input_info.name);
}
return op_list;
})
.def("get_attr_names",
[](Operation &self) -> py::list {
py::list op_list;
paddle::dialect::OpYamlInfoInterface yaml_interface =
self.dyn_cast<paddle::dialect::OpYamlInfoInterface>();
auto attrs_info = std::get<1>(yaml_interface.GetOpInfo());
for (auto &attr_info : attrs_info) {
op_list.append(attr_info.name);
}
return op_list;
})
.def("get_output_names",
[](Operation &self) -> py::list {
py::list op_list;
paddle::dialect::OpYamlInfoInterface yaml_interface =
self.dyn_cast<paddle::dialect::OpYamlInfoInterface>();
auto outputs_info = std::get<2>(yaml_interface.GetOpInfo());
for (auto &output_info : outputs_info) {
op_list.append(output_info.name);
}
return op_list;
})
.def("get_output_intermediate_status",
[](Operation &self) -> py::list {
py::list op_list;
paddle::dialect::OpYamlInfoInterface yaml_interface =
self.dyn_cast<paddle::dialect::OpYamlInfoInterface>();
auto outputs_info = std::get<2>(yaml_interface.GetOpInfo());
for (auto &output_info : outputs_info) {
op_list.append(output_info.intermediate);
}
return op_list;
})
.def("get_input_grad_semantics",
[](Operation &self) -> py::list {
if (self.HasInterface<paddle::dialect::OpYamlInfoInterface>() ==
false) {
PADDLE_THROW(common::errors::InvalidArgument(
"Currently, we can only get input grad semantics of "
"Operation that "
"has OpYamlInfoInterface"));
}
py::list op_list;
paddle::dialect::OpYamlInfoInterface yaml_interface =
self.dyn_cast<paddle::dialect::OpYamlInfoInterface>();
auto inputs_grad_info = std::get<0>(yaml_interface.GetOpInfo());
for (auto &input_grad_info : inputs_grad_info) {
op_list.append(input_grad_info.with_grad_semantic);
}
return op_list;
})
.def("replace_all_uses_with",
[](Operation &self, const std::vector<Value> &values) {
self.ReplaceAllUsesWith(values);
})
.def("as_if_op",
[](Operation &self) { return PyIfOp(self.dyn_cast<IfOp>()); })
.def("as_pylayer_op",
[](Operation &self) -> PyLayerOp {
auto pylayer_op = self.dyn_cast<PyLayerOp>();
if (!pylayer_op) {
PADDLE_THROW(common::errors::InvalidArgument(
"Can't cast non-pylayer_op type Operation to PyLayerOp."));
}
return pylayer_op;
})
.def("as_while_op",
[](Operation &self) { return PyWhileOp(self.dyn_cast<WhileOp>()); })
.def(
"as_tuple_pop_op",
[](Operation &self) -> TuplePopOp {
auto tuple_pop_op = self.dyn_cast<TuplePopOp>();
if (!tuple_pop_op) {
PADDLE_THROW(common::errors::InvalidArgument(
"Can't cast non-tuple_pop_op type Operation to TuplePopOp."));
}
return tuple_pop_op;
})
.def("__repr__",
[](Operation &self) {
std::ostringstream print_stream;
print_stream << "Operation(";
self.Print(print_stream);
print_stream << ")";
return print_stream.str();
})
.def(
"clone",
[](Operation &self, IrMapping &ir_mapping, CloneOptions options) {
auto op = self.Clone(ir_mapping, options);
return ApiBuilder::Instance().GetBuilder()->Insert(op);
},
return_value_policy::reference)
.def("erase", &Operation::Erase)
.def("move_before",
[](Operation &self, Operation &other) {
self.MoveTo(other.GetParent(), Block::Iterator{other});
})
.def_property(
"callstack",
[](Operation &self) -> py::list {
py::list callstack_list;
if (!self.HasAttribute(paddle::framework::OpProtoAndCheckerMaker::
OpCreationCallstackAttrName())) {
return callstack_list;
}
pir::Attribute op_callstack = self.attribute<pir::Attribute>(
paddle::framework::OpProtoAndCheckerMaker::
OpCreationCallstackAttrName());
PADDLE_ENFORCE(op_callstack.isa<pir::ArrayAttribute>(),
common::errors::PreconditionNotMet(
"The callstack of operation `%s` should be an "
"array attribute.",
self.name()));
auto op_callstack_array_attr =
op_callstack.dyn_cast<pir::ArrayAttribute>();
for (size_t i = 0; i < op_callstack_array_attr.size(); ++i) {
PADDLE_ENFORCE(
op_callstack_array_attr.at(i).isa<StrAttribute>(),
common::errors::PreconditionNotMet(
"The callstack info of operation `%s` should be array of "
"string attribute.",
self.name()));
callstack_list.append(op_callstack_array_attr.at(i)
.dyn_cast<StrAttribute>()
.AsString());
}
return callstack_list;
},
[](Operation &self,
const std::vector<std::string> &callstack) -> void {
std::vector<pir::Attribute> op_callstack_infos;
for (auto str : callstack) {
op_callstack_infos.push_back(
StrAttribute::get(IrContext::Instance(), str));
}
self.set_attribute(paddle::framework::OpProtoAndCheckerMaker::
OpCreationCallstackAttrName(),
pir::ArrayAttribute::get(IrContext::Instance(),
op_callstack_infos));
})
.def_property(
"dist_attr",
[](Operation &self) -> py::object {
if (self.HasAttribute(kAttrOpDistAttr)) {
return py::cast(
self.attribute<OperationDistAttribute>(kAttrOpDistAttr));
} else {
return py::cast<py::none>(Py_None);
}
},
[](Operation &self, OperationDistAttribute op_dist_attr) {
self.set_attribute(kAttrOpDistAttr, op_dist_attr);
})
.def_property(
"op_role",
[](Operation &self) -> py::object {
auto int_attr = self.attribute<Int32Attribute>("op_role");
if (int_attr) {
return py::cast(int_attr.data());
} else {
return py::cast(-1);
}
},
[](Operation &self, const int &op_role) {
self.set_attribute(
"op_role", Int32Attribute::get(IrContext::Instance(), op_role));
})
.def_property(
"chunk_id",
[](Operation &self) -> py::object {
auto int_attr = self.attribute<Int32Attribute>("chunk_id");
if (int_attr) {
return py::cast(int_attr.data());
} else {
return py::cast(-1);
}
},
[](Operation &self, const int &chunk_id) {
self.set_attribute(
"chunk_id",
Int32Attribute::get(IrContext::Instance(), chunk_id));
})
.def("is_no_need_buffer",
[](Operation &self, const Value &operand_source) -> bool {
paddle::dialect::OpYamlInfoInterface op_info_interface =
self.dyn_cast<paddle::dialect::OpYamlInfoInterface>();
std::unique_ptr<paddle::dialect::OpYamlInfoParser> info_parser(
nullptr);
if (op_info_interface) {
info_parser =
std::make_unique<paddle::dialect::OpYamlInfoParser>(
op_info_interface.GetOpInfo(),
paddle::dialect::IsLegacyOp(self.name()));
auto &no_need_buffer_ids = info_parser->NoNeedBufferIds();
for (auto no_need_buffer_id : no_need_buffer_ids) {
if (operand_source == self.operand_source(no_need_buffer_id)) {
return true;
}
}
}
return false;
});
py::class_<Operation::BlockContainer> block_container(
*m, "Operation_BlockContainer", R"DOC(
The Operation_BlockContainer only use to walk all blocks in the operation.
)DOC");
block_container.def(
"__iter__",
[](Operation::BlockContainer &self) {
return py::make_iterator(self.begin(), self.end());
},
py::keep_alive<0, 1>());
}
py::str Value2String(Value self) {
std::ostringstream print_stream;
print_stream << "Value(";
print_stream << GetValueInfo(self);
print_stream << ")";
return print_stream.str();
}
const phi::DDim &GetTensorDims(Type type) {
if (!type) {
PADDLE_THROW(common::errors::InvalidArgument(
"The type used to get dims is nullptr."));
}
if (auto dense_type = type.dyn_cast<DenseTensorType>()) {
return dense_type.dims();
} else if (auto select_rows_type = type.dyn_cast<SelectedRowsType>()) {
return select_rows_type.dims();
} else if (auto sparse_coo_tensor_type =
type.dyn_cast<SparseCooTensorType>()) {
return sparse_coo_tensor_type.dims();
} else if (auto sparse_csr_tensor_type =
type.dyn_cast<SparseCsrTensorType>()) {
return sparse_csr_tensor_type.dims();
} else if (auto dense_array_type = type.dyn_cast<DenseTensorArrayType>()) {
return dense_array_type.dims();
} else {
PADDLE_THROW(common::errors::InvalidArgument(
"Currently, we can only get shape for dense and select rows type."));
}
}
const phi::DDim &GetValueDims(Value value) {
return GetTensorDims(value.type());
}
Value apply(Value self, py::object func) {
py::gil_scoped_acquire gil;
auto stop_gradient = self.attribute<BoolAttribute>(kAttrStopGradients);
if (stop_gradient && !stop_gradient.data()) {
PADDLE_THROW(common::errors::Unavailable(
"Cannot apply function on a tensor that required gradient."));
}
PyObject *py_func = func.release().ptr();
Py_INCREF(py_func);
PyObject *res = nullptr;
try {
py::object obj = py::cast(self);
PyObject *tmp_self = obj.release().ptr();
Py_INCREF(tmp_self);
res = PyObject_CallFunctionObjArgs(py_func, tmp_self, nullptr);
Py_DECREF(tmp_self);
} catch (std::exception &e) {
PADDLE_THROW(common::errors::Unavailable(
"Apply function of Tensor raises an exception: %s.", e.what()));
} catch (...) {
PADDLE_THROW(common::errors::Fatal(
"Apply function of Tensor raises an unknown exception."));
}
if (res == Py_None) {
return self;
}
auto out = CastPyArg2Value(res, "", 0, false);
Py_DECREF(py_func);
Py_DECREF(res);
return out;
}
#define DEF_VALUE_BOOL_PROPERTY(name) \
def_property( \
name, \
[](Value self) { \
auto bool_data = self.attribute<BoolAttribute>(name); \
return bool_data && bool_data.data(); \
}, \
[](Value self, bool bool_data) { \
self.set_attribute( \
name, BoolAttribute::get(IrContext::Instance(), bool_data)); \
})
#define DEF_VALUE_STOP_GRADIENT_PROPERTY(name) \
def_property( \
name, \
[](Value self) { \
auto bool_data = self.attribute<BoolAttribute>(name); \
return !bool_data || bool_data.data(); \
}, \
[](Value self, bool bool_data) { \
self.set_attribute( \
name, BoolAttribute::get(IrContext::Instance(), bool_data)); \
})
#define DEF_VALUE_POINTER_PROPERTY(name) \
def_property( \
name, \
[](Value self) -> py::object { \
auto prop_ptr = self.property(name); \
if (!prop_ptr) { \
return py::cast<py::none>(Py_None); \
} \
auto py_data = reinterpret_cast<PyObject *>(prop_ptr); \
py::object obj = \
py::reinterpret_borrow<py::object>(py::handle(py_data)); \
return obj; \
}, \
[](Value self, py::object obj) { \
pir::PropertiesDeleter deleter = [](void *python_obj) { \
Py_DECREF(python_obj); \
}; \
PyObject *pointer_data = obj.release().ptr(); \
pir::Property value_property(reinterpret_cast<void *>(pointer_data), \
deleter); \
self.set_property(name, value_property); \
})
void BindValue(py::module *m) {
py::class_<Value> value(*m,
"Value",
R"DOC(
Value class represents the SSA value in the IR system. It is a directed edge
and a base class.
Notes:
The constructor of Value should not be invoked directly. Value can be automatically constructed
when build network.
)DOC");
g_ir_value_pytype = reinterpret_cast<PyTypeObject *>(value.ptr());
value.def(py::init<>())
.def(py::init([](Value value) { return value; }))
.def_property_readonly(
"block",
[](Value self) {
if (auto op_result = self.dyn_cast<OpResult>()) {
return op_result.owner()->GetParent();
}
return self.dyn_cast<BlockArgument>().owner();
},
return_value_policy::reference)
.def_property_readonly(
"id",
[](Value self) {
if (self.impl() == nullptr) {
PADDLE_THROW(common::errors::InvalidArgument(
"Currently, we can only get id of Value whose impl "
"is not nullptr"));
} else {
std::stringstream ss;
ss << std::hex << self.impl();
return ss.str();
}
})
.def_property(
"name",
[](Value self) -> std::string {
return name_analysis::GetValueFirstName(self);
},
[](Value self, const std::string &name) {
name_analysis::SetValueName(self, name);
})
.def_property_readonly(
"has_name",
[](Value self) {
return name_analysis::TryGetValueFirstName(self).has_value();
})
// Return all Maybe names of given Value, for example:
// DataOp("var_1") -> %0 -> shadow_output("output_2")
// Return ["var_1", "output_2"]
.def_property_readonly("_names",
[](Value self) -> py::list {
std::vector<std::string> names =
name_analysis::GetValueAllNames(self);
return py::cast(names);
})
.def_property(
"shape",
[](Value self) {
auto array = phi::vectorize(GetValueDims(self));
auto ptr =
Paddle_Size_NewFromInt64Array(array.data(), array.size());
if (!ptr) {
throw py::error_already_set();
}
return py::reinterpret_steal<py::object>(ptr);
},
[](Value self, const std::vector<int> &shape) {
PADDLE_THROW(common::errors::InvalidArgument(
"can't set shape when building static graph"));
})
.def_property(
"_local_shape",
[](Value self) {
if (!self.type().isa<DistDenseTensorType>()) {
PADDLE_THROW(common::errors::InvalidArgument(
"_local_shape is only for distdense tensor."));
}
return phi::vectorize(
self.type().dyn_cast<DistDenseTensorType>().local_ddim());
},
[](Value self, const std::vector<int> &shape) {
PADDLE_THROW(common::errors::InvalidArgument(
"can't set _local_shape when building static graph"));
})
.def_property(
"dtype",
[](Value self) { return pir::GetValueDtype(self); },
[](Value self, DataType dtype) {
PADDLE_THROW(common::errors::InvalidArgument(
"can't set dtype when building static graph"));
})
.def_property(
"place_attr",
[](Value self) -> Place {
auto place_attr = self.attribute<PlaceAttribute>("place");
return place_attr ? place_attr.data() : Place();
},
[](Value self, const Place &place) {
// auto place = CastPyArg2Place(place_obj.release().ptr(), 1);
auto place_attr =
dialect::PlaceAttribute::get(IrContext::Instance(), place);
self.set_attribute("place", place_attr);
})
.def("initialized",
[](Value self) {
if (self.impl() == nullptr || self.type().storage() == nullptr) {
return false;
} else {
return true;
}
})
.DEF_VALUE_STOP_GRADIENT_PROPERTY("stop_gradient")
.DEF_VALUE_BOOL_PROPERTY("trainable")
.DEF_VALUE_BOOL_PROPERTY("persistable")
.DEF_VALUE_BOOL_PROPERTY("need_clip")
.DEF_VALUE_BOOL_PROPERTY("is_distributed")
.DEF_VALUE_BOOL_PROPERTY("is_parameter")
.DEF_VALUE_POINTER_PROPERTY("optimize_attr")
.DEF_VALUE_POINTER_PROPERTY("regularizer")
.DEF_VALUE_POINTER_PROPERTY("do_model_average")
.def("all_used_ops",
[](Value &self) -> py::list {
py::list op_list;
for (auto it = self.use_begin(); it != self.use_end(); ++it) {
op_list.append(it.owner());
}
return op_list;
})
.def("all_used_ops_in_same_block",
[](Value &self) -> py::list {
py::list op_list;
for (auto it = self.use_begin(); it != self.use_end(); ++it) {
Operation *used_op = it.owner();
while (used_op->GetParent() != self.defining_op()->GetParent() &&
used_op->GetParent()->GetParentOp()) {
used_op = used_op->GetParent()->GetParentOp();
}
op_list.append(used_op);
}
return op_list;
})
.def(
"get_defining_op",
[](Value self) -> Operation * { return self.defining_op(); },
return_value_policy::reference)
.def("type", &Value::type)
.def("index",
[](Value self) -> uint32_t {
if (auto op_result = self.dyn_cast<OpResult>()) {
return op_result.index();
} else if (auto arg = self.dyn_cast<BlockArgument>()) {
if (!arg.is_kwarg()) {
return arg.index();
}
}
PADDLE_THROW(common::errors::InvalidArgument(
"only support accessing index from op_result or positional "
"block arg."));
})
.def("is_dense_tensor_type",
[](Value self) { return self.type().isa<DenseTensorType>(); })
.def("is_selected_row_type",
[](Value self) { return self.type().isa<SelectedRowsType>(); })
.def("is_sparse_coo_tensor_type",
[](Value self) { return self.type().isa<SparseCooTensorType>(); })
.def("is_sparse_csr_tensor_type",
[](Value self) { return self.type().isa<SparseCsrTensorType>(); })
.def("is_dense_tensor_array_type",
[](Value self) { return self.type().isa<DenseTensorArrayType>(); })
.def("is_dist_dense_tensor_type",
[](Value self) { return self.type().isa<DistDenseTensorType>(); })
.def("value_assign", [](Value &self, Value value) { self = value; })
.def("replace_all_uses_with",
[](Value self, Value value) { self.ReplaceAllUsesWith(value); })
.def("replace_grad_users_with",
[](Value self,
Value value,
std::unordered_set<Operation *> &grad_ops) {
for (auto it = self.use_begin(); it != self.use_end();) {
auto use_op = it.owner();
if (grad_ops.find(use_op) != grad_ops.end()) {
(it++)->set_source(value);
} else {
it++;
}
}
})
.def("set_type", [](Value self, Type type) { self.set_type(type); })
.def("first_use", &Value::first_use, return_value_policy::reference)
.def("has_one_use", &Value::HasOneUse)
.def("use_empty", &Value::use_empty)
.def("apply", &apply)
.def("is_same", &Value::operator==)
.def("hash", [](Value self) { return std::hash<Value>{}(self); })
.def("element_size",
[](Value self) { return phi::SizeOf(pir::GetValueDtype(self)); })
.def(
"stride",
[](Value self, py::object dim_obj = py::none()) {
const auto &dims = paddle::pybind::GetValueDims(self);
std::vector<int64_t> strides;
int64_t step = 1;
for (int i = static_cast<int>(dims.size()) - 1; i >= 0; --i) {
strides.insert(strides.begin(), step);
step *= dims[i];
}
if (dim_obj.is_none()) {
return py::cast(strides);
}
int dim = py::cast<int>(dim_obj);
dim = dim < 0 ? dim + static_cast<int>(dims.size()) : dim;
PADDLE_ENFORCE_EQ(dim >= 0 && dim < static_cast<int>(dims.size()),
true,
common::errors::InvalidArgument(
"Dimension out of range (expected to be in "
"range of [%d, %d], "
"but got %d)",
-static_cast<int>(dims.size()),
static_cast<int>(dims.size()) - 1,
dim));
return py::cast(strides[dim]);
},
py::arg("dim") = py::none())
.def("_rename", &name_analysis::RenameValue)
.def("_has_only_one_name",
[](Value self) -> bool {
return name_analysis::HasOnlyOneValueName(self);
})
.def("detach",
[](Value self) {
auto share_data_op =
ApiBuilder::Instance()
.GetBuilder()
->Build<paddle::dialect::ShareData_Op>(self);
auto out = share_data_op.out();
out.set_attribute(kAttrStopGradients,
BoolAttribute::get(IrContext::Instance(), true));
return out;
})
.def("__repr__", &Value2String)
.def("is_combine",
[](Value self) { return self.type().isa<pir::VectorType>(); })
.def("is_dist",
[](Value self) { return self.type().isa<DistTypeInterface>(); })
// The function will calculate the new local shape based on the global
// shape and the dist_attr argument.
.def("update_dist_attr",
[](Value &self, Attribute dist_attr) {
self.set_type(dialect::CvtToPirDistType(self.type(), dist_attr));
})
.def("is_coalesced",
[](Value self) {
auto sparse_coo_tensor_type =
self.type().dyn_cast<SparseCooTensorType>();
if (sparse_coo_tensor_type) {
return sparse_coo_tensor_type.coalesced();
} else {
PADDLE_THROW(common::errors::InvalidType(
"Method is_coalesced only support sparse coo tensor."));
}
})
.def_property_readonly(
"process_mesh",
[](Value &self) -> py::object {
auto type = self.type();
if (auto dist_type = type.dyn_cast<DistTypeInterface>()) {
return py::cast(dist_type.tensor_dist_attr()
.process_mesh_attr()
.process_mesh());
} else {
return py::cast<py::none>(Py_None);
}
})
.def("_clone",
[](Value self) {
// Return a new value owned by python side
return self;
})
.def("sparse_dim",
[](Value self) -> int32_t {
auto op_result = self.dyn_cast<OpResult>();
Operation *operation = op_result.owner();
if (self.type().isa<SparseCooTensorType>() &&
operation->name() == "pd_op.sparse_coo_tensor_sp") {
std::vector<Value> sources = operation->operands_source();
Value non_zero_indices = sources[1];
return phi::vectorize(GetValueDims(non_zero_indices))[0];
} else if (self.type().isa<SparseCsrTensorType>()) {
PADDLE_THROW(common::errors::InvalidType(
"SparseCsrTensor is unsupported in pir mode."));
} else {
return 0;
}
})
.def("dense_dim", [](Value self) -> int32_t {
auto op_result = self.dyn_cast<OpResult>();
Operation *operation = op_result.owner();
if (self.type().isa<SparseCooTensorType>() &&
operation->name() == "pd_op.sparse_coo_tensor_sp") {
std::vector<Value> sources = operation->operands_source();
Value non_zero_indices = sources[1];
int32_t dims = phi::vectorize(GetValueDims(self)).size();
return dims - phi::vectorize(GetValueDims(non_zero_indices))[0];
} else if (self.type().isa<SparseCsrTensorType>()) {
PADDLE_THROW(common::errors::InvalidType(
"SparseCsrTensor is unsupported in pir mode."));
} else {
return phi::vectorize(GetValueDims(self)).size();
}
});
}
void BindOpOperand(py::module *m) {
py::class_<OpOperand> op_operand(*m,
"OpOperand",
R"DOC(
OpOperand class represents the op_operand (input) of operation.
Notes:
The constructor of OpOperand should not be invoked directly. OpOperand can be automatically constructed
when build network.
)DOC");
op_operand.def("source", [](OpOperand &self) { return self.source(); })
.def("set_source",
[](OpOperand &self, Value *value) {
value ? self.set_source(*value) : self.set_source(nullptr);
})
.def("owner", &OpOperand::owner, return_value_policy::reference)
.def("index", &OpOperand::index);
}
bool GetValueBoolAttr(Value value, const std::string &attr_name) {
auto bool_attr = value.attribute<BoolAttribute>(attr_name);
return !bool_attr || bool_attr.data();
}
std::string GetAttrsMapJson(Operation *op) {
if (!op) {
PADDLE_THROW(common::errors::InvalidArgument(
"Operation pointer cannot be nullptr."));
}
auto attributes = op->attributes();
pir::ProgramWriter writer(1, false);
auto attrs_map_info = writer.GetAttributesMapJson(op->attributes()).dump();
return attrs_map_info;
}
pir::AttributeMap ConvertAttrsToAttributeMap(py::dict attrs) {
IrContext *ctx = IrContext::Instance();
pir::AttributeMap attrs_map;
for (auto item : attrs) {
std::string key = py::cast<std::string>(item.first);
py::handle value = item.second;
if (py::isinstance<py::bool_>(value)) {
attrs_map[key] = pir::BoolAttribute::get(ctx, py::cast<bool>(value));
} else if (py::isinstance<py::float_>(value)) {
attrs_map[key] = pir::FloatAttribute::get(ctx, py::cast<float>(value));
} else if (py::isinstance<py::str>(value)) {
attrs_map[key] =
pir::StrAttribute::get(ctx, py::cast<std::string>(value));
} else if (py::isinstance<py::list>(value)) {
py::list list_value = py::cast<py::list>(value);
std::vector<pir::Attribute> attr_list;
if (list_value.size() > 0) {
auto first_elem = list_value[0];
if (py::isinstance<py::bool_>(first_elem)) {
for (auto elem : list_value) {
attr_list.push_back(
pir::BoolAttribute::get(ctx, py::cast<bool>(elem)));
}
} else if (py::isinstance<py::str>(first_elem)) {
for (auto elem : list_value) {
attr_list.push_back(
pir::StrAttribute::get(ctx, py::cast<std::string>(elem)));
}
} else if (py::isinstance<py::int_>(first_elem)) {
for (auto elem : list_value) {
int64_t val = py::cast<int64_t>(elem);
attr_list.push_back(pir::Int64Attribute::get(ctx, val));
}
} else {
PADDLE_THROW(common::errors::InvalidArgument(
"Unsupported list element type, key: %s", key));
}
}
attrs_map[key] = pir::ArrayAttribute::get(ctx, attr_list);
} else {
PADDLE_THROW(common::errors::InvalidArgument(
"Unsupported attribute type, key: %s", key));
}
}
return attrs_map;
}
std::string GetAttrsMapJson(py::dict attrs) {
pir::AttributeMap attrs_map = ConvertAttrsToAttributeMap(attrs);
pir::ProgramWriter writer(1, false);
return writer.GetAttributesMapJson(attrs_map).dump();
}
std::string GetTypeJson(Operation *op, bool is_input) {
if (!op) {
PADDLE_THROW(
common::errors::InvalidArgument("Operation pointer cannot be nullptr"));
}
pir::ProgramWriter writer(1, false);
std::stringstream type_info_ss;
if (is_input) {
for (auto operand : op->operands_source()) {
type_info_ss << (writer.GetTypeJson(operand.type()).dump())
<< '\n'; // use '\n' as separator
}
} else {
for (auto result : op->results()) {
type_info_ss << (writer.GetTypeJson(result.type()).dump())
<< '\n'; // use '\n' as separator
}
}
return type_info_ss.str();
}
std::string GetInputsTypeJson(Operation *op) { return GetTypeJson(op, true); }
std::string GetOutputsTypeJson(Operation *op) { return GetTypeJson(op, false); }
void BindType(py::module *m) {
py::class_<Type> ir_type(*m, "Type");
ir_type.def("__eq__", &Type::operator==)
.def_property(
"shape",
[](Type self) { return phi::vectorize(GetTensorDims(self)); },
[](Type self, const std::vector<int> &shape) {
PADDLE_THROW(common::errors::InvalidArgument(
"can't set shape when building static graph"));
})
.def_property(
"dtype",
[](Type self) { return GetTensorDtype(self); },
[](Type self, DataType dtype) {
PADDLE_THROW(common::errors::InvalidArgument(
"can't set dtype when building static graph"));
})
.def_property(
"_local_shape",
[](Type self) {
if (!self.isa<DistDenseTensorType>()) {
PADDLE_THROW(common::errors::InvalidArgument(
"_local_shape is only for distdense tensor."));
}
return phi::vectorize(
self.dyn_cast<DistDenseTensorType>().local_ddim());
},
[](Type self, const std::vector<int> &shape) {
PADDLE_THROW(common::errors::InvalidArgument(
"can't set _local_shape when building static graph"));
})
.def("as_vec_type",
[](Type self) -> py::object {
if (auto vec_type = self.dyn_cast<VectorType>()) {
return py::cast(vec_type);
}
return py::cast<py::none>(Py_None);
})
.def("as_dist_type",
[](Type &self) -> py::object {
if (auto dist_type = self.dyn_cast<DistTypeInterface>()) {
return py::cast(dist_type);
}
return py::cast<py::none>(Py_None);
})
.def("__str__", [](Type &self) {
std::ostringstream print_stream;
print_stream << self;
return print_stream.str();
});
m->def("create_shaped_type",
[](Type &type, const std::vector<int64_t> &shape) -> Type {
if (type.isa<DenseTensorType>()) {
DenseTensorType src_type = type.dyn_cast<DenseTensorType>();
DenseTensorType dst_type =
DenseTensorType::get(IrContext::Instance(),
src_type.dtype(),
phi::make_ddim(shape),
src_type.data_layout(),
src_type.lod(),
src_type.offset());
return dst_type;
} else if (type.isa<SelectedRowsType>()) {
SelectedRowsType src_type = type.dyn_cast<SelectedRowsType>();
SelectedRowsType dst_type =
SelectedRowsType::get(IrContext::Instance(),
src_type.dtype(),
phi::make_ddim(shape),
src_type.data_layout(),
src_type.lod(),
src_type.offset());
return dst_type;
} else {
PADDLE_THROW(common::errors::InvalidArgument(
"Currently, we can only set shape for dense tensor"));
}
});
}
void BindVectorType(py::module *m) {
py::class_<VectorType, Type> vec_type(*m, "VectorType");
vec_type.def("as_list", &VectorType::data);
m->def("create_vec_type", [](std::vector<Type> &types) {
return VectorType::get(IrContext::Instance(), types);
});
}
void BindAttribute(py::module *m) {
py::class_<Attribute> ir_attr(*m, "Attribute", py::module_local());
ir_attr.def(py::init<>())
.def("__bool__", [](Attribute &self) { return static_cast<bool>(self); })
.def("__eq__", &Attribute::operator==)
.def("__str__",
[](Attribute &self) {
std::ostringstream print_stream;
print_stream << self;
return print_stream.str();
})
.def("as_tensor_dist_attr",
[](Attribute &self) -> py::object {
if (auto dist_attr = self.dyn_cast<TensorDistAttribute>()) {
return py::cast(dist_attr);
}
return py::cast<py::none>(Py_None);
})
.def("as_array_attr", [](Attribute &self) -> py::object {
if (auto array_attr = self.dyn_cast<ArrayAttribute>()) {
return py::cast(array_attr);
}
return py::cast<py::none>(Py_None);
});
py::class_<ArrayAttribute, Attribute> array_attr(*m, "ArrayAttribute");
array_attr.def("__len__", [](ArrayAttribute &self) { return self.size(); })
.def("__getitem__",
[](ArrayAttribute &self, int idx) { return self.at(idx); });
}
struct PyInsertionPoint {
pir::InsertionPoint value;
};
void BindInsertionPoint(pybind11::module *m) {
py::class_<PyInsertionPoint> ir_insertion_point(*m, "InsertionPoint", R"DOC(
InsertionPoint class represents the insertion point in the Builder.)DOC");
ir_insertion_point
.def(
"next",
[](PyInsertionPoint &self) -> Operation & {
if (self.value.second == self.value.first->end()) {
PADDLE_THROW(common::errors::InvalidArgument(
"The insertion point is already at the end and can't call "
"next()."));
}
return *(self.value.second++);
},
return_value_policy::reference)
.def(
"prev",
[](PyInsertionPoint &self) -> Operation & {
if (self.value.second == self.value.first->begin()) {
PADDLE_THROW(common::errors::InvalidArgument(
"The insertion point is already at the begin and can't call "
"prev()."));
}
return *(--self.value.second);
},
return_value_policy::reference)
.def(
"get_operation",
[](PyInsertionPoint &self) -> Operation & {
if (self.value.second == self.value.first->begin()) {
PADDLE_THROW(common::errors::InvalidArgument(
"The insertion point is already at the begin."));
} else if (self.value.second == self.value.first->end()) {
PADDLE_THROW(common::errors::InvalidArgument(
"The insertion point is already at the end."));
}
return *(self.value.second);
},
return_value_policy::reference)
.def(
"block",
[](PyInsertionPoint &self) { return self.value.first; },
return_value_policy::reference);
}
template <typename F, typename S>
void range_block_do(const Block *block,
std::pair<size_t, size_t> range,
F fn,
S skip_fn) {
auto [start, end] = range;
if (start >= end) {
return;
}
auto it = block->begin();
std::advance(it, start);
for (size_t i = start; i < end && it != block->end(); ++i, ++it) {
if (skip_fn(it)) {
continue;
}
fn(it);
}
}
template <typename F>
void range_block_do(const Block *block, std::pair<size_t, size_t> range, F fn) {
range_block_do(block, range, fn, [](Operation *op) { return false; });
}
std::map<int, int> GetOpInplaceInfo(const Operation *op) {
std::map<int, int> inplace_info;
if (!op->HasTrait<paddle::dialect::InplaceTrait>()) {
return inplace_info;
}
IrContext *ctx = IrContext::Instance();
std::string op_name = op->name();
if (op->attributes().count("op_name")) {
op_name =
op->attributes().at("op_name").dyn_cast<StrAttribute>().AsString();
}
pir::OpInfo op_info = ctx->GetRegisteredOpInfo(op_name);
paddle::dialect::OpYamlInfoParser yaml_parser(
op_info.GetInterfaceImpl<paddle::dialect::OpYamlInfoInterface>()
->get_op_info_(op_name),
paddle::dialect::IsLegacyOp(op_name));
for (size_t i = 0; i < op->num_results(); ++i) {
std::string value_name = yaml_parser.OutputNames()[i];
if (yaml_parser.HasInplace(value_name)) {
const std::string &inplace_name = yaml_parser.InplaceName(value_name);
inplace_info[i] = yaml_parser.InputName2Id().at(inplace_name);
}
if (yaml_parser.HasView(value_name)) {
const std::string &view_name = yaml_parser.ViewName(value_name);
inplace_info[i] = yaml_parser.InputName2Id().at(view_name);
}
}
return inplace_info;
}
std::pair<std::vector<Value>, std::unordered_set<Value>> AnalysisMiddleVariable(
const Program &program,
const std::vector<Value> &forward_inputs,
const std::vector<Value> &backward_outputs,
const std::pair<size_t, size_t> &forward_range,
const std::pair<size_t, size_t> &backward_range) {
std::vector<Value> middle_values;
std::unordered_set<Value> backward_used_values;
std::unordered_set<Value> x_or_param(forward_inputs.begin(),
forward_inputs.end());
for (const auto &value : backward_outputs) {
backward_used_values.insert(value);
}
range_block_do(
program.block(), backward_range, [&backward_used_values](Operation *op) {
pir::Walk(op, [&](Operation *inner_op) {
for (auto &t : inner_op->operands()) {
backward_used_values.insert(t.source());
}
});
});
range_block_do(
program.block(),
forward_range,
[&middle_values, &backward_used_values, &x_or_param](Operation *op) {
pir::Walk(op, [&](Operation *inner_op) {
for (auto &t : inner_op->results()) {
auto v = Value(t.Value::impl());
if (backward_used_values.count(v) && !x_or_param.count(v)) {
middle_values.push_back(v);
}
}
});
});
return std::make_pair(middle_values, backward_used_values);
}
void mapping_value(const std::vector<Value> &origin,
const std::unordered_map<Value, Value> &value_map,
std::vector<Value> &out) { // NOLINT
std::transform(origin.begin(),
origin.end(),
std::back_inserter(out),
[&value_map](const Value &v) {
if (v.impl() == nullptr) return Value(nullptr);
if (!value_map.count(v)) {
VLOG(2) << "mapping value found v is not exist. may not "
"used by backward program.";
return Value(nullptr);
}
return value_map.at(v);
});
}
using SplitedProgram = std::vector<std::shared_ptr<Program>>;
using SplitedAttribute = std::map<std::string, std::vector<Value>>;
using SplitedResult = std::pair<SplitedProgram, SplitedAttribute>;
static auto GetNoNeedBufferValue(const pir::Block *whole_block,
std::pair<size_t, size_t> range) {
// filter no need buffer values.
std::unordered_set<Value> need_buffer_values;
std::unordered_set<Value> no_need_buffer_values;
range_block_do(whole_block, range, [&need_buffer_values](Operation *op) {
// NOTE(SigureMo): We should process the CombineOp in it's users.
if (op->isa<pir::CombineOp>()) {
return;
}
if (op->HasInterface<paddle::dialect::OpYamlInfoInterface>() == false) {
// not a OpYamlInfoInterface, can't have no_need_buffer.
for (const auto &operand : op->operands_source()) {
need_buffer_values.insert(operand);
}
} else {
auto opinfo =
op->dyn_cast<paddle::dialect::OpYamlInfoInterface>().GetOpInfo();
int counter = 0;
for (const auto &op_input_info : std::get<0>(opinfo)) {
auto value = op->operand_source(counter);
if (!op_input_info.no_need_buffer) {
need_buffer_values.insert(value);
if (!IsFakeValue(value) && value.defining_op() &&
value.defining_op()->isa<pir::CombineOp>()) {
for (const auto &combine_value :
value.defining_op()->operands_source()) {
need_buffer_values.insert(combine_value);
}
}
}
counter += 1;
}
}
});
range_block_do(
whole_block,
range,
[&need_buffer_values, &no_need_buffer_values](const Operation *op) {
for (const auto &operand : op->operands_source()) {
if (need_buffer_values.count(operand) == 0) {
no_need_buffer_values.insert(operand);
}
}
});
return std::vector<Value>(no_need_buffer_values.begin(),
no_need_buffer_values.end());
}
using ValueMap = std::pair<std::vector<Value>, std::vector<Value>>;
std::pair<std::shared_ptr<Program>, ValueMap> CloneProgram(
const Program &program) {
// Limitation of this function:
// 1. don't support Parameters.
pir::IrMapping mapper;
auto cloned_program = program.Clone(mapper);
std::vector<Value> associated_array_key, associated_array_value;
for (auto &pair : mapper.GetMap<Value>()) {
associated_array_key.push_back(pair.first);
associated_array_value.push_back(pair.second);
}
return std::make_pair(
cloned_program,
std::make_pair(associated_array_key, associated_array_value));
}
void AppendPrintOp(Program *program,
const Value &value,
int first_n,
std::string message,
int summarize,
bool print_tensor_name,
bool print_tensor_type,
bool print_tensor_shape,
bool print_tensor_layout,
bool print_tensor_lod,
std::string print_phase,
bool is_forward,
int start_point) {
std::unordered_set<std::string> print_phase_set{
"FORWARD", "BACKWARD", "BOTH"};
if (!print_phase_set.count(print_phase)) {
PADDLE_THROW(common::errors::InvalidArgument(
"The attribute 'print_phase' must be one of 'FORWARD', 'BACKWARD', "
"'BOTH' but got '%s'.",
print_phase));
}
IrContext *ctx = IrContext::Instance();
auto op_info = ctx->GetRegisteredOpInfo(paddle::dialect::PrintOp::name());
pir::AttributeMap attribute_map = {
{"first_n", Int32Attribute::get(ctx, first_n)},
{"message", StrAttribute::get(ctx, message)},
{"summarize", Int32Attribute::get(ctx, summarize)},
{"print_tensor_name", BoolAttribute::get(ctx, print_tensor_name)},
{"print_tensor_type", BoolAttribute::get(ctx, print_tensor_type)},
{"print_tensor_shape", BoolAttribute::get(ctx, print_tensor_shape)},
{"print_tensor_layout", BoolAttribute::get(ctx, print_tensor_layout)},
{"print_tensor_lod", BoolAttribute::get(ctx, print_tensor_lod)},
{"print_phase", StrAttribute::get(ctx, print_phase)},
{"is_forward", BoolAttribute::get(ctx, is_forward)},
};
std::vector<pir::Type> output_types{value.type()};
Operation *operation =
Operation::Create({value}, attribute_map, output_types, op_info);
auto block = value.defining_op()->GetParent();
auto position = block->begin();
std::advance(position, start_point);
if (position == block->end()) {
block->push_back(operation);
} else {
block->insert(position, operation);
}
}
void AppendPrintOps(Program *program,
const std::vector<Value> &values,
int first_n,
std::string message,
int summarize,
bool print_tensor_name,
bool print_tensor_type,
bool print_tensor_shape,
bool print_tensor_layout,
bool print_tensor_lod,
std::string print_phase,
bool is_forward,
int start_point) {
int counter = 0;
std::unordered_set<Value> added_values;
for (const auto &value : values) {
if (!added_values.count(value)) {
AppendPrintOp(program,
value,
first_n,
message,
summarize,
print_tensor_name,
print_tensor_type,
print_tensor_shape,
print_tensor_layout,
print_tensor_lod,
print_phase,
is_forward,
start_point + counter);
++counter;
added_values.insert(value);
}
}
}
void AppendShadowOutput(Program *program,
const Value &value,
const std::string &name,
size_t start_point) {
IrContext *ctx = IrContext::Instance();
auto op_info = ctx->GetRegisteredOpInfo(pir::ShadowOutputOp::name());
pir::AttributeMap attribute_map = {
{"output_name", StrAttribute::get(ctx, name)},
};
Operation *operation = Operation::Create({value}, attribute_map, {}, op_info);
auto position = program->block()->begin();
std::advance(position, start_point);
if (position == program->block()->end()) {
program->block()->push_back(operation);
} else {
program->block()->insert(position, operation);
}
}
int AppendShadowOutputs(Program *program,
const std::vector<Value> &outputs,
int start_point,
std::string name_prefix) {
int counter = 0;
std::unordered_set<Value> added_value;
for (const auto &value : outputs) {
if (!added_value.count(value) || IsFakeValue(value)) {
std::string shadow_output_name =
name_analysis::TryGetValueFirstName(value).value_or(
name_prefix + std::to_string(counter));
AppendShadowOutput(
program, value, shadow_output_name, start_point + counter);
counter += 1;
added_value.insert(value);
}
}
// return the inserted op.
return counter;
}
SplitedResult SplitForwardBackward(
const Program &program,
const std::vector<Value> &forward_inputs,
const std::vector<Value> &forward_params,
const std::vector<Value> &forward_outputs,
const std::vector<Value> &forward_inputs_grads,
const std::vector<Value> &forward_params_grads,
const std::vector<Value> &forward_outputs_grads,
const std::pair<size_t, size_t> &forward_range,
const std::pair<size_t, size_t> &backward_range) {
std::vector<Value> forward_in_out_values;
for (auto &v :
std::vector({&forward_inputs, &forward_outputs, &forward_params})) {
forward_in_out_values.insert(
forward_in_out_values.end(), v->begin(), v->end());
}
std::vector<Value> backward_out_values;
for (auto &v : std::vector({&forward_inputs_grads, &forward_params_grads})) {
backward_out_values.insert(backward_out_values.end(), v->begin(), v->end());
}
std::vector<Value> fx, fp, fm, fo, bx, bp, bm, bo_g, bx_g, bp_g, bo;
std::vector<Value> no_need_buffer_values;
IrContext *ctx = IrContext::Instance();
auto forward_program = std::make_shared<Program>(ctx);
auto backward_program = std::make_shared<Program>(ctx);
std::vector<Value> middle_values;
std::unordered_set<Value> backward_used_values;
std::tie(middle_values, backward_used_values) =
AnalysisMiddleVariable(program,
forward_in_out_values,
backward_out_values,
forward_range,
backward_range);
pir::Block &backward_block = *backward_program->block();
bool has_backward = forward_inputs_grads.size() > 0 ||
forward_params_grads.size() > 0 ||
forward_outputs_grads.size() > 0;
// forward program construct.
VLOG(4) << "start create forward program.";
pir::IrMapping forward_mapper;
auto clone_options = pir::CloneOptions::All();
range_block_do(
program.block(),
forward_range,
[&forward_mapper, &forward_program, &clone_options](Operation *op) {
auto *cloned_op = op->Clone(forward_mapper, clone_options);
forward_program->block()->push_back(cloned_op);
},
// Skip the ShadowOutputOp.
/*skip_fn=*/[](Operation *op) { return op->isa<pir::ShadowOutputOp>(); });
auto &forward_value_map = forward_mapper.GetMutableMap<Value>();
// backward program construct.
// Step1. insert data op for inputs_values and middle_values
pir::IrMapping backward_mapper;
auto &backward_value_map = backward_mapper.GetMutableMap<Value>();
auto create_output_fn = [&ctx](
const std::unordered_map<Value, Value> &value_map,
const std::shared_ptr<Program> &program,
const std::string &prefix) {
auto counter = std::make_shared<size_t>(0);
return [&ctx, &value_map, &program, &prefix, counter](const Value &v) {
// NOTE(SigureMo): Ensure counter++ executed in each iteration.
auto default_name = prefix + std::to_string((*counter)++);
if (v.impl() == nullptr) {
return;
}
const Value &new_value = value_map.at(v);
std::string shadow_output_name =
name_analysis::TryGetValueFirstName(new_value).value_or(default_name);
auto op_info = ctx->GetRegisteredOpInfo(pir::ShadowOutputOp::name());
pir::AttributeMap attribute_map = {
{"output_name", StrAttribute::get(ctx, shadow_output_name)},
};
Operation *operation =
Operation::Create({new_value}, attribute_map, {}, op_info);
program->block()->push_back(operation);
};
};
VLOG(4) << "start create forward outputs, inserting shadow_output ops.";
std::for_each(
middle_values.begin(),
middle_values.end(),
create_output_fn(forward_value_map, forward_program, "middle_"));
std::for_each(
forward_outputs.begin(),
forward_outputs.end(),
create_output_fn(forward_value_map, forward_program, "output_"));
auto create_kwarg_fn = [&backward_block,
&backward_used_values,
&backward_value_map,
&forward_value_map](const std::string &prefix) {
auto counter = std::make_shared<size_t>(0);
return [&backward_block,
&backward_used_values,
&backward_value_map,
&forward_value_map,
&prefix,
counter](const Value &v) {
// NOTE(SigureMo): Ensure counter++ executed in each iteration.
auto default_name = prefix + std::to_string((*counter)++);
if (v && !backward_value_map.count(v) &&
(backward_used_values.count(v))) {
backward_value_map[v] = backward_block.AddKwarg(
name_analysis::TryGetValueFirstName(forward_value_map[v])
.value_or(default_name),
v.type());
}
};
};
if (has_backward) {
VLOG(4) << "start create backward inputs, creating keyword argument.";
VLOG(4) << "Create keyword argument for backward program: fo";
std::for_each(forward_outputs.begin(),
forward_outputs.end(),
create_kwarg_fn("output_"));
VLOG(4) << "Create keyword argument for backward program: fx";
std::for_each(forward_inputs.begin(),
forward_inputs.end(),
create_kwarg_fn("input_"));
VLOG(4) << "Create keyword argument for backward program: fp";
std::for_each(forward_params.begin(),
forward_params.end(),
create_kwarg_fn("param_"));
VLOG(4) << "Create keyword argument for backward program: fm";
std::for_each(
middle_values.begin(), middle_values.end(), create_kwarg_fn("middle_"));
VLOG(4) << "Create keyword argument for backward program: fo_g";
std::for_each(forward_outputs_grads.begin(),
forward_outputs_grads.end(),
create_kwarg_fn("output_grad_"));
VLOG(4) << "Create keyword argument for backward program end.";
}
// Step2. copy backward ops .
VLOG(4) << "start copy backward ops";
range_block_do(
program.block(),
backward_range,
[&backward_mapper, &backward_program, &clone_options](Operation *op) {
auto *cloned_op = op->Clone(backward_mapper, clone_options);
backward_program->block()->push_back(cloned_op);
},
// Skip the ShadowOutputOp.
/*skip_fn=*/[](Operation *op) { return op->isa<pir::ShadowOutputOp>(); });
VLOG(4) << "start create backward outputs, inserting shadow_output ops.";
if (has_backward) {
std::for_each(
forward_inputs_grads.begin(),
forward_inputs_grads.end(),
create_output_fn(backward_value_map, backward_program, "input_grad_"));
std::for_each(
forward_params_grads.begin(),
forward_params_grads.end(),
create_output_fn(backward_value_map, backward_program, "param_grad_"));
}
VLOG(4) << "forward_value_map.size() is " << forward_value_map.size();
VLOG(4) << "backward_value_map.size() is " << backward_value_map.size();
if (FLAGS_print_ir) {
std::ostringstream print_stream;
print_stream << "ForwardProgram is :\n";
forward_program->Print(print_stream);
print_stream << "BackwardProgram is:\n";
backward_program->Print(print_stream);
std::cout << "Splited Program (fwd | bwd): \n"
<< print_stream.str() << std::endl;
}
// construct all attributes we needed.
mapping_value(middle_values, forward_value_map, fm); // write 'fm'
mapping_value(middle_values, backward_value_map, bm); // write 'bm'
mapping_value(forward_inputs, forward_value_map, fx); // write 'fx'
mapping_value(forward_inputs, backward_value_map, bx); // write 'bx'
mapping_value(forward_params, forward_value_map, fp); // write 'fp'
mapping_value(forward_params, backward_value_map, bp); // write 'bp'
mapping_value(forward_outputs, forward_value_map, fo); // write 'fo'
mapping_value(
forward_inputs_grads, backward_value_map, bx_g); // write 'bx_g'
mapping_value(
forward_params_grads, backward_value_map, bp_g); // write 'bp_g'
mapping_value(
forward_outputs_grads, backward_value_map, bo_g); // write 'bo_g'
mapping_value(forward_outputs, backward_value_map, bo); // write 'bo'
mapping_value(GetNoNeedBufferValue(program.block(), backward_range),
forward_value_map,
no_need_buffer_values); // write 'no_need_buffers'
std::map<std::string, std::vector<Value>> attr = {
{"fx", fx},
{"fp", fp},
{"fm", fm},
{"fo", fo},
{"bx", bx},
{"bp", bp},
{"bm", bm},
{"bo_g", bo_g},
{"bx_g", bx_g},
{"bp_g", bp_g},
{"no_need_buffers", no_need_buffer_values},
{"bo", bo}};
std::vector<std::shared_ptr<Program>> programs = {forward_program,
backward_program};
return std::make_pair(programs, attr);
}
pir::Type CreateSelectedRowsTypeByDenseTensor(pir::Type dense_tensor_type) {
if (dense_tensor_type.isa<DenseTensorType>()) {
DenseTensorType type = dense_tensor_type.dyn_cast<DenseTensorType>();
return SelectedRowsType::get(IrContext::Instance(),
type.dtype(),
type.dims(),
type.data_layout(),
type.lod(),
type.offset());
} else {
PADDLE_THROW(common::errors::InvalidArgument(
"Currently, input is not a dense tensor type."));
}
}
pir::Type CreateDistDenseTensorTypeByDenseTensor(
const pir::Type &gdense_tensor_type,
const std::vector<int> &lshape,
const phi::distributed::ProcessMesh &mesh,
const std::vector<int64_t> &dims_mapping) {
if (gdense_tensor_type.isa<DenseTensorType>()) {
DenseTensorType type = gdense_tensor_type.dyn_cast<DenseTensorType>();
paddle::flat_hash_map<int64_t, phi::ReduceType> partial_status;
paddle::dialect::TensorDistAttribute tensor_dist_attr =
paddle::dialect::TensorDistAttribute::get(
IrContext::Instance(), mesh, dims_mapping, partial_status);
return DistDenseTensorType::get(
IrContext::Instance(), type, tensor_dist_attr, phi::make_ddim(lshape));
} else {
PADDLE_THROW(common::errors::InvalidArgument(
"Currently, input is not a dense tensor type are not supported."));
}
}
static void inline CreateVariableIfNotExist(
const std::vector<Value> &var_list,
framework::Scope *scope,
const framework::Executor *exe = nullptr) {
size_t len = var_list.size();
for (size_t i = 0; i < len; ++i) {
Value value = var_list[i];
std::string para_name = name_analysis::GetValueFirstName(value);
auto var = scope->FindVar(para_name);
if (var == nullptr) {
PADDLE_ENFORCE_NOT_NULL(exe,
common::errors::InvalidArgument(
"Parameter not Initialized, "
"Please set argument [executor] not None "
"or run startup program first"));
var = scope->Var(para_name);
auto *tensor_temp = var->GetMutable<DenseTensor>();
tensor_temp->Resize(
common::make_ddim(phi::vectorize(GetValueDims(value))));
phi::DeviceContextPool &pool = phi::DeviceContextPool::Instance();
const phi::DeviceContext *dev_ctx = nullptr;
dev_ctx = pool.Get(exe->GetPlace());
dev_ctx->Alloc(tensor_temp, pir::GetValueDtype(value));
}
}
return;
}
void BindUtils(pybind11::module *m) {
m->def("create_loaded_parameter", CreateVariableIfNotExist);
m->def("clone_program", CloneProgram);
m->def("get_op_inplace_info", GetOpInplaceInfo);
m->def("split_program", SplitForwardBackward);
m->def("append_shadow_outputs", AppendShadowOutputs);
m->def("append_shadow_output", AppendShadowOutput);
m->def("append_print", AppendPrintOp);
m->def("append_prints", AppendPrintOps);
m->def("fake_value", FakeValue);
m->def("get_fake_value_name",
[]() -> std::string { return paddle::framework::kFakeVarName; });
m->def("is_fake_value", IsFakeValue);
m->def("get_current_insertion_point", []() -> PyInsertionPoint {
return {ApiBuilder::Instance().GetCurrentInsertionPoint()};
});
m->def("set_insertion_point", [](const PyInsertionPoint &insertion_point) {
ApiBuilder::Instance().SetInsertionPoint(insertion_point.value);
});
m->def("set_insertion_point",
[](Operation *op) { ApiBuilder::Instance().SetInsertionPoint(op); });
m->def("set_insertion_point_after", [](Operation *op) {
ApiBuilder::Instance().SetInsertionPointAfter(op);
});
m->def("set_insertion_point_to_block_end", [](Block *block) {
ApiBuilder::Instance().SetInsertionPointToBlockEnd(block);
});
m->def("reset_insertion_point_to_start",
[]() { ApiBuilder::Instance().ResetInsertionPointToStart(); });
m->def("reset_insertion_point_to_end",
[]() { ApiBuilder::Instance().ResetInsertionPointToEnd(); });
m->def("set_chunk_id",
[](int chunk_id) { ApiBuilder::Instance().SetChunkId(chunk_id); });
m->def("get_chunk_id", []() { return ApiBuilder::Instance().GetChunkId(); });
m->def("set_op_role",
[](int op_role) { ApiBuilder::Instance().SetOpRole(op_role); });
m->def("get_op_role", []() { return ApiBuilder::Instance().GetOpRole(); });
m->def("set_comp_op_name", [](std::string comp_op_name) {
ApiBuilder::Instance().SetCompOpName(comp_op_name);
});
m->def("get_comp_op_name",
[]() { return ApiBuilder::Instance().GetCompOpName(); });
m->def("register_paddle_dialect", []() {
IrContext::Instance()
->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
});
m->def("register_dist_dialect", []() {
IrContext::Instance()->GetOrRegisterDialect<paddle::dialect::DistDialect>();
});
m->def("create_selected_rows_type_by_dense_tensor",
CreateSelectedRowsTypeByDenseTensor);
m->def("create_dist_dense_tensor_type_by_dense_tensor",
CreateDistDenseTensorTypeByDenseTensor);
m->def(
"translate_to_pir",
[](const ::paddle::framework::ProgramDesc &legacy_program) {
std::shared_ptr<Program> ret =
paddle::TranslateLegacyProgramToProgram(legacy_program);
return ret;
},
R"DOC(
Convert Fluid Program to New IR Program.
Args:
legacy_program (ProgramDesc): The Fluid Program that will be converted.
Returns:
Program: The New IR Program
Raises:
PreconditionNotMet: If legacy_program has multi block will raise error.
Examples:
.. code-block:: pycon
>>> import os
>>> # Paddle will remove this flag in the next version
>>> pir_flag = 'FLAGS_enable_pir_in_executor'
>>> os.environ[pir_flag] = 'True'
>>> import paddle
>>> from paddle import pir
>>> paddle.enable_static()
>>> x = paddle.randn([4, 4])
>>> main_program, start_program = (
... paddle.static.Program(),
... paddle.static.Program(),
... )
>>> with paddle.static.program_guard(main_program, start_program):
... x_s = paddle.static.data('x', [4, 4], x.dtype)
... x_s.stop_gradient = False
... y_s = paddle.matmul(x_s, x_s)
... z_s = paddle.add(y_s, y_s)
... k_s = paddle.tanh(z_s)
>>> pir_program = pir.translate_to_pir(main_program.desc)
>>> print(pir_program)
{
(%0) = "pd_op.data" () {dtype:(pd_op.DataType)float32,is_persistable:[false],name:"x",place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[4,4],stop_gradient:[false]} : () -> builtin.tensor<4x4xf32>
(%1) = "pd_op.matmul" (%0, %0) {is_persistable:[false],stop_gradient:[false],transpose_x:false,transpose_y:false} : (builtin.tensor<4x4xf32>, builtin.tensor<4x4xf32>) -> builtin.tensor<4x4xf32>
(%2) = "pd_op.add" (%1, %1) {is_persistable:[false],stop_gradient:[false]} : (builtin.tensor<4x4xf32>, builtin.tensor<4x4xf32>) -> builtin.tensor<4x4xf32>
(%3) = "pd_op.tanh" (%2) {is_persistable:[false],stop_gradient:[false]} : (builtin.tensor<4x4xf32>) -> builtin.tensor<4x4xf32>
}
)DOC");
m->def(
"check_unregistered_ops",
[](const framework::ProgramDesc &legacy_program) {
IrContext *ctx = IrContext::Instance();
return paddle::translator::CheckUnregisteredOperation(ctx,
legacy_program);
},
R"DOC(
Check unregistered operators in paddle dialect.
Args:
legacy_program (ProgramDesc): The Fluid Program that need checked.
Returns:
list[str] : List of unregistered operators in paddle dialect, the name is expressed by origin op name.
)DOC");
m->def(
"translate_to_pir_with_param_map",
[](const framework::ProgramDesc &legacy_program) {
auto ir_ctx = IrContext::Instance();
auto program = std::make_shared<pir::Program>(ir_ctx);
translator::ProgramTranslator program_translator(&legacy_program,
program.get());
program_translator.Translate();
return std::make_pair(program, program_translator.VarDesc2Value());
},
R"DOC(
Convert Fluid Program to New IR Program and get the mappings of VarDesc -> Value.
Args:
legacy_program (ProgramDesc): The Fluid Program that will be converted.
Returns:
Program: The New IR Program
dict[str, Value]: Mapping between VarDesc(by name) and Value.
Raises:
PreconditionNotMet: If legacy_program has multi block will raise error.
Examples:
.. code-block:: pycon
>>> import os
>>> # Paddle will remove this flag in the next version
>>> pir_flag = 'FLAGS_enable_pir_in_executor'
>>> os.environ[pir_flag] = 'True'
>>> import paddle
>>> from paddle import pir
>>> paddle.enable_static()
>>> x = paddle.randn([4, 4])
>>> main_program, start_program = (
... paddle.static.Program(),
... paddle.static.Program(),
... )
>>> with paddle.static.program_guard(main_program, start_program):
... x_s = paddle.static.data('x', [4, 4], x.dtype)
... x_s.stop_gradient = False
... y_s = paddle.matmul(x_s, x_s)
... z_s = paddle.add(y_s, y_s)
... k_s = paddle.tanh(z_s)
>>> pir_program, mappings = pir.translate_to_pir_with_param_map(main_program.desc)
>>> print(pir_program)
{
(%0) = "pd_op.data" () {dtype:(pd_op.DataType)float32,is_persistable:[false],name:"x",place:(pd_op.Place)Place(undefined:0),shape:(pd_op.IntArray)[4,4],stop_gradient:[false]} : () -> builtin.tensor<4x4xf32>
(%1) = "pd_op.matmul" (%0, %0) {is_persistable:[false],stop_gradient:[false],transpose_x:false,transpose_y:false} : (builtin.tensor<4x4xf32>, builtin.tensor<4x4xf32>) -> builtin.tensor<4x4xf32>
(%2) = "pd_op.add" (%1, %1) {is_persistable:[false],stop_gradient:[false]} : (builtin.tensor<4x4xf32>, builtin.tensor<4x4xf32>) -> builtin.tensor<4x4xf32>
(%3) = "pd_op.tanh" (%2) {is_persistable:[false],stop_gradient:[false]} : (builtin.tensor<4x4xf32>) -> builtin.tensor<4x4xf32>
}
>>> print(mappings)
{'matmul_v2_0.tmp_0': [Value(define_op_name=pd_op.matmul, index=0, dtype=builtin.tensor<4x4xf32>)], 'x': [Value(define_op_name=pd_op.data, index=0, dtype=builtin.tensor<4x4xf32>)], 'tanh_0.tmp_0': [Value(define_op_name=pd_op.tanh, index=0, dtype=builtin.tensor<4x4xf32>)], 'elementwise_add_0': [Value(define_op_name=pd_op.add, index=0, dtype=builtin.tensor<4x4xf32>)]}
)DOC");
m->def("clear_cinn_compilation_cache", []() {
#ifdef PADDLE_WITH_CINN
pybind11::gil_scoped_release release;
VLOG(4) << "clear CINN CompilationCache and free BackendResource.";
cinn::hlir::framework::CompilationCache::Instance().Clear();
#endif
});
m->def("cinn_compilation_cache_size", []() {
#ifdef PADDLE_WITH_CINN
pybind11::gil_scoped_release release;
VLOG(4) << "clear CINN CompilationCache and free BackendResource.";
return cinn::hlir::framework::CompilationCache::Instance().Size();
#endif
});
m->def("get_attrs_map_json",
py::overload_cast<Operation *>(&GetAttrsMapJson),
py::arg("op"));
m->def("get_attrs_map_json",
py::overload_cast<py::dict>(&GetAttrsMapJson),
py::arg("attrs"));
m->def("get_inputs_type_json",
&GetInputsTypeJson,
"Get operation input types as JSON string.");
m->def("get_outputs_type_json",
&GetOutputsTypeJson,
"Get operation output types as JSON string.");
}
namespace {
void ApplyCinnPass(Program &program) { // NOLINT
#ifdef PADDLE_WITH_CINN
auto CreatePassManager = [&]() -> std::shared_ptr<pir::PassManager> {
IrContext *ctx = IrContext::Instance();
ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
ctx->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
ctx->GetOrRegisterDialect<ap::dialect::OperatorDialect>();
ctx->GetOrRegisterDialect<pir::shape::ShapeDialect>();
auto pass_manager = std::make_shared<pir::PassManager>(ctx);
if (FLAGS_print_ir && VLOG_IS_ON(4)) {
pass_manager->EnableIRPrinting();
}
auto &shape_analysis = pir::ShapeAnalysisManager::Instance().Get(&program);
pass_manager->SetValueReplacedHook([&](Value from, Value to) {
shape_analysis.ShareShapeOrData(from, to);
});
return pass_manager;
};
cinn::dialect::ir::ApplyCinnPass(&program, CreatePassManager);
#else
PADDLE_THROW(common::errors::Unimplemented(
"Currently we only support CINN Pass for Pir under @to_static, please "
"compile PaddlePaddle with CINN"));
#endif
}
void ApplyPccPass(Program &program) { // NOLINT
#ifdef PADDLE_WITH_CINN
auto CreatePassManager = [&]() -> std::shared_ptr<pir::PassManager> {
IrContext *ctx = IrContext::Instance();
ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
ctx->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
ctx->GetOrRegisterDialect<ap::dialect::OperatorDialect>();
ctx->GetOrRegisterDialect<pir::shape::ShapeDialect>();
auto pass_manager = std::make_shared<pir::PassManager>(ctx);
if (FLAGS_print_ir && VLOG_IS_ON(4)) {
pass_manager->EnableIRPrinting();
}
auto &shape_analysis = pir::ShapeAnalysisManager::Instance().Get(&program);
pass_manager->SetValueReplacedHook([&](Value from, Value to) {
shape_analysis.ShareShapeOrData(from, to);
});
return pass_manager;
};
ap::paddle::ApplyPccPass(&program, CreatePassManager);
#else
PADDLE_THROW(common::errors::Unimplemented(
"Currently we only support CINN Pass for Pir under @to_static, please "
"compile PaddlePaddle with CINN"));
#endif
}
void CheckInferSymbolicIfNeed(Program &program) { // NOLINT
#ifdef PADDLE_WITH_CINN
auto CreatePassManager = [&]() -> std::shared_ptr<pir::PassManager> {
IrContext *ctx = IrContext::Instance();
ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
ctx->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
ctx->GetOrRegisterDialect<pir::shape::ShapeDialect>();
auto pass_manager = std::make_shared<pir::PassManager>(ctx);
if (FLAGS_print_ir) {
pass_manager->EnableIRPrinting();
}
return pass_manager;
};
cinn::dialect::ir::CheckInferSymbolicIfNeed(&program, CreatePassManager);
#else
// Do nothing.
#endif
}
} // namespace
void InferSymbolicShapePass(
std::shared_ptr<pir::PassManager> &pass_manager, // NOLINT
pir::Program &program) { // NOLINT
IrContext *ctx = IrContext::Instance();
ctx->GetOrRegisterDialect<pir::shape::ShapeDialect>();
pir::OriginalAttributesFilter::Instance().SetOriginalAttributesMap(
paddle::dialect::GetAllOpOriginalAttributes());
pass_manager->AddPass(pir::CreateShapeOptimizationPass());
}
std::shared_ptr<Program> ApplyCommonSubexpressionEliminationPass(
std::shared_ptr<Program> program) {
pir::PassManager pm(IrContext::Instance(), 2);
pm.AddPass(pir::CreateCommonSubexpressionEliminationPass());
pm.Run(program.get());
if (FLAGS_print_ir) {
std::cout
<< "IR After CommonSubexpressionEliminationPass -------------------"
<< std::endl;
std::cout << *program << std::endl;
}
return program;
}
void ApplyReduceAsToSumPass(
std::shared_ptr<pir::PassManager> &pass_manager, // NOLINT
pir::Program &program) { // NOLINT
#ifdef PADDLE_WITH_CINN
pass_manager->AddPass(cinn::dialect::ir::CreateReduceAsToSumPass());
pass_manager->AddPass(pir::CreateDeadCodeEliminationPass());
#else
PADDLE_THROW(common::errors::Unimplemented(
"Currently we only support ReduceAsToSumPass Pass for Pir under "
"@to_static, please "
"compile PaddlePaddle with CINN"));
#endif
}
std::shared_ptr<Program> ApplyFusedBnAddActPass(
std::shared_ptr<Program> program) {
pir::PassManager pm(IrContext::Instance(), 3);
pm.AddPass(pir::CreateFusedBnAddActPass());
pm.Run(program.get());
if (FLAGS_print_ir) {
std::cout << "IR After FusedBnAddActPass -------------------" << std::endl;
std::cout << *program << std::endl;
}
return program;
}
void BindIrPass(pybind11::module *m) {
m->def("apply_cinn_pass", ApplyCinnPass);
m->def("apply_pcc_pass", ApplyPccPass);
m->def("check_infer_symbolic_if_need", CheckInferSymbolicIfNeed);
m->def("infer_symbolic_shape_pass", InferSymbolicShapePass);
m->def("apply_cse_pass", ApplyCommonSubexpressionEliminationPass);
m->def("apply_bn_add_act_pass", ApplyFusedBnAddActPass);
m->def("reduce_as_sum_pass", ApplyReduceAsToSumPass);
py::class_<Pass, std::shared_ptr<Pass>> pass(*m,
"Pass",
R"DOC(
Pass class.
)DOC");
pass.def("name", &Pass::name)
.def("opt_level",
[](const Pass &self) { return self.pass_info().opt_level; })
.def("dependents",
[](const Pass &self) { return self.pass_info().dependents; });
}
void BindPassManager(pybind11::module *m) {
py::class_<PassManager, std::shared_ptr<PassManager>> pass_manager(
*m,
"PassManager",
R"DOC(
A class that manages all passes.
)DOC");
pass_manager
.def(py::init([](uint8_t opt_level) {
return std::make_unique<PassManager>(IrContext::Instance(),
opt_level);
}),
py::arg("opt_level") = 2)
.def("add_pass",
[](PassManager &self,
const std::string &pass_name,
const std::unordered_map<std::string, py::object> attrs = {}) {
auto pass = pir::PassRegistry::Instance().Get(pass_name);
for (const auto &attr : attrs) {
if (py::isinstance<py::str>(attr.second)) {
pass->Set(attr.first,
new std::string(attr.second.cast<std::string>()));
} else if (py::isinstance<py::bool_>(attr.second)) {
pass->Set(attr.first, new bool(attr.second.cast<bool>()));
} else if (py::isinstance<py::int_>(attr.second)) {
pass->Set(attr.first, new int(attr.second.cast<int>()));
} else if (py::isinstance<py::float_>(attr.second)) {
pass->Set(attr.first, new float(attr.second.cast<float>()));
} else if (py::isinstance<framework::Scope>(attr.second)) {
pass->SetNotOwned(attr.first,
attr.second.cast<framework::Scope *>());
} else if (py::isinstance<GPUPlace>(attr.second)) {
pass->Set(attr.first, new Place(attr.second.cast<GPUPlace>()));
} else {
PADDLE_THROW(common::errors::InvalidArgument(
"The pass attr is not supported this type."));
}
}
self.AddPass(std::move(pass));
})
.def("register_pass",
[](PassManager &self,
const std::string &pass_name,
std::shared_ptr<paddle::drr::DrrPatternContext> pattern_ctx) {
using AutoFinalPass =
paddle::drr::AutoDrrPass<paddle::drr::AutoDrrPattern>;
// Instead of using static PassRegistrar which may cause lifetime
// issues during program termination, directly register the pass to
// PassRegistry. This approach provides better control over object
// lifetime management and avoids potential segmentation faults
// during static destruction.
self.AddPass(
std::make_unique<AutoFinalPass>(pass_name, pattern_ctx));
})
.def("passes",
[](PassManager &self) {
std::vector<std::string> pass_names;
for (const auto &pass : self.passes()) {
pass_names.emplace_back(pass->name());
}
return pass_names;
})
.def("run", [](PassManager &self, Program *p) { self.Run(p); })
.def("empty", &PassManager::empty)
.def("clear", &PassManager::clear)
.def("enable_ir_printing",
[](PassManager &self) { self.EnableIRPrinting(); })
.def("enable_print_statistics",
[](PassManager &self) { self.EnablePrintStatistics(); });
}
void BindDrrPatternContext(pybind11::module *m) {
// bind NormalAttribute
pybind11::class_<drr::NormalAttribute> normal_attribute(*m,
"NormalAttribute");
// bind ComputeAttribute
pybind11::class_<drr::ComputeAttribute> compute_attribute(*m,
"ComputeAttribute");
// bind Tensor
pybind11::class_<drr::Tensor> tensor(*m,
"Tensor",
R"DOC(
register Tensor for DRR.
)DOC");
// bind Op
pybind11::class_<drr::Op> op(*m,
"Op",
R"DOC(
Represents an operation in the DRR framework.
)DOC");
op.def(
"__call__",
[](drr::Op *self,
const std::vector<drr::Tensor> &input_tensors,
const std::vector<drr::Tensor> &output_tensors) {
std::vector<const drr::Tensor *> input_ptrs;
std::vector<const drr::Tensor *> output_ptrs;
for (const auto &t : input_tensors) {
input_ptrs.push_back(&t);
}
for (const auto &t : output_tensors) {
output_ptrs.push_back(&t);
}
(*self)(input_ptrs, output_ptrs);
},
pybind11::arg("input_tensors"),
pybind11::arg("output_tensors"),
"Call the operation with an input tensor and return the output tensor.");
// bind DrrPatternContext
pybind11::class_<drr::DrrPatternContext,
std::shared_ptr<drr::DrrPatternContext>>
drr_pattern_context(*m,
"DrrPatternContext",
R"DOC(
A class that manages DRR (Dynamic Rewrite Rule) pattern context.
)DOC");
drr_pattern_context.def(pybind11::init<>())
.def("SourcePattern", &drr::DrrPatternContext::SourcePattern);
// bind drr::SourcePattern
pybind11::class_<drr::SourcePattern, std::shared_ptr<drr::SourcePattern>>
source_pattern(*m,
"SourcePattern",
R"DOC(
Represents a source pattern for matching in the DRR framework.
)DOC");
source_pattern.def("ResultPattern", &drr::SourcePattern::ResultPattern)
.def(
"Op",
[](drr::SourcePattern &self,
const std::string &op_type,
const std::unordered_map<std::string, drr::Attribute> &attributes =
{}) { return self.Op(op_type, attributes); },
pybind11::return_value_policy::reference_internal,
pybind11::arg("op_type"),
pybind11::arg("attributes") =
std::unordered_map<std::string, drr::Attribute>())
.def(
"Tensor",
[](drr::SourcePattern &self, const std::string &name) {
return self.Tensor(name);
},
pybind11::return_value_policy::reference_internal,
pybind11::arg("name"))
.def(
"InputNoneTensor",
[](drr::ResultPattern &self) { return self.InputNoneTensor(); },
pybind11::return_value_policy::reference_internal)
.def(
"OutputNoneTensor",
[](drr::ResultPattern &self) { return self.OutputNoneTensor(); },
pybind11::return_value_policy::reference_internal)
.def(
"Attr",
[](drr::SourcePattern &self, const std::string &attr_name) {
return self.Attr(attr_name);
},
pybind11::return_value_policy::reference_internal,
pybind11::arg("attr_name"))
.def("AddConstraint",
[](drr::SourcePattern &self, const pybind11::function &py_func) {
// wrap pyfunction -> cpp function
paddle::drr::ConstraintFunction cpp_func =
[py_func](const paddle::drr::MatchContext &context) -> bool {
try {
pybind11::object py_context = pybind11::cast(context);
pybind11::object result = py_func(py_context);
bool ret = result.cast<bool>();
return ret;
} catch (const pybind11::error_already_set &e) {
std::cerr << "Python error in AddConstraint callback: "
<< e.what() << std::endl;
throw;
}
};
self.AddConstraint(cpp_func);
})
.def("AddPostProcess",
[](drr::SourcePattern &self, const pybind11::function &py_func) {
// wrap pyfunction -> cpp function
paddle::drr::PostProcessFunction cpp_func =
[py_func](const paddle::drr::MatchContext &context) -> void {
try {
pybind11::object py_context = pybind11::cast(context);
py_func(py_context);
} catch (const pybind11::error_already_set &e) {
std::cerr << "Python error in AddPostProcess callback: "
<< e.what() << std::endl;
throw;
}
};
self.AddPostProcess(cpp_func);
});
// bind MatchContext
pybind11::class_<drr::MatchContext, std::shared_ptr<drr::MatchContext>>
match_context(*m,
"MatchContext",
R"DOC(
Represents the context of a match in the DRR framework.
)DOC");
match_context
.def(
"Tensor",
[](drr::MatchContext &self, std::string &tensor_name) -> Value {
return self.Tensor(tensor_name);
},
pybind11::return_value_policy::reference_internal,
pybind11::arg("tensor_name"))
// Attr
.def(
"StrAttr",
[](drr::MatchContext &self, const std::string &value_name) {
return self.Attr<std::string>(value_name);
},
pybind11::arg("value_name"))
.def(
"BoolAttr",
[](drr::MatchContext &self, const std::string &value_name) {
return self.Attr<bool>(value_name);
},
pybind11::arg("value_name"))
.def(
"Int32Attr",
[](drr::MatchContext &self, const std::string &value_name) {
return self.Attr<int32_t>(value_name);
},
pybind11::arg("value_name"))
.def(
"Int64Attr",
[](drr::MatchContext &self, const std::string &value_name) {
return self.Attr<int64_t>(value_name);
},
pybind11::arg("value_name"))
.def(
"Float32Attr",
[](drr::MatchContext &self, const std::string &value_name) {
return self.Attr<float>(value_name);
},
pybind11::arg("value_name"))
.def(
"DoubleAttr",
[](drr::MatchContext &self, const std::string &value_name) {
return self.Attr<double>(value_name);
},
pybind11::arg("value_name"))
.def(
"VectorInt32Attr",
[](drr::MatchContext &self, const std::string &value_name) {
return self.Attr<std::vector<int32_t>>(value_name);
},
pybind11::arg("value_name"))
.def(
"VectorInt64Attr",
[](drr::MatchContext &self, const std::string &value_name) {
return self.Attr<std::vector<int64_t>>(value_name);
},
pybind11::arg("value_name"))
.def(
"VectorFloat32Attr",
[](drr::MatchContext &self, const std::string &value_name) {
return self.Attr<std::vector<int32_t>>(value_name);
},
pybind11::arg("value_name"))
.def(
"DataTypeAttr",
[](drr::MatchContext &self, const std::string &value_name) {
return self.Attr<DataType>(value_name);
},
pybind11::arg("value_name"))
.def(
"PlaceAttr",
[](drr::MatchContext &self, const std::string &value_name) {
return self.Attr<Place>(value_name);
},
pybind11::arg("value_name"));
// bind drr::ResultPattern
pybind11::class_<drr::ResultPattern, std::shared_ptr<drr::ResultPattern>>
result_pattern(*m,
"ResultPattern",
R"DOC(
Represents a result pattern for matching in the DRR framework
)DOC");
result_pattern
.def(
"Op",
[](drr::ResultPattern &self,
const std::string &op_type,
const std::unordered_map<std::string, drr::Attribute> &attributes =
{}) { return self.Op(op_type, attributes); },
pybind11::return_value_policy::reference_internal,
pybind11::arg("op_type"),
pybind11::arg("attributes") =
std::unordered_map<std::string, drr::Attribute>())
.def(
"InputNoneTensor",
[](drr::ResultPattern &self) { return self.InputNoneTensor(); },
pybind11::return_value_policy::reference_internal)
.def(
"OutputNoneTensor",
[](drr::ResultPattern &self) { return self.OutputNoneTensor(); },
pybind11::return_value_policy::reference_internal)
.def(
"Tensor",
[](drr::ResultPattern &self, const std::string &name) {
return self.Tensor(name);
},
pybind11::return_value_policy::reference_internal,
pybind11::arg("name"))
// Attr
.def(
"StrAttr",
[](drr::ResultPattern &self, const std::string &value) {
return self.StrAttr(value);
},
pybind11::arg("value"))
.def(
"BoolAttr",
[](drr::ResultPattern &self, bool value) {
return self.BoolAttr(value);
},
pybind11::arg("value"))
.def(
"Int32Attr",
[](drr::ResultPattern &self, int32_t value) {
return self.Int32Attr(value);
},
pybind11::arg("value"))
.def(
"Int64Attr",
[](drr::ResultPattern &self, int64_t value) {
return self.Int64Attr(value);
},
pybind11::arg("value"))
.def(
"Float32Attr",
[](drr::ResultPattern &self, float value) {
return self.Float32Attr(value);
},
pybind11::arg("value"))
.def(
"DoubleAttr",
[](drr::ResultPattern &self, double value) {
return self.DoubleAttr(value);
},
pybind11::arg("value"))
.def(
"VectorInt32Attr",
[](drr::ResultPattern &self, const std::vector<int32_t> &value) {
return self.VectorInt32Attr(value);
},
pybind11::arg("value"))
.def(
"VectorInt64Attr",
[](drr::ResultPattern &self, const std::vector<int64_t> &value) {
return self.VectorInt64Attr(value);
},
pybind11::arg("value"))
.def(
"VectorFloat32Attr",
[](drr::ResultPattern &self, const std::vector<float> &value) {
return self.VectorFloatAttr(value);
},
pybind11::arg("value"))
.def(
"DataTypeAttr",
[](drr::ResultPattern &self, const std::string &value) {
return self.DataTypeAttr(value);
},
pybind11::arg("value"))
.def(
"PlaceAttr",
[](drr::ResultPattern &self, const std::string &value) {
return self.PlaceAttr(value);
},
pybind11::arg("value"))
.def(
"DataLayoutAttr",
[](drr::ResultPattern &self, const std::string &value) {
return self.DataLayoutAttr(value);
},
pybind11::arg("value"))
.def(
"ComputeAttr",
[](drr::ResultPattern &self, pybind11::function py_func) {
paddle::drr::AttrComputeFunc cpp_func =
[py_func](
const paddle::drr::MatchContext &context) -> std::any {
try {
pybind11::object py_context = pybind11::cast(context);
pybind11::object py_result = py_func(py_context);
pybind11::tuple result_tuple =
py_result.cast<pybind11::tuple>();
pybind11::object result = result_tuple[0];
std::string type_name = result_tuple[1].cast<std::string>();
auto any_result = CastPyObjectToAny(result, type_name);
return std::visit(
[](auto &&value) -> std::any { return std::any(value); },
any_result);
} catch (const pybind11::error_already_set &e) {
std::cerr << "Python error in ComputeAttr callback: "
<< e.what() << std::endl;
throw;
}
};
return self.ComputeAttr(cpp_func);
},
pybind11::arg("py_func"));
m->def("value_is_persistable",
[](const Value &value) { return pir::ValueIsPersistable(value); });
}
void BindShapeOrDataDimExprs(pybind11::module *m) {
py::class_<symbol::ShapeOrDataDimExprs,
std::shared_ptr<symbol::ShapeOrDataDimExprs>>
shape_or_data_dim_exprs(*m, "ShapeOrDataDimExprs", R"DOC(
A class that store the shape or data of value.
)DOC");
shape_or_data_dim_exprs
.def("shape",
&symbol::ShapeOrDataDimExprs::shape,
return_value_policy::reference)
.def("data",
&symbol::ShapeOrDataDimExprs::data,
return_value_policy::reference)
.def(
"is_equal",
[](symbol::ShapeOrDataDimExprs &self,
std::vector<int64_t> expect_shape,
std::vector<int64_t> expect_data) -> bool {
VLOG(3) << "Start compare shape and data.";
const auto &CompareFunc =
[&](const std::vector<int64_t> &expect,
const std::vector<symbol::DimExpr> &actual,
const std::string &compare_type) -> bool {
const auto PrintExpectAndActual = [&](const std::string &prefix) {
std::ostringstream sout;
sout << prefix << " expect: [";
std::copy(expect.begin(),
expect.end(),
std::ostream_iterator<int64_t>(sout, ","));
sout << "]" << std::endl;
sout << prefix << " actual:" << actual << std::endl;
LOG(ERROR) << sout.str();
};
if (actual.size() != expect.size()) {
LOG(ERROR) << compare_type << " expect size " << expect.size()
<< " is not equal to actual size " << actual.size()
<< " . The detailed information is as follows:";
PrintExpectAndActual(compare_type);
return false;
} else if (actual.empty()) {
return true;
}
for (size_t i = 0; i < actual.size(); i++) {
if (!actual.at(i).isa<int64_t>()) {
PrintExpectAndActual(compare_type);
PADDLE_THROW(common::errors::InvalidArgument(
"In OpTest, only supports cases where the type of "
"DimExpr "
"is int64_t."));
return false;
}
if (actual.at(i) != expect.at(i)) {
LOG(ERROR)
<< compare_type << " expect[" << i
<< "]: " << expect.at(i) << " is not equal to actual["
<< i << "]: " << actual.at(i)
<< " . The detailed information is as follows:";
PrintExpectAndActual(compare_type);
return false;
}
}
return true;
};
// compare shape
const std::vector<symbol::DimExpr> &actual_shape = self.shape();
const bool shape_status =
CompareFunc(expect_shape, actual_shape, "shape");
// compare data
const std::optional<std::vector<symbol::DimExpr>> &actual_data_ =
self.data();
if (actual_data_.has_value()) {
PADDLE_ENFORCE_LE(actual_shape.size(),
1,
common::errors::Unimplemented(
"Now data dim expr is not supported for "
"multi-dim shape."));
const std::vector<symbol::DimExpr> actual_data =
actual_data_.value();
const bool data_status =
CompareFunc(expect_data, actual_data, "data");
return shape_status && data_status;
}
return shape_status;
},
py::arg("expect_shape"),
py::arg("expect_data") = py::list());
}
void BindShapeConstraintIRAnalysis(pybind11::module *m) {
m->def(
"get_shape_constraint_ir_analysis",
[](const pir::Program *program) -> pir::ShapeConstraintIRAnalysis & {
return pir::ShapeAnalysisManager::Instance().Get(program);
},
return_value_policy::reference);
m->def("all_ops_defined_symbol_infer",
[](const pir::Program *program) -> bool {
// check that all ops have defined the InferSymbolicShapeInterface
bool flag = true;
for (Operation &op : *(program->block())) {
pir::InferSymbolicShapeInterface infer_interface =
op.dyn_cast<pir::InferSymbolicShapeInterface>();
if (!infer_interface) {
LOG(ERROR) << "The op: " << op.name()
<< " does not implement InferSymbolicShapeInterface.";
flag = false;
}
}
return flag;
});
#ifdef PADDLE_WITH_CINN
m->def(
"bind_symbolic_constraints",
[](pir::Program *program, const py::handle &constraints) -> void {
// Check input is sequence
PADDLE_ENFORCE_EQ(
py::isinstance<py::sequence>(constraints),
true,
common::errors::InvalidArgument(
"constraints for SOT symbolic variables must be a sequence."));
const py::sequence constraints_seq =
py::cast<py::sequence>(constraints);
if (py::len(constraints_seq) == 0) {
return;
}
// Process constraints
std::vector<std::tuple<std::string,
std::tuple<int64_t,
std::optional<int64_t>,
std::optional<int64_t>>>>
raw_constraints;
for (size_t idx = 0; idx < constraints_seq.size(); ++idx) {
const auto &constraint = constraints_seq[idx];
// Check constraint item is tuple
PADDLE_ENFORCE_EQ(
py::isinstance<py::tuple>(constraint),
true,
common::errors::InvalidArgument("Constraint[%zu] must be a tuple "
"of (name, dimension_triplet).",
idx));
const py::tuple constraint_tuple = py::cast<py::tuple>(constraint);
// Check tuple has 2 elements
PADDLE_ENFORCE_EQ(
constraint_tuple.size(),
2,
common::errors::InvalidArgument(
"Constraint[%zu] must have exactly 2 elements (got %zu).",
idx,
constraint_tuple.size()));
// Check and get input spec name
const py::handle name_handle = constraint_tuple[0];
PADDLE_ENFORCE_EQ(
py::isinstance<py::str>(name_handle),
true,
common::errors::InvalidArgument(
"Constraint[%zu][0] must be a string (got %s)",
idx,
py::str(name_handle.get_type()).cast<std::string>().c_str()));
const std::string input_spec_name =
py::cast<std::string>(name_handle);
// Check and get dimension triplet
const py::handle triplet_handle = constraint_tuple[1];
PADDLE_ENFORCE_EQ(py::isinstance<py::tuple>(triplet_handle),
true,
common::errors::InvalidArgument(
"Constraint[%zu][1] must be a tuple.", idx));
const py::tuple triplet = py::cast<py::tuple>(triplet_handle);
PADDLE_ENFORCE_EQ(
triplet.size(),
3,
common::errors::InvalidArgument(
"Constraint[%zu][1] must have 3 elements (got %zu).",
idx,
triplet.size()));
// Validate and convert elements
auto convert_optional = [idx](const py::handle &h,
int pos) -> std::optional<int64_t> {
if (h.is_none()) return std::nullopt;
PADDLE_ENFORCE_EQ(
py::isinstance<py::int_>(h),
true,
"Constraint[%zu][1][%d] must be int or None (got %s).",
idx,
pos,
py::str(h.get_type()).cast<std::string>().c_str());
return py::cast<int64_t>(h);
};
// Check dim_idx
PADDLE_ENFORCE_EQ(
py::isinstance<py::int_>(triplet[0]),
true,
common::errors::InvalidArgument(
"Constraint[%zu][1][0] (dim_idx) must be int (got %s).",
idx,
py::str(triplet[0].get_type()).cast<std::string>().c_str()));
const int64_t dim_idx = py::cast<int64_t>(triplet[0]);
// Convert min/max with position info
std::optional<int64_t> min_val = convert_optional(triplet[1], 1);
std::optional<int64_t> max_val = convert_optional(triplet[2], 2);
// Add to constraints
raw_constraints.emplace_back(
std::move(input_spec_name),
std::make_tuple(dim_idx, min_val, max_val));
}
::cinn::dialect::ir::SpecifyInputDynamicDimFromPython(program,
raw_constraints);
},
py::arg("program"),
py::arg("constraints").noconvert());
#endif
py::class_<pir::ShapeConstraintIRAnalysis,
std::shared_ptr<pir::ShapeConstraintIRAnalysis>>
shape_constraint_ir_analysis(*m, "ShapeConstraintIRAnalysis", R"DOC(
A class that store the shape information of all operators.
)DOC");
shape_constraint_ir_analysis
.def("get_shape_or_data_for_var",
&pir::ShapeConstraintIRAnalysis::GetShapeOrDataForValue,
return_value_policy::reference)
.def("set_shape_or_data_for_var",
&pir::ShapeConstraintIRAnalysis::SetShapeOrDataForValue)
.def("register_symbol_cstr_from_shape_analysis",
&pir::ShapeConstraintIRAnalysis::
RegisterSymbolConstraintFromShapeAnalysis);
}
void BindPir(pybind11::module *module) {
auto ir_module = module->def_submodule("pir");
BindProgram(&ir_module);
BindBlock(&ir_module);
BindValue(&ir_module);
BindIrMapping(&ir_module);
BindCloneOptions(&ir_module);
BindOperation(&ir_module);
BindOpOperand(&ir_module);
BindType(&ir_module);
BindVectorType(&ir_module);
BindAttribute(&ir_module);
BindInsertionPoint(&ir_module);
BindUtils(&ir_module);
BindIrPass(&ir_module);
BindPassManager(&ir_module);
BindControlFlowApi(&ir_module);
BindShapeOrDataDimExprs(&ir_module);
BindShapeConstraintIRAnalysis(&ir_module);
auto ops_modules = ir_module.def_submodule("ops");
BindOpsAPI(&ops_modules);
BindDrrPatternContext(&ir_module);
}
} // namespace pybind
} // namespace paddle
Reference in New Issue View Git Blame Copy Permalink