wehub-resource-sync
417 lines
17 KiB
C++
417 lines
17 KiB
C++
/* Copyright (c) 2022 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/jit.h"
|
|
#include "glog/logging.h"
|
|
#include "paddle/fluid/framework/variable.h"
|
|
#include "paddle/fluid/imperative/layer.h"
|
|
#include "paddle/fluid/jit/function.h"
|
|
#include "paddle/fluid/jit/function_schema.h"
|
|
#include "paddle/fluid/jit/layer.h"
|
|
#include "paddle/fluid/jit/serializer.h"
|
|
#include "paddle/fluid/pybind/sot/eval_frame.h"
|
|
#include "paddle/fluid/pybind/sot/eval_frame_tools.h"
|
|
#include "paddle/fluid/pybind/sot/frame_proxy.h"
|
|
#include "paddle/fluid/pybind/sot/guards.h"
|
|
#include "paddle/fluid/pybind/sot/macros.h"
|
|
#include "paddle/phi/common/data_type.h"
|
|
#include "paddle/phi/common/place.h"
|
|
#include "paddle/utils/pybind.h"
|
|
#include "pybind11/pybind11.h"
|
|
#include "pybind11/pytypes.h"
|
|
|
|
namespace py = pybind11;
|
|
|
|
namespace paddle::pybind {
|
|
|
|
PyTypeObject *g_jit_function_pytype = nullptr;
|
|
using Variable = paddle::framework::Variable;
|
|
|
|
void BindJit(pybind11::module *m) {
|
|
py::class_<jit::Layer>(*m, "Layer", R"DOC(Layer Class.)DOC")
|
|
.def("function_names", &jit::Layer::FunctionNames)
|
|
.def("function", &jit::Layer::Function)
|
|
.def("function_info", &jit::Layer::FunctionInfo);
|
|
|
|
py::class_<jit::Function, std::shared_ptr<jit::Function>> function(
|
|
*m, "Function", R"DOC(Function Class.)DOC");
|
|
g_jit_function_pytype = reinterpret_cast<PyTypeObject *>(function.ptr());
|
|
|
|
py::class_<jit::BaseFunctionInfo, std::shared_ptr<jit::BaseFunctionInfo>>(
|
|
*m, "FunctionInfo", R"DOC(BaseFunctionInfo Class.)DOC")
|
|
.def("name", &jit::BaseFunctionInfo::FunctionName)
|
|
.def("input_names", &jit::BaseFunctionInfo::InputArgNames)
|
|
.def("output_names", &jit::BaseFunctionInfo::OutputArgNames);
|
|
|
|
m->def("Load", [](const std::string &path, const CPUPlace &cpu_place) {
|
|
return paddle::jit::Load(path, cpu_place);
|
|
});
|
|
|
|
m->def("Load", [](const std::string &path, const GPUPlace &cuda_place) {
|
|
return paddle::jit::Load(path, cuda_place);
|
|
});
|
|
}
|
|
|
|
void BindGuard(pybind11::module *m) {
|
|
#if SOT_IS_SUPPORTED
|
|
py::class_<GuardBase, std::shared_ptr<GuardBase>>(
|
|
*m, "GuardBase", R"DOC(GuardBase Class.)DOC")
|
|
.def("check", &GuardBase::check_pybind);
|
|
py::class_<LambdaGuard, GuardBase, std::shared_ptr<LambdaGuard>>(
|
|
*m, "LambdaGuard", R"DOC(LambdaGuard Class.)DOC")
|
|
.def(py::init<const py::function &>(), py::arg("guard_check_fn"));
|
|
py::class_<GuardGroup, GuardBase, std::shared_ptr<GuardGroup>>(
|
|
*m, "GuardGroup", R"DOC(GuardGroup Class.)DOC")
|
|
.def(py::init<const std::vector<std::shared_ptr<GuardBase>> &>(),
|
|
py::arg("guards"));
|
|
py::class_<TypeMatchGuard, GuardBase, std::shared_ptr<TypeMatchGuard>>(
|
|
*m, "TypeMatchGuard", R"DOC(TypeMatchGuard Class.)DOC")
|
|
.def(py::init<const py::type &>(), py::arg("py_type"));
|
|
py::class_<IdMatchGuard, GuardBase, std::shared_ptr<IdMatchGuard>>(
|
|
*m, "IdMatchGuard", R"DOC(IdMatchGuard Class.)DOC")
|
|
.def(py::init<const py::object &>(), py::arg("py_obj"));
|
|
py::class_<LengthMatchGuard, GuardBase, std::shared_ptr<LengthMatchGuard>>(
|
|
*m, "LengthMatchGuard", R"DOC(LengthMatchGuard Class.)DOC")
|
|
.def(py::init<const Py_ssize_t &>(), py::arg("length"));
|
|
py::class_<ValueMatchGuard, GuardBase, std::shared_ptr<ValueMatchGuard>>(
|
|
*m, "ValueMatchGuard", R"DOC(ValueMatchGuard Class.)DOC")
|
|
.def(py::init<const py::object &>(), py::arg("py_value"));
|
|
py::class_<DtypeMatchGuard, GuardBase, std::shared_ptr<DtypeMatchGuard>>(
|
|
*m, "DtypeMatchGuard", R"DOC(DtypeMatchGuard Class.)DOC")
|
|
.def(py::init<const paddle::framework::proto::VarType &>(),
|
|
py::arg("dtype"))
|
|
.def(py::init<const DataType &>(), py::arg("dtype"));
|
|
py::class_<AttributeMatchGuard,
|
|
GuardBase,
|
|
std::shared_ptr<AttributeMatchGuard>>(
|
|
*m, "AttributeMatchGuard", R"DOC(AttributeMatchGuard Class.)DOC")
|
|
.def(py::init<const py::object &, const std::string &>(),
|
|
py::arg("obj"),
|
|
py::arg("attr_name"));
|
|
py::class_<ShapeMatchGuard, GuardBase, std::shared_ptr<ShapeMatchGuard>>(
|
|
*m, "ShapeMatchGuard", R"DOC(ShapeMatchGuard Class.)DOC")
|
|
.def(py::init<const std::vector<py::object> &, int64_t>(),
|
|
py::arg("shape"),
|
|
py::arg("min_non_specialized_number"));
|
|
py::class_<LayerMatchGuard, GuardBase, std::shared_ptr<LayerMatchGuard>>(
|
|
*m, "LayerMatchGuard", R"DOC(LayerMatchGuard Class.)DOC")
|
|
.def(py::init<const py::object &>(), py::arg("layer_obj"));
|
|
py::class_<InstanceCheckGuard,
|
|
GuardBase,
|
|
std::shared_ptr<InstanceCheckGuard>>(
|
|
*m, "InstanceCheckGuard", R"DOC(InstanceCheckGuard Class.)DOC")
|
|
.def(py::init<const py::object &>(), py::arg("isinstance_obj"));
|
|
py::class_<NumPyDtypeMatchGuard,
|
|
GuardBase,
|
|
std::shared_ptr<NumPyDtypeMatchGuard>>(
|
|
*m, "NumPyDtypeMatchGuard", R"DOC(NumPyDtypeMatchGuard Class.)DOC")
|
|
.def(py::init<const py::object &>(), py::arg("dtype"));
|
|
py::class_<NumPyArrayValueMatchGuard,
|
|
GuardBase,
|
|
std::shared_ptr<NumPyArrayValueMatchGuard>>(
|
|
*m,
|
|
"NumPyArrayValueMatchGuard",
|
|
R"DOC(NumPyArrayValueMatchGuard Class.)DOC")
|
|
.def(py::init<const py::object &>(), py::arg("array"));
|
|
py::class_<NumPyArrayShapeMatchGuard,
|
|
GuardBase,
|
|
std::shared_ptr<NumPyArrayShapeMatchGuard>>(
|
|
*m,
|
|
"NumPyArrayShapeMatchGuard",
|
|
R"DOC(NumPyArrayShapeMatchGuard Class.)DOC")
|
|
.def(py::init<const std::vector<py::object> &, int64_t>(),
|
|
py::arg("shape"),
|
|
py::arg("min_non_specialized_number"));
|
|
py::class_<WeakRefMatchGuard, GuardBase, std::shared_ptr<WeakRefMatchGuard>>(
|
|
*m, "WeakRefMatchGuard", R"DOC(WeakRefMatchGuard Class.)DOC")
|
|
.def(py::init<const py::object &>(), py::arg("func"));
|
|
py::class_<IsNotDenseTensorHoldAllocationMatchGuard,
|
|
GuardBase,
|
|
std::shared_ptr<IsNotDenseTensorHoldAllocationMatchGuard>>(
|
|
*m,
|
|
"IsNotDenseTensorHoldAllocationMatchGuard",
|
|
R"DOC(IsNotDenseTensorHoldAllocationMatchGuard Class.)DOC")
|
|
.def(py::init());
|
|
|
|
m->def(
|
|
"merge_guard",
|
|
[](const std::vector<std::shared_ptr<GuardBase>> &py_guards) {
|
|
return GuardGroup(py_guards);
|
|
},
|
|
py::arg("py_guards"));
|
|
#endif
|
|
}
|
|
|
|
void BindGuardTree(pybind11::module *m) {
|
|
#if SOT_IS_SUPPORTED
|
|
py::class_<GuardTree, std::shared_ptr<GuardTree>>(
|
|
*m, "GuardTree", R"DOC(GuardTree Class.)DOC")
|
|
.def(py::init<const std::vector<
|
|
std::vector<std::shared_ptr<GuardNodeBase>>> &>(),
|
|
py::arg("guard_chain_list"))
|
|
.def_property_readonly(
|
|
"guard_nodes", [](GuardTree &self) { return self.get_guard_nodes(); })
|
|
.def(
|
|
"lookup",
|
|
[](GuardTree &self, py::object frame) {
|
|
return self.lookup(reinterpret_cast<FrameProxy *>(frame.ptr()));
|
|
},
|
|
py::arg("frame"))
|
|
.def(
|
|
"add_guard_chain",
|
|
[](GuardTree &self,
|
|
const std::vector<std::shared_ptr<GuardNodeBase>> &guard_chain) {
|
|
self.add_guard_chain(guard_chain);
|
|
},
|
|
py::arg("guard_chain"))
|
|
.def("stringify", &GuardTree::stringify);
|
|
|
|
py::class_<GuardNodeBase, std::shared_ptr<GuardNodeBase>>(
|
|
*m, "GuardNodeBase", R"DOC(GuardNodeBase Class.)DOC")
|
|
.def_property(
|
|
"return_cache_index",
|
|
[](GuardNodeBase &self) { return self.return_cache_index; },
|
|
[](GuardNodeBase &self, int index) {
|
|
self.return_cache_index = index;
|
|
})
|
|
.def_property(
|
|
"next_guard_nodes",
|
|
[](GuardNodeBase &self) { return self.next_guard_nodes; },
|
|
[](GuardNodeBase &self,
|
|
const std::vector<std::shared_ptr<GuardNodeBase>>
|
|
&next_guard_nodes) {
|
|
self.next_guard_nodes = next_guard_nodes;
|
|
})
|
|
.def(
|
|
"lookup",
|
|
[](GuardNodeBase &self, py::object frame) {
|
|
return self.lookup(reinterpret_cast<FrameProxy *>(frame.ptr()));
|
|
},
|
|
py::arg("frame"));
|
|
|
|
py::class_<CheckGuardNode<0>,
|
|
GuardNodeBase,
|
|
std::shared_ptr<CheckGuardNode<0>>>(*m, "CheckGuardNode0")
|
|
.def_property_readonly("exprs",
|
|
[](CheckGuardNode<0> &self) { return self.exprs; })
|
|
.def("get_guard_name", &CheckGuardNode<0>::get_guard_name);
|
|
|
|
py::class_<CheckGuardNode<1>,
|
|
GuardNodeBase,
|
|
std::shared_ptr<CheckGuardNode<1>>>(*m, "CheckGuardNode1")
|
|
.def_property_readonly("exprs",
|
|
[](CheckGuardNode<1> &self) { return self.exprs; })
|
|
.def("get_guard_name", &CheckGuardNode<1>::get_guard_name);
|
|
|
|
py::class_<CheckGuardNode<2>,
|
|
GuardNodeBase,
|
|
std::shared_ptr<CheckGuardNode<2>>>(*m, "CheckGuardNode2")
|
|
.def_property_readonly("exprs",
|
|
[](CheckGuardNode<2> &self) { return self.exprs; })
|
|
.def("get_guard_name", &CheckGuardNode<2>::get_guard_name);
|
|
|
|
py::class_<IsGradEnabledGuardNode,
|
|
CheckGuardNode<0>,
|
|
std::shared_ptr<IsGradEnabledGuardNode>>(
|
|
*m, "IsGradEnabledGuardNode", R"DOC(IsGradEnabledGuardNode Class.)DOC")
|
|
.def(py::init<bool,
|
|
const std::vector<std::shared_ptr<GuardNodeBase>> &,
|
|
const std::optional<int> &>(),
|
|
py::arg("is_grad_enabled"),
|
|
py::arg("next_guard_nodes") = py::list(),
|
|
py::arg("return_cache_index") = py::none());
|
|
|
|
py::class_<LegacyGuardNode,
|
|
CheckGuardNode<1>,
|
|
std::shared_ptr<LegacyGuardNode>>(
|
|
*m, "GuardNode", R"DOC(GuardNode Class.)DOC")
|
|
.def(py::init<const std::shared_ptr<GuardBase> &,
|
|
const std::array<std::shared_ptr<ExprNodeBase>, 1> &,
|
|
const std::vector<std::shared_ptr<GuardNodeBase>> &,
|
|
const std::optional<int> &>(),
|
|
py::arg("guard"),
|
|
py::arg("exprs"),
|
|
py::arg("next_guard_nodes") = py::list(),
|
|
py::arg("return_cache_index") = py::none())
|
|
.def_property_readonly("guard",
|
|
[](LegacyGuardNode &self) { return self.guard; });
|
|
|
|
py::class_<TensorDistMetaMatchGuardNode,
|
|
CheckGuardNode<2>,
|
|
std::shared_ptr<TensorDistMetaMatchGuardNode>>(
|
|
*m,
|
|
"TensorDistMetaMatchGuardNode",
|
|
R"DOC(TensorDistMetaMatchGuardNode Class.)DOC")
|
|
.def(py::init<const py::object &,
|
|
const std::array<std::shared_ptr<ExprNodeBase>, 2> &,
|
|
const std::vector<std::shared_ptr<GuardNodeBase>> &,
|
|
const std::optional<int> &>(),
|
|
py::arg("dist_info"),
|
|
py::arg("exprs"),
|
|
py::arg("next_guard_nodes") = py::list(),
|
|
py::arg("return_cache_index") = py::none());
|
|
|
|
py::class_<ExprGuardNode, GuardNodeBase, std::shared_ptr<ExprGuardNode>>(
|
|
*m, "ExprGuardNode", R"DOC(ExprGuardNode Class.)DOC")
|
|
.def(py::init<const std::shared_ptr<ExprNodeBase> &,
|
|
const std::vector<std::shared_ptr<GuardNodeBase>> &,
|
|
const std::optional<int> &>(),
|
|
py::arg("expr"),
|
|
py::arg("next_guard_nodes") = py::list(),
|
|
py::arg("return_cache_index") = py::none());
|
|
|
|
py::class_<DummyGuardNode, GuardNodeBase, std::shared_ptr<DummyGuardNode>>(
|
|
*m, "DummyGuardNode", R"DOC(DummyGuardNode Class.)DOC")
|
|
.def(py::init<bool,
|
|
const std::vector<std::shared_ptr<GuardNodeBase>> &,
|
|
const std::optional<int> &>(),
|
|
py::arg("return_true") = true,
|
|
py::arg("next_guard_nodes") = py::list(),
|
|
py::arg("return_cache_index") = py::none());
|
|
|
|
py::class_<ExprNodeBase, std::shared_ptr<ExprNodeBase>>(
|
|
*m, "ExprNodeBase", R"DOC(ExprNodeBase Class.)DOC")
|
|
.def(
|
|
"eval",
|
|
[](ExprNodeBase &self, py::object frame) {
|
|
return self.eval(reinterpret_cast<FrameProxy *>(frame.ptr()));
|
|
},
|
|
py::arg("frame"))
|
|
.def("stringify", &ExprNodeBase::stringify, py::arg("indent") = 0);
|
|
|
|
py::class_<ConstantExprNode, ExprNodeBase, std::shared_ptr<ConstantExprNode>>(
|
|
*m, "ConstantExprNode", R"DOC(ConstantExprNode Class.)DOC")
|
|
.def(py::init<const py::object &>(), py::arg("value_ptr"));
|
|
|
|
py::class_<ExternVarExprNode,
|
|
ExprNodeBase,
|
|
std::shared_ptr<ExternVarExprNode>>(
|
|
*m, "ExternVarExprNode", R"DOC(ExternVarExprNode Class.)DOC")
|
|
.def(py::init<const std::string &, const py::object &>(),
|
|
py::arg("var_name"),
|
|
py::arg("value_ptr"));
|
|
|
|
py::class_<LocalVarExprNode, ExprNodeBase, std::shared_ptr<LocalVarExprNode>>(
|
|
*m, "LocalVarExprNode", R"DOC(LocalVarExprNode Class.)DOC")
|
|
.def(py::init<const std::string &>(), py::arg("var_name"));
|
|
|
|
py::class_<GlobalVarExprNode,
|
|
ExprNodeBase,
|
|
std::shared_ptr<GlobalVarExprNode>>(
|
|
*m, "GlobalVarExprNode", R"DOC(GlobalVarExprNode Class.)DOC")
|
|
.def(py::init<const std::string &>(), py::arg("var_name"));
|
|
|
|
py::class_<AttributeExprNode,
|
|
ExprNodeBase,
|
|
std::shared_ptr<AttributeExprNode>>(
|
|
*m, "AttributeExprNode", R"DOC(AttributeExprNode Class.)DOC")
|
|
.def(py::init<std::shared_ptr<ExprNodeBase>, const std::string &>(),
|
|
py::arg("var_expr"),
|
|
py::arg("attr_name"));
|
|
|
|
py::class_<ItemExprNode, ExprNodeBase, std::shared_ptr<ItemExprNode>>(
|
|
*m, "ItemExprNode", R"DOC(ItemExprNode Class.)DOC")
|
|
.def(py::init<std::shared_ptr<ExprNodeBase>,
|
|
std::shared_ptr<ExprNodeBase>>(),
|
|
py::arg("var_expr"),
|
|
py::arg("key_expr"));
|
|
|
|
py::class_<BinaryExprNode, ExprNodeBase, std::shared_ptr<BinaryExprNode>>(
|
|
*m, "BinaryExprNode", R"DOC(BinaryExprNode Class.)DOC")
|
|
.def(py::init<std::shared_ptr<ExprNodeBase>,
|
|
std::shared_ptr<ExprNodeBase>,
|
|
const std::string &>(),
|
|
py::arg("lhs"),
|
|
py::arg("rhs"),
|
|
py::arg("op"));
|
|
|
|
py::class_<UnaryExprNode, ExprNodeBase, std::shared_ptr<UnaryExprNode>>(
|
|
*m, "UnaryExprNode", R"DOC(UnaryExprNode Class.)DOC")
|
|
.def(py::init<std::shared_ptr<ExprNodeBase>, const std::string &>(),
|
|
py::arg("expr"),
|
|
py::arg("op"));
|
|
#endif
|
|
}
|
|
|
|
void BindSot(pybind11::module *m) {
|
|
#if SOT_IS_SUPPORTED
|
|
PyInit__eval_frame();
|
|
#if PY_3_11_PLUS
|
|
PyInit__frame_proxy();
|
|
#endif
|
|
m->def(
|
|
"set_eval_frame",
|
|
[](const py::object &py_func) {
|
|
VLOG(5) << "start call set_eval_frame_py.";
|
|
auto ret = set_eval_frame_py(py_func.ptr());
|
|
auto obj = py::reinterpret_borrow<py::object>(ret);
|
|
return obj;
|
|
},
|
|
py::arg("callback"));
|
|
|
|
m->def("has_custom_getattro", [](py::object obj) {
|
|
PyObject *py_obj = obj.ptr();
|
|
|
|
if (!PyType_Check(py_obj)) {
|
|
PADDLE_THROW(common::errors::InvalidArgument(
|
|
"The input object should be a type object, but got %s.",
|
|
py::str(py_obj).cast<std::string>()));
|
|
}
|
|
PyTypeObject *type = reinterpret_cast<PyTypeObject *>(py_obj);
|
|
|
|
return type->tp_getattro != PyObject_GenericGetAttr;
|
|
});
|
|
|
|
m->def(
|
|
"sot_setup_codes_with_graph",
|
|
[](const py::object &py_codes) {
|
|
auto ret = setup_codes_with_graph(py_codes.ptr());
|
|
auto obj = py::reinterpret_borrow<py::object>(ret);
|
|
return obj;
|
|
},
|
|
py::arg("py_codes"));
|
|
|
|
m->def(
|
|
"sot_set_with_graph",
|
|
[](const py::object &py_codes) {
|
|
auto ret = set_with_graph(py_codes.ptr());
|
|
auto obj = py::reinterpret_borrow<py::object>(ret);
|
|
return obj;
|
|
},
|
|
py::arg("py_codes"));
|
|
|
|
m->def(
|
|
"eval_frame_no_skip_codes",
|
|
[](const py::object &py_codes) {
|
|
auto ret = no_skip_codes(py_codes.ptr());
|
|
auto obj = py::reinterpret_borrow<py::object>(ret);
|
|
return obj;
|
|
},
|
|
py::arg("py_codes"));
|
|
|
|
m->def(
|
|
"eval_frame_skip_file_prefix",
|
|
[](const py::object &py_codes) {
|
|
auto ret = skip_file_prefix(py_codes.ptr());
|
|
auto obj = py::reinterpret_borrow<py::object>(ret);
|
|
return obj;
|
|
},
|
|
py::arg("py_codes"));
|
|
BindGuard(m);
|
|
BindGuardTree(m);
|
|
#endif
|
|
}
|
|
|
|
} // namespace paddle::pybind
|