WilliamLian
5 years ago
parent
0a3bf64b79
commit
9682d08d96
@ -1,71 +0,0 @@
|
|||||||
/**
|
|
||||||
* Copyright 2020 Huawei Technologies Co., Ltd
|
|
||||||
*
|
|
||||||
* 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 "ir/primitive_base.h"
|
|
||||||
|
|
||||||
#include <utility>
|
|
||||||
|
|
||||||
namespace mindspore {
|
|
||||||
bool Primitive::operator==(const Value &other) const {
|
|
||||||
if (other.isa<Primitive>()) {
|
|
||||||
auto other_prim = static_cast<const Primitive &>(other);
|
|
||||||
return *this == other_prim;
|
|
||||||
} else {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
bool Primitive::operator==(const Primitive &other) const {
|
|
||||||
if (name() != other.name()) {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
if (attrs_.size() != other.attrs_.size()) {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
auto all = std::all_of(attrs_.begin(), attrs_.end(), [&other](const std::pair<std::string, ValuePtr> &item) -> bool {
|
|
||||||
if (item.second == nullptr) {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
auto iter = other.attrs_.find(item.first);
|
|
||||||
if (iter == other.attrs_.end()) {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
return *item.second == *iter->second;
|
|
||||||
});
|
|
||||||
return all;
|
|
||||||
}
|
|
||||||
|
|
||||||
std::string Primitive::GetAttrsText() const {
|
|
||||||
if (attrs_.empty()) {
|
|
||||||
return "";
|
|
||||||
}
|
|
||||||
|
|
||||||
std::ostringstream oss;
|
|
||||||
oss << "[";
|
|
||||||
bool is_first = true;
|
|
||||||
for (auto &attr : attrs_) {
|
|
||||||
if (is_first) {
|
|
||||||
is_first = false;
|
|
||||||
} else {
|
|
||||||
oss << ", ";
|
|
||||||
}
|
|
||||||
oss << attr.first << "=" << attr.second->DumpText();
|
|
||||||
}
|
|
||||||
oss << "]";
|
|
||||||
|
|
||||||
return oss.str();
|
|
||||||
}
|
|
||||||
} // namespace mindspore
|
|
||||||
@ -1,150 +0,0 @@
|
|||||||
/**
|
|
||||||
* Copyright 2020 Huawei Technologies Co., Ltd
|
|
||||||
*
|
|
||||||
* 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.
|
|
||||||
*/
|
|
||||||
|
|
||||||
#ifndef MINDSPORE_CCSRC_IR_PRIMITIVE_BASE_H_
|
|
||||||
#define MINDSPORE_CCSRC_IR_PRIMITIVE_BASE_H_
|
|
||||||
|
|
||||||
#include <unordered_map>
|
|
||||||
#include <vector>
|
|
||||||
#include <memory>
|
|
||||||
#include <string>
|
|
||||||
#include <tuple>
|
|
||||||
|
|
||||||
#include "ir/dtype/type.h"
|
|
||||||
#include "pybind11/pybind11.h"
|
|
||||||
|
|
||||||
namespace py = pybind11;
|
|
||||||
|
|
||||||
namespace mindspore {
|
|
||||||
// Supported meta type
|
|
||||||
enum PrimType {
|
|
||||||
kPrimTypeUnknown = 0,
|
|
||||||
kPrimTypeBegin = kTypeUnknown,
|
|
||||||
kPrimTypeBuiltIn, // Built-in primitive operator
|
|
||||||
kPrimTypePyInferShape, // Primitive operator defined by custom
|
|
||||||
kPrimTypePyInferTensor, // Primitive operator defined by custom
|
|
||||||
kPrimTypeUserCustom
|
|
||||||
};
|
|
||||||
|
|
||||||
class Primitive : public Named {
|
|
||||||
public:
|
|
||||||
explicit Primitive(const std::string &name, const bool is_base = true, const PrimType prim_type = kPrimTypeBuiltIn)
|
|
||||||
: Named(name),
|
|
||||||
is_base_(is_base),
|
|
||||||
has_signature_(false),
|
|
||||||
prim_type_(prim_type),
|
|
||||||
record_evaluate_add_attr_(false) {}
|
|
||||||
|
|
||||||
Primitive(const Primitive &prim)
|
|
||||||
: Named(prim),
|
|
||||||
attrs_(prim.attrs_),
|
|
||||||
instance_name_(prim.instance_name_),
|
|
||||||
is_base_(prim.is_base_),
|
|
||||||
has_signature_(prim.has_signature_),
|
|
||||||
prim_type_(prim.prim_type_),
|
|
||||||
record_evaluate_add_attr_(false) {}
|
|
||||||
|
|
||||||
MS_DECLARE_PARENT(Primitive, Named);
|
|
||||||
|
|
||||||
abstract::AbstractBasePtr ToPrimAbstract(const AnfNodePtr &anf_node);
|
|
||||||
std::string ToString() const override { return name(); }
|
|
||||||
void BeginRecordAddAttr() {
|
|
||||||
evaluate_added_attrs_.clear();
|
|
||||||
record_evaluate_add_attr_ = true;
|
|
||||||
}
|
|
||||||
void EndRecordAddAttr() { record_evaluate_add_attr_ = false; }
|
|
||||||
Primitive &AddAttr(const std::string &name, const ValuePtr &attr) {
|
|
||||||
attrs_[name] = attr;
|
|
||||||
if (record_evaluate_add_attr_) {
|
|
||||||
evaluate_added_attrs_[name] = attr;
|
|
||||||
}
|
|
||||||
return *this;
|
|
||||||
}
|
|
||||||
|
|
||||||
Primitive &SetAttrs(const std::unordered_map<std::string, ValuePtr> &attrs) {
|
|
||||||
for (auto &attr : attrs) {
|
|
||||||
attrs_[attr.first] = attr.second;
|
|
||||||
}
|
|
||||||
return *this;
|
|
||||||
}
|
|
||||||
|
|
||||||
void set_attr(const std::string &attrName, const ValuePtr &attr) { attrs_[attrName] = attr; }
|
|
||||||
void EraseAttr(const std::string &attrName) { (void)attrs_.erase(attrName); }
|
|
||||||
|
|
||||||
ValuePtr GetAttr(const std::string &attrName) const {
|
|
||||||
auto iter = attrs_.find(attrName);
|
|
||||||
return iter == attrs_.cend() ? nullptr : iter->second;
|
|
||||||
}
|
|
||||||
|
|
||||||
const std::unordered_map<std::string, ValuePtr> &attrs() const { return attrs_; }
|
|
||||||
const std::unordered_map<std::string, ValuePtr> &evaluate_added_attrs() const { return evaluate_added_attrs_; }
|
|
||||||
|
|
||||||
// if Primitive has any attribute, for Primitives like scalar_add, return, etc, don't have any attribute.
|
|
||||||
bool HasAttr() const { return !attrs_.empty(); }
|
|
||||||
bool HasAttr(const std::string &attrName) const {
|
|
||||||
auto iter = attrs_.find(attrName);
|
|
||||||
return !(iter == attrs_.cend());
|
|
||||||
}
|
|
||||||
void set_prim_type(const PrimType t) { prim_type_ = t; }
|
|
||||||
void set_instance_name(const std::string s) { instance_name_ = s; }
|
|
||||||
bool HasPyEvaluator() const { return prim_type_ == kPrimTypePyInferShape || prim_type_ == kPrimTypeUserCustom; }
|
|
||||||
bool HasPyInferTensor() const { return prim_type_ == kPrimTypePyInferTensor; }
|
|
||||||
bool IsCustomPrim() const { return prim_type_ == kPrimTypeUserCustom; }
|
|
||||||
|
|
||||||
PrimType prim_type() const { return prim_type_; }
|
|
||||||
std::string instance_name() const { return instance_name_; }
|
|
||||||
std::string GetAttrsText() const;
|
|
||||||
bool operator==(const Value &other) const override;
|
|
||||||
bool operator==(const Primitive &other) const;
|
|
||||||
~Primitive() override = default;
|
|
||||||
|
|
||||||
void set_has_signature(bool has_signature) { has_signature_ = has_signature; }
|
|
||||||
bool has_signature() const { return has_signature_; }
|
|
||||||
bool is_base() const { return is_base_; }
|
|
||||||
|
|
||||||
protected:
|
|
||||||
std::unordered_map<std::string, ValuePtr> attrs_;
|
|
||||||
std::unordered_map<std::string, ValuePtr> evaluate_added_attrs_;
|
|
||||||
|
|
||||||
private:
|
|
||||||
std::string instance_name_;
|
|
||||||
bool is_base_;
|
|
||||||
bool has_signature_;
|
|
||||||
PrimType prim_type_;
|
|
||||||
bool record_evaluate_add_attr_;
|
|
||||||
};
|
|
||||||
|
|
||||||
inline std::ostream &operator<<(std::ostream &os, const PrimitivePtr &p) {
|
|
||||||
os << *p;
|
|
||||||
return os;
|
|
||||||
}
|
|
||||||
|
|
||||||
struct PrimitiveEqual {
|
|
||||||
bool operator()(PrimitivePtr const &t1, PrimitivePtr const &t2) const {
|
|
||||||
MS_EXCEPTION_IF_NULL(t1);
|
|
||||||
MS_EXCEPTION_IF_NULL(t2);
|
|
||||||
return t1->name() == t2->name();
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
struct PrimitiveHasher {
|
|
||||||
std::size_t operator()(PrimitivePtr const &prim) const {
|
|
||||||
MS_EXCEPTION_IF_NULL(prim);
|
|
||||||
return prim->Hash();
|
|
||||||
}
|
|
||||||
};
|
|
||||||
} // namespace mindspore
|
|
||||||
#endif // MINDSPORE_CCSRC_IR_PRIMITIVE_BASE_H_
|
|
||||||
@ -0,0 +1,195 @@
|
|||||||
|
/**
|
||||||
|
* Copyright 2019-2020 Huawei Technologies Co., Ltd
|
||||||
|
*
|
||||||
|
* 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 "ir/primitive_py.h"
|
||||||
|
#include <mutex>
|
||||||
|
#include <utility>
|
||||||
|
#include "ir/signature.h"
|
||||||
|
#include "operator/ops.h"
|
||||||
|
#include "./common.h"
|
||||||
|
#include "pipeline/parse/python_adapter.h"
|
||||||
|
#include "pipeline/parse/data_converter.h"
|
||||||
|
#include "pybind11/pytypes.h"
|
||||||
|
#include "utils/convert_utils_base.h"
|
||||||
|
#include "utils/primitive_utils.h"
|
||||||
|
#include "utils/base_ref_py.h"
|
||||||
|
#include "pybind_api/api_register.h"
|
||||||
|
#include "pybind_api/export_flags.h"
|
||||||
|
|
||||||
|
namespace mindspore {
|
||||||
|
namespace {
|
||||||
|
constexpr auto kBpropAttrName = "bprop";
|
||||||
|
constexpr auto kCellHookAttrName = "cell_hook";
|
||||||
|
constexpr auto kCellIDAttrName = "cell_id";
|
||||||
|
void SyncData(const py::object &arg) {
|
||||||
|
if (py::isinstance<py::tuple>(arg)) {
|
||||||
|
py::tuple arg_list = py::cast<py::tuple>(arg);
|
||||||
|
for (size_t i = 0; i < arg_list.size(); i++) {
|
||||||
|
SyncData(arg_list[i]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if (py::isinstance<tensor::Tensor>(arg)) {
|
||||||
|
auto tensor = py::cast<tensor::TensorPtr>(arg);
|
||||||
|
(void)tensor->data_sync();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
} // namespace
|
||||||
|
std::map<std::string, py::object> PrimitivePy::hook_grad_;
|
||||||
|
static ValuePtr PyArgToValue(const py::object &arg) {
|
||||||
|
if (py::isinstance<SignatureEnumKind>(arg) &&
|
||||||
|
py::cast<SignatureEnumKind>(arg) == SignatureEnumKind::kKindEmptyDefaultValue) {
|
||||||
|
return nullptr;
|
||||||
|
}
|
||||||
|
return parse::data_converter::PyDataToValue(arg);
|
||||||
|
}
|
||||||
|
|
||||||
|
void PrimitivePy::set_signatures(
|
||||||
|
std::vector<std::tuple<std::string, SignatureEnumRW, SignatureEnumKind, py::object, SignatureEnumDType>> signatures) {
|
||||||
|
signatures_.clear();
|
||||||
|
for (auto &signature : signatures) {
|
||||||
|
auto [name, rw, kind, arg_default, dtype] = signature;
|
||||||
|
auto default_value = PyArgToValue(arg_default);
|
||||||
|
signatures_.emplace_back(name, rw, kind, default_value, dtype);
|
||||||
|
}
|
||||||
|
set_has_signature(true);
|
||||||
|
}
|
||||||
|
|
||||||
|
py::function PrimitivePy::GetBpropFunction() {
|
||||||
|
static const char *const get_bprop_func_name = "get_bprop";
|
||||||
|
if (py::hasattr(python_obj_, get_bprop_func_name)) {
|
||||||
|
py::function fn = python_obj_.attr(get_bprop_func_name)().cast<py::function>();
|
||||||
|
return fn;
|
||||||
|
} else {
|
||||||
|
auto fn = GetBpropFunctionByObj(python_obj_);
|
||||||
|
return fn;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
BaseRef PrimitivePy::RunHookFunction(const VectorRef &args) const {
|
||||||
|
auto py_args = py::tuple(args.size());
|
||||||
|
size_t i = 0;
|
||||||
|
for (auto &arg : args) {
|
||||||
|
py_args[i] = BaseRefToPyData(arg);
|
||||||
|
MS_LOG(DEBUG) << "arg:" << i << ":";
|
||||||
|
i++;
|
||||||
|
}
|
||||||
|
py::object obj;
|
||||||
|
bool is_bprop = this->HasAttr(kBpropAttrName);
|
||||||
|
if (is_bprop) {
|
||||||
|
SyncData(py_args);
|
||||||
|
obj = hook_(*py_args);
|
||||||
|
return std::make_shared<PyObjectRef>(obj);
|
||||||
|
}
|
||||||
|
SyncData(py_args[2]);
|
||||||
|
bool is_cell = this->HasAttr(kCellHookAttrName);
|
||||||
|
if (is_cell) {
|
||||||
|
auto cell_id = GetValue<std::string>(this->GetAttr(kCellIDAttrName));
|
||||||
|
auto iter = hook_grad_.find(cell_id);
|
||||||
|
if (iter != hook_grad_.end()) {
|
||||||
|
auto hook_args = py::tuple(3);
|
||||||
|
hook_args[0] = cell_id;
|
||||||
|
hook_args[1] = py::make_tuple(iter->second);
|
||||||
|
hook_args[2] = py::make_tuple(py_args[2]);
|
||||||
|
obj = hook_(*hook_args);
|
||||||
|
if (py::isinstance<py::none>(obj)) {
|
||||||
|
obj = py_args[2];
|
||||||
|
}
|
||||||
|
hook_grad_.erase(cell_id);
|
||||||
|
} else {
|
||||||
|
hook_grad_[cell_id] = py_args[2];
|
||||||
|
obj = py_args[2];
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// Hook operator for execute variable hook function
|
||||||
|
obj = hook_(py::make_tuple(py_args[2]));
|
||||||
|
if (py::isinstance<py::none>(obj)) {
|
||||||
|
obj = py_args[2];
|
||||||
|
}
|
||||||
|
}
|
||||||
|
obj = py::make_tuple(obj);
|
||||||
|
return std::make_shared<PyObjectRef>(obj);
|
||||||
|
}
|
||||||
|
|
||||||
|
py::function PrimitivePy::GetComputeFunction() {
|
||||||
|
static const char *const compute_func_name = "vm_impl";
|
||||||
|
|
||||||
|
if (py::hasattr(python_obj_, compute_func_name)) {
|
||||||
|
MS_LOG(INFO) << name() << " compute_func_name";
|
||||||
|
py::function fn = python_obj_.attr(compute_func_name).cast<py::function>();
|
||||||
|
return fn;
|
||||||
|
}
|
||||||
|
|
||||||
|
static const std::string vm_module = "mindspore.ops.vm_impl_registry";
|
||||||
|
static const std::string get_vm_impl_fn = "get_vm_impl_fn";
|
||||||
|
MS_LOG(INFO) << name() << ": get_vm_impl_fn";
|
||||||
|
py::function get_fn = parse::python_adapter::GetPyFn(vm_module, get_vm_impl_fn);
|
||||||
|
py::function vm_fn = get_fn(python_obj_);
|
||||||
|
|
||||||
|
if (py::isinstance<py::none>(vm_fn)) {
|
||||||
|
MS_LOG(WARNING) << "Cannot find " << python_obj_.attr("__class__").attr("__name__").cast<std::string>();
|
||||||
|
vm_fn = mindspore::GetComputeFunction(Primitive::name());
|
||||||
|
}
|
||||||
|
return vm_fn;
|
||||||
|
}
|
||||||
|
|
||||||
|
void PrimitivePy::AddPyAttr(const py::str &name, const py::object &obj) {
|
||||||
|
std::string attr_name = name;
|
||||||
|
ValuePtr converted_ret = nullptr;
|
||||||
|
if (py::isinstance<py::module>(obj)) {
|
||||||
|
MS_LOG(EXCEPTION) << "AddPyAttr failed, obj should not be py::module";
|
||||||
|
}
|
||||||
|
bool converted = parse::ConvertData(obj, &converted_ret);
|
||||||
|
if (!converted) {
|
||||||
|
MS_LOG(EXCEPTION) << "Attribute convert error with type: " << std::string(py::str(obj));
|
||||||
|
}
|
||||||
|
(void)this->AddAttr(attr_name, converted_ret);
|
||||||
|
}
|
||||||
|
|
||||||
|
py::dict PrimitivePy::GetAttrDict() {
|
||||||
|
py::dict attr_dict;
|
||||||
|
for (auto &attr : attrs_) {
|
||||||
|
attr_dict[py::str(attr.first)] = ValuePtrToPyData(attr.second);
|
||||||
|
}
|
||||||
|
return attr_dict;
|
||||||
|
}
|
||||||
|
|
||||||
|
void PrimitivePy::CopyHookFunction(const PrimitivePtr &primitive) {
|
||||||
|
MS_EXCEPTION_IF_NULL(primitive);
|
||||||
|
if (!primitive->isa<PrimitivePy>()) {
|
||||||
|
MS_LOG(EXCEPTION) << "Cannot copy a primtive which is not python primitive hook function to python primitive!";
|
||||||
|
}
|
||||||
|
auto primitive_py = primitive->cast<PrimitivePyPtr>();
|
||||||
|
MS_EXCEPTION_IF_NULL(primitive_py);
|
||||||
|
this->set_hook(primitive_py->hook());
|
||||||
|
}
|
||||||
|
|
||||||
|
REGISTER_PYBIND_DEFINE(Primitive_, ([](const py::module *m) {
|
||||||
|
(void)py::enum_<PrimType>(*m, "prim_type", py::arithmetic())
|
||||||
|
.value("unknown", PrimType::kPrimTypeUnknown)
|
||||||
|
.value("builtin", PrimType::kPrimTypeBuiltIn)
|
||||||
|
.value("py_infer_shape", PrimType::kPrimTypePyInferShape)
|
||||||
|
.value("user_custom", PrimType::kPrimTypeUserCustom);
|
||||||
|
(void)py::class_<PrimitivePy, std::shared_ptr<PrimitivePy>>(*m, "Primitive_")
|
||||||
|
.def_readonly(PYTHON_PRIMITIVE_FLAG, &PrimitivePy::parse_info_)
|
||||||
|
.def(py::init<py::str &, py::object>())
|
||||||
|
.def("add_attr", &PrimitivePy::AddPyAttr, "add primitive attr")
|
||||||
|
.def("get_attr_dict", &PrimitivePy::GetAttrDict, "get primitive attr")
|
||||||
|
.def("set_prim_type", &PrimitivePy::set_prim_type, "Set primitive type.")
|
||||||
|
.def("set_signatures", &PrimitivePy::set_signatures, "Set primitive inputs signature.")
|
||||||
|
.def("register_hook", &PrimitivePy::set_hook, "Set primitive hook function.")
|
||||||
|
.def("set_instance_name", &PrimitivePy::set_instance_name, "Set primitive instance name.");
|
||||||
|
}));
|
||||||
|
} // namespace mindspore
|
||||||
@ -0,0 +1,72 @@
|
|||||||
|
/**
|
||||||
|
* Copyright 2019 Huawei Technologies Co., Ltd
|
||||||
|
*
|
||||||
|
* 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.
|
||||||
|
*/
|
||||||
|
|
||||||
|
#ifndef MINDSPORE_CCSRC_IR_PRIMITIVE_PY_H_
|
||||||
|
#define MINDSPORE_CCSRC_IR_PRIMITIVE_PY_H_
|
||||||
|
|
||||||
|
#include <unordered_map>
|
||||||
|
#include <vector>
|
||||||
|
#include <memory>
|
||||||
|
#include <string>
|
||||||
|
#include <tuple>
|
||||||
|
#include <map>
|
||||||
|
|
||||||
|
#include "pipeline/static_analysis/abstract_value.h"
|
||||||
|
#include "utils/misc.h"
|
||||||
|
#include "pybind11/pybind11.h"
|
||||||
|
#include "utils/log_adapter.h"
|
||||||
|
#include "ir/primitive.h"
|
||||||
|
#include "ir/signature.h"
|
||||||
|
#include "parallel/ops_info/operator_info.h"
|
||||||
|
namespace py = pybind11;
|
||||||
|
namespace mindspore {
|
||||||
|
class PrimitivePy : public Primitive {
|
||||||
|
public:
|
||||||
|
PrimitivePy(const py::str &name, const py::object &python_obj)
|
||||||
|
: Primitive(name, false), python_obj_(python_obj), signatures_() {}
|
||||||
|
~PrimitivePy() override = default;
|
||||||
|
MS_DECLARE_PARENT(PrimitivePy, Primitive);
|
||||||
|
py::function GetBpropFunction();
|
||||||
|
py::function GetComputeFunction();
|
||||||
|
|
||||||
|
void set_signatures(
|
||||||
|
std::vector<std::tuple<std::string, SignatureEnumRW, SignatureEnumKind, py::object, SignatureEnumDType>>
|
||||||
|
signatures);
|
||||||
|
|
||||||
|
const std::vector<Signature> &signatures() const { return signatures_; }
|
||||||
|
|
||||||
|
void CopyHookFunction(const PrimitivePtr &primitive) override;
|
||||||
|
|
||||||
|
void AddPyAttr(const py::str &name, const py::object &obj);
|
||||||
|
|
||||||
|
py::dict GetAttrDict();
|
||||||
|
void set_hook(const py::function &hook) { hook_ = hook; }
|
||||||
|
py::function hook() const { return hook_; }
|
||||||
|
BaseRef RunHookFunction(const VectorRef &args) const override;
|
||||||
|
const bool parse_info_ = true;
|
||||||
|
const py::object &GetPyObj() const { return python_obj_; }
|
||||||
|
bool is_tuple_input_ = false;
|
||||||
|
|
||||||
|
private:
|
||||||
|
py::object python_obj_;
|
||||||
|
py::function hook_;
|
||||||
|
std::vector<Signature> signatures_;
|
||||||
|
static std::map<std::string, py::object> hook_grad_;
|
||||||
|
};
|
||||||
|
|
||||||
|
using PrimitivePyPtr = std::shared_ptr<PrimitivePy>;
|
||||||
|
} // namespace mindspore
|
||||||
|
#endif // MINDSPORE_CCSRC_IR_PRIMITIVE_PY_H_
|
||||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in new issue