colin
/
onnx-mlir
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

initial code for handling custom ops (#288) 

* initial code for handling custom ops

* format
This commit is contained in:
Kevin O'Brien 2020-09-11 10:49:15 -04:00 committed by GitHub
parent ac67900baf
commit 1fcf97ef8d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 15 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
src/Builder/FrontendDialectTransformer.cpp
View File

@ -485,18 +485,24 @@ private:
buildOutputAndOperation<mlir::ONNXSliceOp>(node, in, nIn, nOut);
}
void ImportCustomNode(const onnx::NodeProto &node) {
llvm::StringRef opName = node.op_type();
mlir::emitWarning(UnknownLoc(), "Could not find op importer: assuming this "
"represents a custom operator.");
}
void ImportNode(const onnx::NodeProto &node) {
llvm::StringRef opName = node.op_type();
// the following code is generated by gen_doc.py
// refer to Dialect/ONNX/ONNXOps.td for details
// when the input or output of then op does not match the specification,
// the generic operator is used
// one known reeason is the optional input
auto found = import_handler_map_.find(opName.str());
assert(found != import_handler_map_.end() && "Could not find op importer");
(this->*(found->second))(node);
// look up handler for the opName. If not found, create a node
// for a custom op, and issue a warning.
auto handler = import_handler_map_.find(opName.str());
if (handler != import_handler_map_.end()) {
(this->*(handler->second))(node);
} else {
ImportCustomNode(node);
}
}
void InitHandlerMap() {