MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / make_custom_op

Function make_custom_op

imperative/python/src/ops.cpp:680–731  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

678 }
679
680PyObject* make_custom_op(PyObject* self, PyObject** args, Py_ssize_t nargs) {
681#if MGB_CUSTOM_OP
682 auto op_name = py::handle(args[0]).cast<std::string>();
683 auto kwargs = py::handle(args[1]).cast<py::dict>();
684
685 std::shared_ptr<OpDef> opdef = CustomOpDefFactory::inst()->create_opdef(op_name);
686 auto& custom_opdef = static_cast<mgb::imperative::CustomOpDef&>(*opdef);
687 auto& param = custom_opdef.param();
688
689 for (auto&& kv : kwargs) {
690 std::string param_name = py::handle(kv.first).cast<std::string>();
691 std::string type_name = py::handle(kv.second).ptr()->ob_type->tp_name;
692
693 if (!param.exist(param_name)) {
694 mgb_log_warn(
695 "op %s have no param named %s, ignore this param parsed from "
696 "python",
697 op_name.c_str(), param_name.c_str());
698 continue;
699 }
700
701 auto& param_val = param[param_name];
702 switch (param_val.type()) {
703 CUSTOM_FOR_EACH_BASIC_PARAMTYPE(CUSTOM_CASE_TO_PARSE_NON_LIST)
704 CUSTOM_FOR_STRING_PARAMTYPE(CUSTOM_CASE_TO_PARSE_NON_LIST)
705 CUSTOM_FOR_EACH_BASIC_LIST_PARAMTYPE(CUSTOM_CASE_TO_PARSE_LIST)
706 CUSTOM_FOR_BOOL_LIST_PARAMTYPE(CUSTOM_CASE_TO_PARSE_LIST)
707 CUSTOM_FOR_STRING_LIST_PARAMTYPE(CUSTOM_CASE_TO_PARSE_LIST)
708 case custom::ParamDynType::Device: {
709 param_val =
710 to_custom_device(py::handle(kv.second).cast<mgb::CompNode>());
711 break;
712 }
713 default: {
714 mgb_assert(
715 false, "param dtype of %s:%s is invalid", op_name.c_str(),
716 param_name.c_str());
717 }
718 }
719 }
720
721 PyTypeObject* pytype;
722 pytype = &PyOpType(OpDef);
723 PyObject* obj = pytype->tp_alloc(pytype, 0);
724 reinterpret_cast<PyOp(OpDef)*>(obj)->op = opdef;
725
726 return obj;
727#else
728 mgb_assert(false, "CustomOp disabled, please build megengine with CustomOp open");
729 return nullptr;
730#endif
731}
732
733#undef CUSTOM_CASE_TO_PARSE_LIST
734#undef CUSTOM_CASE_TO_PARSE_NON_LIST

Callers 1

py35_make_custom_opFunction · 0.85

Calls 6

instFunction · 0.85
create_opdefMethod · 0.80
existMethod · 0.80
paramMethod · 0.45
ptrMethod · 0.45
typeMethod · 0.45

Tested by

no test coverage detected