fix custom op bug and add custom op check
This commit is contained in:
parent
c24252b2cc
commit
c42e53ae5b
|
@ -279,6 +279,31 @@ class OpAdapter : public BaseOpAdapter {
|
|||
}
|
||||
|
||||
OutHandler getOutput(const OperatorPtr& op, int index) override {
|
||||
MS_EXCEPTION_IF_NULL(op);
|
||||
if (IsCustomOp(op)) {
|
||||
return getCustomOutput(op, index);
|
||||
}
|
||||
return getNormalOutput(op, index);
|
||||
}
|
||||
|
||||
OutHandler getCustomOutput(const OperatorPtr& op, int index) {
|
||||
MS_EXCEPTION_IF_NULL(op);
|
||||
auto it = cus_output_map_.find(op->GetOpType());
|
||||
if (it == cus_output_map_.end()) {
|
||||
MS_LOG(ERROR) << "OpAdpator(" << op->GetName() << ") has both OUTPUT is not supported!";
|
||||
return OutHandler();
|
||||
}
|
||||
|
||||
std::unordered_map<int, std::string>& output_map = it->second;
|
||||
|
||||
if ((output_map.find(index) != output_map.end())) {
|
||||
return OutHandler(op, output_map[index]);
|
||||
}
|
||||
MS_LOG(ERROR) << "OpAdpator(" << op->GetName() << ") has no OUTPUT index(" << index << ")!";
|
||||
return OutHandler();
|
||||
}
|
||||
|
||||
OutHandler getNormalOutput(const OperatorPtr& op, int index) {
|
||||
MS_EXCEPTION_IF_NULL(op);
|
||||
if (!dyn_output_map_.empty() && !output_map_.empty()) {
|
||||
MS_LOG(ERROR) << "OpAdpator(" << op->GetName() << ") has both OUTPUT and DYN_OUTPUT is not supported!";
|
||||
|
|
|
@ -223,7 +223,13 @@ bool IsCustomPrim(const PrimitivePtr& prim) {
|
|||
return false;
|
||||
}
|
||||
|
||||
return GetValue<bool>(flag);
|
||||
bool is_custom_op = GetValue<bool>(flag);
|
||||
if (!is_custom_op && prim->GetAttr("_custom_op_impl_config_path") != nullptr) {
|
||||
MS_LOG(EXCEPTION) << "The custom op flag is false, but the op information config path is not null, non-custom op "
|
||||
"can not assign the op information config path.";
|
||||
}
|
||||
|
||||
return is_custom_op;
|
||||
}
|
||||
|
||||
bool IsCustomCNode(const AnfNodePtr& anf) {
|
||||
|
|
Loading…
Reference in New Issue