You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.

primitive_py.cc 23 kB

5 years ago
5 years ago
5 years ago
5 years ago
4 years ago
4 years ago
4 years ago
4 years ago
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550
  1. /**
  2. * Copyright 2019-2020 Huawei Technologies Co., Ltd
  3. *
  4. * Licensed under the Apache License, Version 2.0 (the "License");
  5. * you may not use this file except in compliance with the License.
  6. * You may obtain a copy of the License at
  7. *
  8. * http://www.apache.org/licenses/LICENSE-2.0
  9. *
  10. * Unless required by applicable law or agreed to in writing, software
  11. * distributed under the License is distributed on an "AS IS" BASIS,
  12. * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. * See the License for the specific language governing permissions and
  14. * limitations under the License.
  15. */
  16. #include "pybind_api/ir/primitive_py.h"
  17. #include <mutex>
  18. #include <map>
  19. #include <utility>
  20. #include "ir/signature.h"
  21. #include "pipeline/jit/parse/data_converter.h"
  22. #include "pipeline/jit/parse/python_adapter.h"
  23. #include "pybind11/pytypes.h"
  24. #include "pybind_api/api_register.h"
  25. #include "pybind_api/export_flags.h"
  26. #include "pybind_api/ir/base_ref_py.h"
  27. #include "utils/convert_utils_base.h"
  28. #include "utils/convert_utils_py.h"
  29. #include "utils/ms_context.h"
  30. #include "utils/primitive_utils.h"
  31. #include "utils/check_convert_utils.h"
  32. #include "pipeline/jit/resource.h"
  33. #include "pipeline/pynative/pynative_execute.h"
  34. namespace mindspore {
  35. namespace {
  36. constexpr auto kBpropAttrName = "bprop";
  37. constexpr auto kCellHookAttrName = "cell_hook";
  38. constexpr auto kCellIDAttrName = "cell_id";
  39. std::map<std::string, std::string> kOpAttrNameReplaceMap = {
  40. {"data_format", "format"},
  41. };
  42. void SyncData(const py::object &arg) {
  43. if (py::isinstance<py::tuple>(arg)) {
  44. py::tuple arg_list = py::cast<py::tuple>(arg);
  45. for (size_t i = 0; i < arg_list.size(); i++) {
  46. SyncData(arg_list[i]);
  47. }
  48. }
  49. if (py::isinstance<tensor::Tensor>(arg)) {
  50. auto tensor = py::cast<tensor::TensorPtr>(arg);
  51. tensor->data_sync();
  52. }
  53. }
  54. } // namespace
  55. std::map<std::string, py::object> PrimitivePy::hook_grad_;
  56. PrimitivePy::PrimitivePy(const std::string &name) : Primitive(name, false), python_obj_(py::none()) {}
  57. PrimitivePy::PrimitivePy(const py::object &python_obj, const PrimitivePyAdapterPtr &adapter)
  58. : Primitive(adapter->name_, false), python_obj_(python_obj), adapter_(adapter) {
  59. MS_LOG(DEBUG) << "New primitive:" << adapter->name_;
  60. set_signatures(adapter->signatures_);
  61. (void)Primitive::SetAttrs(adapter->attrs_);
  62. Primitive::set_prim_type(adapter->prim_type_);
  63. Primitive::set_const_prim(adapter->is_const_prim_);
  64. Primitive::set_const_input_indexes(adapter->const_input_indexes_);
  65. set_hook(adapter->hook_);
  66. set_instance_name(adapter->instance_name_);
  67. }
  68. PrimitivePy::~PrimitivePy() {}
  69. void PrimitivePy::set_signatures(const std::vector<Signature> &signatures) {
  70. signatures_ = signatures;
  71. set_has_signature(!signatures.empty());
  72. }
  73. py::function PrimitivePy::GetBpropFunction() {
  74. static const char *const get_bprop_func_name = "get_bprop";
  75. if (py::hasattr(python_obj_, get_bprop_func_name)) {
  76. py::function fn = python_obj_.attr(get_bprop_func_name)().cast<py::function>();
  77. return fn;
  78. } else {
  79. auto fn = GetBpropFunctionByObj(python_obj_);
  80. return fn;
  81. }
  82. }
  83. py::tuple check_bprop_out(const py::object &grads_obj, const py::tuple &py_args, const std::string &bprop_cls_name) {
  84. py::tuple grads;
  85. if (!py::isinstance<py::tuple>(grads_obj)) {
  86. grads = py::make_tuple(grads_obj);
  87. } else {
  88. grads = py::cast<py::tuple>(grads_obj);
  89. }
  90. if (!MsContext::GetInstance()->get_param<bool>(MS_CTX_CHECK_BPROP_FLAG)) {
  91. return grads;
  92. }
  93. constexpr int filter_args_size = 2;
  94. if (grads.size() != py_args.size() - filter_args_size) {
  95. MS_EXCEPTION(TypeError) << "For user defined method 'bprop' of net '" << bprop_cls_name
  96. << "', the number of return values(gradients) should be equal to the number of input "
  97. "arguments except 'out' and 'dout', which is: "
  98. << (py_args.size() - filter_args_size) << ", but got:" << grads.size() << ".";
  99. }
  100. for (size_t i = 0; i < grads.size(); i++) {
  101. if (py::isinstance<tensor::Tensor>(py_args[i])) {
  102. if (!py::isinstance<tensor::Tensor>(grads[i])) {
  103. MS_EXCEPTION(ValueError) << "For user defined method 'bprop' of net '" << bprop_cls_name << "', the " << i
  104. << "th return value(gradient of the " << i << "th argument) should be Tensor, but got "
  105. << py::cast<std::string>(grads[i].attr("__class__").attr("__name__"))
  106. << ", and the value is " << py::cast<py::str>(grads[i]) << ".";
  107. }
  108. py::object arg_dtype = py_args[i].attr("dtype");
  109. py::object grad_dtype = grads[i].attr("dtype");
  110. py::tuple arg_shape = py_args[i].attr("shape");
  111. py::tuple grad_shape = grads[i].attr("shape");
  112. if (!grad_dtype.equal(arg_dtype)) {
  113. MS_EXCEPTION(TypeError) << "For user defined method 'bprop' of net '" << bprop_cls_name << "', the " << i
  114. << "th return value(gradient of the " << i
  115. << "th argument) should have the same dtype as the " << i
  116. << "th argument, which is:" << py::cast<py::str>(arg_dtype)
  117. << ", but got: " << py::cast<py::str>(grad_dtype) << ".";
  118. }
  119. if (!grad_shape.equal(arg_shape)) {
  120. MS_EXCEPTION(ValueError) << "For user defined method 'bprop' of net '" << bprop_cls_name << "', the " << i
  121. << "th return value(gradient of the " << i
  122. << "th argument) should have the same shape as the " << i
  123. << "th argument, which is:" << py::cast<py::str>(arg_shape)
  124. << ", but got: " << py::cast<py::str>(grad_shape) << ".";
  125. }
  126. }
  127. }
  128. return grads;
  129. }
  130. void PrimitivePy::ConvertCTensorToPyTensor(const py::tuple &input_args, py::tuple *convert_args) const {
  131. MS_EXCEPTION_IF_NULL(convert_args);
  132. if (input_args.size() != (*convert_args).size()) {
  133. MS_LOG(EXCEPTION) << "The size of input_args: " << input_args.size()
  134. << " should be equal to the size of convert_args: " << (*convert_args).size();
  135. }
  136. for (size_t i = 0; i < input_args.size(); ++i) {
  137. (*convert_args)[i] = py::isinstance<tensor::Tensor>(input_args[i])
  138. ? parse::python_adapter::CallPyFn(parse::PYTHON_MOD_PARSE_MODULE,
  139. parse::PYTHON_MOD_CONVERT_TO_MS_TENSOR, input_args[i])
  140. : input_args[i];
  141. }
  142. }
  143. void PrimitivePy::CheckHookConsistency(const py::object &grad_out, const py::object &expected_grad_out) const {
  144. if (py::isinstance<py::tuple>(expected_grad_out)) {
  145. if (!py::isinstance<py::tuple>(grad_out)) {
  146. hook_grad_.clear();
  147. MS_EXCEPTION(TypeError) << "The output gradient should be a tuple!";
  148. }
  149. auto actual_out_tuple = py::cast<py::tuple>(grad_out);
  150. auto expected_out_tuple = py::cast<py::tuple>(expected_grad_out);
  151. if (actual_out_tuple.size() != expected_out_tuple.size()) {
  152. hook_grad_.clear();
  153. MS_EXCEPTION(ValueError) << "The tuple size of output gradient should be " << expected_out_tuple.size()
  154. << ", but it is " << actual_out_tuple.size();
  155. }
  156. for (size_t i = 0; i < expected_out_tuple.size(); ++i) {
  157. CheckHookConsistency(actual_out_tuple[i], expected_out_tuple[i]);
  158. }
  159. }
  160. if (py::isinstance<tensor::Tensor>(expected_grad_out)) {
  161. if (!py::isinstance<tensor::Tensor>(grad_out)) {
  162. hook_grad_.clear();
  163. py::object code_obj = py::getattr(hook_, "__code__");
  164. py::object co_name = py::getattr(code_obj, "co_name");
  165. MS_EXCEPTION(TypeError) << "The output type of:" << py::str(co_name) << " should be a tensor but got "
  166. << py::cast<std::string>(grad_out.attr("__class__").attr("__name__")) << ".";
  167. }
  168. auto actual_out_tensor = py::cast<tensor::TensorPtr>(grad_out);
  169. auto expected_out_tensor = py::cast<tensor::TensorPtr>(expected_grad_out);
  170. MS_EXCEPTION_IF_NULL(actual_out_tensor);
  171. MS_EXCEPTION_IF_NULL(expected_out_tensor);
  172. if (actual_out_tensor->GetShapeAndDataTypeInfo() != expected_out_tensor->GetShapeAndDataTypeInfo()) {
  173. hook_grad_.clear();
  174. py::object code_obj = py::getattr(hook_, "__code__");
  175. py::object co_name = py::getattr(code_obj, "co_name");
  176. MS_EXCEPTION(ValueError) << "The output type of " << py::str(co_name)
  177. << " is not consistent with the expected, it should be "
  178. << expected_out_tensor->GetShapeAndDataTypeInfo() << ", but got "
  179. << actual_out_tensor->GetShapeAndDataTypeInfo();
  180. }
  181. }
  182. }
  183. BaseRef PrimitivePy::RunCellBpropFunction(const py::tuple &py_args) const {
  184. SyncData(py_args);
  185. auto size = py_args.size();
  186. constexpr size_t grad_param_nums = 2;
  187. py::tuple input_args(size - grad_param_nums);
  188. for (size_t i = 0; i < size - grad_param_nums; ++i) {
  189. input_args[i] = py_args[i];
  190. }
  191. py::tuple convert_args(py_args.size());
  192. ConvertCTensorToPyTensor(py_args, &convert_args);
  193. auto inst = pynative::PynativeExecutor::GetInstance();
  194. MS_EXCEPTION_IF_NULL(inst);
  195. try {
  196. MS_LOG(DEBUG) << "Run bprop function start";
  197. inst->NewGraph(hook_, input_args.cast<py::args>());
  198. py::object grads_obj = hook_(*convert_args);
  199. py::tuple grads = check_bprop_out(grads_obj, py_args, bprop_cls_name_);
  200. inst->EndGraph(hook_, grads_obj, input_args.cast<py::args>());
  201. MS_LOG(DEBUG) << "Run bprop function end";
  202. return std::make_shared<PyObjectRef>(grads);
  203. } catch (std::exception &bt) {
  204. inst->ClearRes();
  205. std::rethrow_exception(std::current_exception());
  206. }
  207. }
  208. BaseRef PrimitivePy::RunCellHookFunction(const py::tuple &py_args) const {
  209. constexpr size_t grad_input_index = 1;
  210. constexpr size_t grad_output_index = 2;
  211. constexpr size_t input_param_nums = 3;
  212. SyncData(py_args[grad_output_index]);
  213. py::object obj;
  214. auto cell_id = GetValue<std::string>(this->GetAttr(kCellIDAttrName));
  215. auto iter = hook_grad_.find(cell_id);
  216. if (iter != hook_grad_.end()) {
  217. py::object code_obj = py::getattr(hook_, "__code__");
  218. py::object co_name = py::getattr(code_obj, "co_name");
  219. if (std::string(py::str(co_name)) == "staging_specialize") {
  220. py::object name_obj = py::getattr(hook_, "__name__");
  221. MS_LOG(EXCEPTION) << "Decorating hook function " << py::str(name_obj) << " with '@ms_function' is not supported.";
  222. }
  223. py::tuple convert_args(input_param_nums - 1);
  224. py::tuple input_args(input_param_nums - 1);
  225. input_args[0] = iter->second;
  226. input_args[1] = py_args[grad_output_index];
  227. ConvertCTensorToPyTensor(input_args, &convert_args);
  228. auto hook_args = py::tuple(input_param_nums);
  229. hook_args[0] = cell_id;
  230. hook_args[grad_input_index] = py::make_tuple(convert_args[0]);
  231. hook_args[grad_output_index] = py::make_tuple(convert_args[1]);
  232. obj = hook_(*hook_args);
  233. if (py::isinstance<py::none>(obj)) {
  234. obj = py_args[grad_output_index];
  235. }
  236. CheckHookConsistency(obj, py_args[grad_output_index]);
  237. (void)hook_grad_.erase(cell_id);
  238. } else {
  239. hook_grad_[cell_id] = py_args[grad_output_index];
  240. obj = py_args[grad_output_index];
  241. }
  242. obj = py::make_tuple(obj);
  243. return std::make_shared<PyObjectRef>(obj);
  244. }
  245. BaseRef PrimitivePy::RunVariableHookFunction(const py::tuple &py_args) const {
  246. py::object code_obj = py::getattr(hook_, "__code__");
  247. py::object co_name = py::getattr(code_obj, "co_name");
  248. if (std::string(py::str(co_name)) == "staging_specialize") {
  249. py::object name_obj = py::getattr(hook_, "__name__");
  250. MS_LOG(EXCEPTION) << "Decorating hook function " << py::str(name_obj) << " with '@ms_function' is not supported.";
  251. }
  252. constexpr size_t grad_output_index = 2;
  253. SyncData(py_args[grad_output_index]);
  254. py::object obj = hook_(py::make_tuple(py_args[grad_output_index]));
  255. if (py::isinstance<py::none>(obj)) {
  256. obj = py_args[grad_output_index];
  257. }
  258. CheckHookConsistency(obj, py_args[grad_output_index]);
  259. obj = py::make_tuple(obj);
  260. return std::make_shared<PyObjectRef>(obj);
  261. }
  262. BaseRef PrimitivePy::RunHookFunction(const VectorRef &args) const {
  263. py::tuple py_args = ConvertDatatoPyTuple(args);
  264. bool is_bprop = this->HasAttr(kBpropAttrName);
  265. if (is_bprop) {
  266. return RunCellBpropFunction(py_args);
  267. }
  268. bool is_cell = this->HasAttr(kCellHookAttrName);
  269. if (is_cell) {
  270. return RunCellHookFunction(py_args);
  271. }
  272. return RunVariableHookFunction(py_args);
  273. }
  274. py::function PrimitivePy::GetComputeFunction() const {
  275. static const char *const compute_func_name = "vm_impl";
  276. if (py::hasattr(python_obj_, compute_func_name)) {
  277. MS_LOG(DEBUG) << name() << " compute_func_name";
  278. py::function fn = python_obj_.attr(compute_func_name).cast<py::function>();
  279. return fn;
  280. }
  281. static const std::string vm_module = "mindspore.ops.vm_impl_registry";
  282. static const std::string get_vm_impl_fn = "get_vm_impl_fn";
  283. MS_LOG(DEBUG) << name() << ": get_vm_impl_fn";
  284. py::function get_fn = parse::python_adapter::GetPyFn(vm_module, get_vm_impl_fn);
  285. py::function vm_fn = get_fn(python_obj_);
  286. if (py::isinstance<py::none>(vm_fn)) {
  287. MS_LOG(DEBUG) << "Cannot find " << python_obj_.attr("__class__").attr("__name__").cast<std::string>();
  288. vm_fn = mindspore::GetComputeFunction(Primitive::name());
  289. }
  290. return vm_fn;
  291. }
  292. py::dict PrimitivePy::GetAttrDict() {
  293. py::dict attr_dict;
  294. for (auto &attr : attrs_) {
  295. attr_dict[py::str(attr.first)] = ValueToPyData(attr.second);
  296. }
  297. return attr_dict;
  298. }
  299. void PrimitivePy::CopyHookFunction(const PrimitivePtr &primitive) {
  300. MS_EXCEPTION_IF_NULL(primitive);
  301. if (!primitive->isa<PrimitivePy>()) {
  302. MS_LOG(EXCEPTION) << "Cannot copy a primitive which is not python primitive hook function to python primitive!";
  303. }
  304. auto primitive_py = primitive->cast<PrimitivePyPtr>();
  305. MS_EXCEPTION_IF_NULL(primitive_py);
  306. this->set_hook(primitive_py->hook());
  307. if (primitive_py->HasAttr(kBpropAttrName)) {
  308. set_bprop_cls_name(primitive_py->bprop_cls_name_);
  309. (void)this->AddAttr(kBpropAttrName, primitive_py->GetAttr(kBpropAttrName));
  310. }
  311. }
  312. BaseRef PrimitivePy::RunComputeFunction(const VectorRef &args) const {
  313. auto py_args = ConvertDatatoPyTuple(args);
  314. auto result = this->RunPyComputeFunction(py_args);
  315. if (py::isinstance<py::none>(result)) {
  316. return std::make_shared<BaseRef>(nullptr);
  317. }
  318. return std::make_shared<PyObjectRef>(result);
  319. }
  320. py::object PrimitivePy::RunPyComputeFunction(const py::tuple &py_args) const {
  321. auto func = this->GetComputeFunction();
  322. if (py::isinstance<py::none>(func)) {
  323. return py::none();
  324. }
  325. auto result = func(*py_args);
  326. return result;
  327. }
  328. bool PrimitivePy::HasComputeFunction() const {
  329. auto func = GetComputeFunction();
  330. return !py::isinstance<py::none>(func);
  331. }
  332. PrimitivePtr PrimitivePy::Clone() {
  333. auto clone_fn = python_obj_.attr("_clone");
  334. py::object obj_adapter = clone_fn();
  335. auto prim_adapter = obj_adapter.cast<PrimitivePyAdapterPtr>();
  336. auto prim = std::make_shared<PrimitivePy>(obj_adapter, prim_adapter);
  337. prim_adapter->set_attached_primitive(prim);
  338. return prim;
  339. }
  340. py::dict PrimitivePy::RunInfer(const py::tuple &args) {
  341. if (!HasPyObj()) {
  342. MS_LOG(EXCEPTION) << "[" << this->ToString() << "]: pyobj is empty";
  343. }
  344. // Python obj could be replaced as None, so it will losed the original info when throw exception in python.
  345. if (!py::hasattr(python_obj_, PY_PRIM_METHOD_INFER)) {
  346. MS_LOG(EXCEPTION) << "prim:" << ToString() << " has no attr:" << PY_PRIM_METHOD_INFER;
  347. }
  348. auto infer_fuc = python_obj_.attr(PY_PRIM_METHOD_INFER);
  349. return infer_fuc(*args);
  350. }
  351. void PrimitivePy::RunCheck(const py::tuple &args) {
  352. if (!HasPyObj()) {
  353. MS_LOG(EXCEPTION) << "[" << this->ToString() << "]: pyobj is empty";
  354. }
  355. // Python obj could be replaced as None, so it will losed the original info when throw exception in python.
  356. if (!py::hasattr(python_obj_, PY_PRIM_METHOD_CHECK)) {
  357. MS_LOG(EXCEPTION) << "prim:" << ToString() << " has no attr:" << PY_PRIM_METHOD_CHECK;
  358. }
  359. auto check_func = python_obj_.attr(PY_PRIM_METHOD_CHECK);
  360. (void)check_func(*args);
  361. }
  362. py::object PrimitivePy::RunInferValue(const py::tuple &args) {
  363. if (!HasPyObj()) {
  364. MS_LOG(EXCEPTION) << "[" << this->ToString() << "]: pyobj is empty";
  365. }
  366. // Python obj could be replaced as None, so it will losed the original info when throw exception in python.
  367. if (!py::hasattr(python_obj_, PY_PRIM_METHOD_INFER_VALUE)) {
  368. MS_LOG(EXCEPTION) << "prim:" << ToString() << " has no attr:" << PY_PRIM_METHOD_INFER_VALUE;
  369. }
  370. auto infer_value = python_obj_.attr(PY_PRIM_METHOD_INFER_VALUE);
  371. return infer_value(*args);
  372. }
  373. void PrimitivePy::ClearHookRes() { hook_grad_.clear(); }
  374. PrimitivePyAdapter::PrimitivePyAdapter(const py::str &name) : name_(name) {}
  375. void PrimitivePyAdapter::AddPyAttr(const py::str &name, const py::object &obj) {
  376. std::string attr_name = name;
  377. ValuePtr converted_ret = nullptr;
  378. if (py::isinstance<py::module>(obj)) {
  379. MS_LOG(EXCEPTION) << "Call 'add_attr' to add attribute to primitive failed,"
  380. << " not support py::module to be attribute value; primitive name: " << this->name_
  381. << ", attribute name: " << attr_name << " attribute value: " << py::str(obj);
  382. }
  383. bool converted = parse::ConvertData(obj, &converted_ret);
  384. if (!converted) {
  385. MS_LOG(EXCEPTION) << "Call 'add_attr' to add attribute to primitive failed,"
  386. << " convert python obj to MindSpore obj failed; primitive name: " << this->name_
  387. << ", attribute name:" << attr_name << ", attribute value:" << py::str(obj)
  388. << ", attribute type:" << py::cast<std::string>(obj.attr("__class__").attr("__name__"));
  389. }
  390. if (kOpAttrNameReplaceMap.find(attr_name) != kOpAttrNameReplaceMap.end()) {
  391. attr_name = kOpAttrNameReplaceMap[attr_name];
  392. }
  393. (void)CheckAndConvertUtils::ConvertAttrValueToInt(this->name_, name, &converted_ret);
  394. if (attr_name == "primitive_target") {
  395. MS_EXCEPTION_IF_NULL(converted_ret);
  396. if (!converted_ret->isa<StringImm>()) {
  397. MS_LOG(EXCEPTION) << "Call 'add_attr' to add attribute to primitive '" << this->name_
  398. << "' failed, value of attribute 'primitive_target' must be CPU|GPU|Ascend but got "
  399. << py::str(obj);
  400. }
  401. auto target = GetValue<std::string>(converted_ret);
  402. if (!target.empty() && target != kCPUDevice && target != kGPUDevice && target != kAscendDevice &&
  403. target != "Device") {
  404. MS_LOG(EXCEPTION) << "Call 'add_attr' to add attribute to primitive '" << this->name_
  405. << "' failed, value of attribute 'primitive_target' must be CPU|GPU|Ascend but got "
  406. << py::str(obj);
  407. }
  408. }
  409. attrs_[attr_name] = converted_ret;
  410. auto prim = attached_primitive_.lock();
  411. if (prim != nullptr) {
  412. (void)prim->AddAttr(attr_name, converted_ret);
  413. }
  414. }
  415. void PrimitivePyAdapter::DelPyAttr(const py::str &name) {
  416. (void)attrs_.erase(name);
  417. auto prim = attached_primitive_.lock();
  418. if (prim != nullptr) {
  419. (void)prim->DelAttr(name);
  420. }
  421. }
  422. py::dict PrimitivePyAdapter::GetAttrDict() {
  423. auto prim = attached_primitive_.lock();
  424. if (prim != nullptr) {
  425. return prim->GetAttrDict();
  426. }
  427. py::dict attr_dict;
  428. for (auto &attr : attrs_) {
  429. attr_dict[py::str(attr.first)] = ValueToPyData(attr.second);
  430. }
  431. return attr_dict;
  432. }
  433. void PrimitivePyAdapter::set_prim_type(const PrimType t) {
  434. prim_type_ = t;
  435. auto prim = attached_primitive_.lock();
  436. if (prim != nullptr) {
  437. prim->set_prim_type(t);
  438. }
  439. }
  440. void PrimitivePyAdapter::set_const_prim(bool is_const_prim) {
  441. is_const_prim_ = is_const_prim;
  442. auto prim = attached_primitive_.lock();
  443. if (prim != nullptr) {
  444. prim->set_const_prim(is_const_prim);
  445. }
  446. }
  447. void PrimitivePyAdapter::set_const_input_indexes(const std::vector<size_t> &const_input_indexes) {
  448. const_input_indexes_ = const_input_indexes;
  449. auto prim = attached_primitive_.lock();
  450. if (prim != nullptr) {
  451. prim->set_const_input_indexes(const_input_indexes);
  452. }
  453. }
  454. void PrimitivePyAdapter::set_signatures(const std::vector<Signature> &signatures) {
  455. signatures_ = signatures;
  456. auto prim = attached_primitive_.lock();
  457. if (prim != nullptr) {
  458. prim->set_signatures(signatures);
  459. }
  460. }
  461. void PrimitivePyAdapter::set_hook(const py::function &hook) {
  462. hook_ = hook;
  463. auto prim = attached_primitive_.lock();
  464. if (prim != nullptr) {
  465. prim->set_hook(hook);
  466. }
  467. }
  468. void PrimitivePyAdapter::set_instance_name(const std::string &s) {
  469. instance_name_ = s;
  470. auto prim = attached_primitive_.lock();
  471. if (prim != nullptr) {
  472. prim->set_instance_name(s);
  473. }
  474. }
  475. void PrimitivePyAdapter::set_attached_primitive(const PrimitivePyPtr &prim) {
  476. if (attached_primitive_.lock() != nullptr) {
  477. MS_LOG(EXCEPTION) << "PrimitivePyAdapter can't attach to multi Primitive.";
  478. }
  479. MS_EXCEPTION_IF_NULL(prim);
  480. attached_primitive_ = prim;
  481. }
  482. REGISTER_PYBIND_DEFINE(Primitive_, ([](const py::module *m) {
  483. (void)py::enum_<PrimType>(*m, "prim_type", py::arithmetic())
  484. .value("unknown", PrimType::kPrimTypeUnknown)
  485. .value("builtin", PrimType::kPrimTypeBuiltIn)
  486. .value("py_infer_shape", PrimType::kPrimTypePyInfer)
  487. .value("user_custom", PrimType::kPrimTypeUserCustom)
  488. .value("py_infer_check", PrimType::kPrimTypePyCheck);
  489. (void)py::class_<PrimitivePyAdapter, std::shared_ptr<PrimitivePyAdapter>>(*m, "Primitive_")
  490. .def_readonly(PYTHON_PRIMITIVE_FLAG, &PrimitivePyAdapter::parse_info_)
  491. .def(py::init<py::str &>())
  492. .def("add_attr", &PrimitivePyAdapter::AddPyAttr, "add primitive attr")
  493. .def("del_attr", &PrimitivePyAdapter::DelPyAttr, "del primitive attr")
  494. .def("get_attr_dict", &PrimitivePyAdapter::GetAttrDict, "get primitive attr")
  495. .def("set_prim_type", &PrimitivePyAdapter::set_prim_type, "Set primitive type.")
  496. .def("set_const_prim", &PrimitivePyAdapter::set_const_prim, "Set primitive is const.")
  497. .def("set_const_input_indexes", &PrimitivePyAdapter::set_const_input_indexes,
  498. "Set primitive const input indexes.")
  499. .def("set_signatures", &PrimitivePyAdapter::set_signatures,
  500. "Set primitive inputs signature.")
  501. .def("register_hook", &PrimitivePyAdapter::set_hook, "Set primitive hook function.")
  502. .def("set_instance_name", &PrimitivePyAdapter::set_instance_name,
  503. "Set primitive instance name.");
  504. }));
  505. } // namespace mindspore