Make parsing of models more robust if useOnnxModelTypes is enabled.

This is done by catching exceptions during shape inference (as they happen for example if the model
 uses custom ops) and by falling back to an onnx-mlir based type mapping for some kinds of invalid
 models.

Pass ImportOptions by const reference instead of value to avoid unnecessay copying.

Signed-off-by: Rickert, Jonas <[email protected]>

Author: jorickert

src/Builder/FrontendDialectTransformer.cpp

@@ -159,20 +159,19 @@ using SymbolToOnnxTypeMapping = SymbolMapping<onnx::TypeProto>;
 
 class FrontendGenImpl {
 public:
-  explicit FrontendGenImpl(MLIRContext &context)
-      : context_(context), builder_(&context) {
+  explicit FrontendGenImpl(MLIRContext &context, const ImportOptions &options)
+      : options_(options), context_(context), builder_(&context) {
     module_ = ModuleOp::create(UnknownLoc::get(&context));
     InitHandlerMap();
   }
 
-  ErrorOr<ModuleOp> ImportONNXModel(const onnx::ModelProto &model,
-      ImportOptions options, std::string &errorMessage) {
-    options_ = options;
+  ErrorOr<ModuleOp> ImportONNXModel(
+      const onnx::ModelProto &model, std::string &errorMessage) {
     modelInputShaper_.setShapeInformation(options_.shapeInformation);
     opset_map_ = GetOpsetImportsFromProto(model); // Which opsets to use.
     in_model_functions_ = GetModelLocalFunctions(model);
-    ErrorOr<mlir::func::FuncOp> importGraphResult = importGraph(
-        model.graph(), options_.allowMissingOutputTypes, errorMessage);
+    ErrorOr<mlir::func::FuncOp> importGraphResult =
+        importGraph(model.graph(), errorMessage);
     if (auto ec = importGraphResult.getError()) {
       return ec;
     }
@@ -193,7 +192,7 @@ class FrontendGenImpl {
   }
 
 private:
-  ImportOptions options_;
+  const ImportOptions &options_;
   MLIRContext &context_;
   ModuleOp module_;
   OpBuilder builder_;
@@ -541,7 +540,7 @@ class FrontendGenImpl {
    */
   ErrorOr<FunctionType> importGraph(const onnx::GraphProto &graph,
       Region &region, Operation *op, bool useReturn,
-      bool allowMissingOutputTypes, std::string &errorMessage) {
+      std::string &errorMessage) {
     frontend_symbols_.pushScope(graph.name());
     onnx_type_map.pushScope(graph.name());
     Block *entryBlock = &region.back();
@@ -648,8 +647,8 @@ class FrontendGenImpl {
     // Import the output tensors
     for (const auto &output : graph.output()) {
       std::string dimParams = "";
-      const auto ec = ImportOutputTensor(output, retTys, retVals, errorMessage,
-          allowMissingOutputTypes, &dimParams);
+      const auto ec =
+          ImportOutputTensor(output, retTys, retVals, errorMessage, &dimParams);
       if (ec) {
         errorMessage +=
             "Failed to import output tensor '" + output.name() + "'.\n";
@@ -854,13 +853,22 @@ class FrontendGenImpl {
       // Optional outputs using empty string.
       if (node.output()[i].empty()) {
         outputTypes.emplace_back(builder_.getNoneType());
-      } else if (auto onnxModelType =
-                     ConvertOnnxType(node.output(i), errorMessage)) {
-        if (auto ec = onnxModelType->getError()) {
-          return ec;
-        }
-        outputTypes.emplace_back(*onnxModelType.value());
       } else {
+        auto onnxModelType = ConvertOnnxType(node.output(i), errorMessage);
+        if (onnxModelType) {
+          const auto ec = onnxModelType->getError();
+          if (!ec) {
+            outputTypes.emplace_back(*onnxModelType.value());
+            continue;
+          }
+          if (!options_.allowMissingOutputTypes || ec != InvalidOnnxFormat) {
+            errorMessage += "Failed to get type for '" + node.output(i) + "\n";
+            return ec;
+          }
+          llvm::errs() << "Warning: "
+                       << "Failed to get type type for '" << node.output(i)
+                       << "', falling back to onnx-mlir based mapping.\n";
+        }
         unsigned int j = i;
         // Variadic output is a single ODS result.
         if (variadicOut)
@@ -910,8 +918,8 @@ class FrontendGenImpl {
         region.push_back(new Block);
         OpBuilder::InsertionGuard guard(builder_);
         builder_.setInsertionPointToStart(&region.back());
-        const ErrorOr<FunctionType> importGraphResult = importGraph(attr.g(),
-            region, op, false, options_.allowMissingOutputTypes, errorMessage);
+        const ErrorOr<FunctionType> importGraphResult =
+            importGraph(attr.g(), region, op, false, errorMessage);
         if (auto ec = importGraphResult.getError()) {
           return ec;
         }
@@ -1443,9 +1451,14 @@ class FrontendGenImpl {
         GetOpsetImportsFromProto(functionProto);
 
     // Populates graph.value_info().
-    onnx::shape_inference::InferShapes(&graph, function_opset_map,
-        onnx::OpSchemaRegistry::Instance(),
-        /*options=*/{}, in_model_functions_);
+    try {
+      onnx::shape_inference::InferShapes(&graph, function_opset_map,
+          onnx::OpSchemaRegistry::Instance(),
+          /*options=*/{}, in_model_functions_);
+    } catch (const std::exception &e) {
+      llvm::errs() << "Warning: Caught exception running onnx shape inference: "
+                   << e.what() << "\n";
+    }
 
     // Save caller context, while generating function body.
     ModelLocalFunctionsMap callerModelFunctions;
@@ -1629,14 +1642,14 @@ class FrontendGenImpl {
       const onnx::ValueInfoProto &output,
       llvm::SmallVectorImpl<Type> &ret_types,
       llvm::SmallVectorImpl<Value> &ret_vals, std::string &errorMessage,
-      bool allowMissingType, std::string *dim_params = nullptr) {
+      std::string *dim_params = nullptr) {
     const Value *valPtr = frontend_symbols_.GetByOnnxName(output.name());
     Value val = *valPtr;
 
     ErrorOr<Type> parsedOutputType =
         ImportType(output.type(), errorMessage, dim_params);
     if (auto ec = parsedOutputType.getError()) {
-      if (!allowMissingType || ec != InvalidOnnxFormat) {
+      if (!options_.allowMissingOutputTypes || ec != InvalidOnnxFormat) {
         errorMessage +=
             "Failed to import output type for '" + output.name() + "\n";
         return ec;
@@ -1716,15 +1729,10 @@ class FrontendGenImpl {
   /*!
    * Import ONNX main computation graph.
    * @param graph onnx graph proto.
-   * @param allowMissingOutputTypes If true, type inference will be used to
-   * infer missing output types. This is done by copying the, potential
-   * inferred, output type of the node connected to the output. According to
-   * ONNX, all outputs MUST have types. Therefore this option has to be
-   * considered as a stretch best effort.
    * @return A function corresponding to the imported computation graph.
    */
-  ErrorOr<func::FuncOp> importGraph(const onnx::GraphProto &graph,
-      bool allowMissingOutputTypes, std::string &errorMessage) {
+  ErrorOr<func::FuncOp> importGraph(
+      const onnx::GraphProto &graph, std::string &errorMessage) {
     const std::string &name = "main_graph";
     auto mainFunc = func::FuncOp::create(UnknownLoc(), name,
         /*type=*/builder_.getFunctionType({}, {}), /*attrs=*/{});
@@ -1735,8 +1743,7 @@ class FrontendGenImpl {
 
     ErrorOr<FunctionType> importedFuncType =
         importGraph(graph, /*region=*/mainFunc.getBody(),
-            /*op=*/mainFunc.getOperation(), /*useReturn=*/true,
-            /*allowMissingOutputTypes=*/allowMissingOutputTypes, errorMessage);
+            /*op=*/mainFunc.getOperation(), /*useReturn=*/true, errorMessage);
     if (auto ec = importedFuncType.getError()) {
       errorMessage +=
           "Failed to import main graph, could not get its function type\n";
@@ -1763,7 +1770,7 @@ class FrontendGenImpl {
 
 [[nodiscard]] std::error_code ImportFrontendModelInternal(
     onnx::ModelProto &model, MLIRContext &context,
-    OwningOpRef<ModuleOp> &module, ImportOptions options,
+    OwningOpRef<ModuleOp> &module, const ImportOptions &options,
     std::string &errorMessage) {
   int originVersion = CURRENT_ONNX_OPSET;
   // Get the version of the model
@@ -1799,21 +1806,35 @@ class FrontendGenImpl {
       originVersion < CURRENT_ONNX_OPSET) {
     onnx::ModelProto convertModel =
         onnx::version_conversion::ConvertVersion(model, CURRENT_ONNX_OPSET);
-    if (options.useOnnxModelTypes)
-      onnx::shape_inference::InferShapes(convertModel);
+    if (options.useOnnxModelTypes) {
+      try {
+        onnx::shape_inference::InferShapes(convertModel);
+      } catch (const std::exception &e) {
+        llvm::errs()
+            << "Warning: Caught exception running onnx shape inference: "
+            << e.what() << "\n";
+      }
+    }
     return ImportFrontendModel(
         convertModel, context, module, errorMessage, options);
   } else {
-    if (options.useOnnxModelTypes)
-      onnx::shape_inference::InferShapes(model);
+    if (options.useOnnxModelTypes) {
+      try {
+        onnx::shape_inference::InferShapes(model);
+      } catch (const std::exception &e) {
+        llvm::errs()
+            << "Warning: Caught exception running onnx shape inference: "
+            << e.what() << "\n";
+      }
+    }
     return ImportFrontendModel(model, context, module, errorMessage, options);
   }
   return CompilerSuccess;
 }
 
 [[nodiscard]] std::error_code ImportFrontendModelArray(const void *onnxBuffer,
     int size, MLIRContext &context, OwningOpRef<ModuleOp> &module,
-    std::string &errorMessage, ImportOptions options) {
+    std::string &errorMessage, const ImportOptions &options) {
   onnx::ModelProto model;
 
   bool parse_success = model.ParseFromArray(onnxBuffer, size);
@@ -1855,7 +1876,7 @@ namespace {
 // Return 0 on success, error otherwise.
 [[nodiscard]] std::error_code ImportFrontendModelFile(StringRef model_fname,
     MLIRContext &context, OwningOpRef<ModuleOp> &module,
-    std::string &errorMessage, ImportOptions options) {
+    std::string &errorMessage, const ImportOptions &options) {
   onnx::ModelProto model;
   if (model_fname.ends_with(".onnxtext")) {
     std::string text;
@@ -1912,11 +1933,11 @@ namespace {
 
 [[nodiscard]] std::error_code ImportFrontendModel(const onnx::ModelProto &model,
     MLIRContext &context, OwningOpRef<ModuleOp> &module,
-    std::string &errorMessage, ImportOptions options) {
+    std::string &errorMessage, const ImportOptions &options) {
 
-  detail::FrontendGenImpl myONNXGen(context);
+  detail::FrontendGenImpl myONNXGen(context, options);
   ErrorOr<ModuleOp> importedModule =
-      myONNXGen.ImportONNXModel(model, options, errorMessage);
+      myONNXGen.ImportONNXModel(model, errorMessage);
   if (auto ec = importedModule.getError()) {
     return ec;
   }

src/Builder/FrontendDialectTransformer.hpp

@@ -13,7 +13,6 @@
 #ifndef ONNX_MLIR_FRONTEND_TRANSFORMER_H
 #define ONNX_MLIR_FRONTEND_TRANSFORMER_H
 
-#include <set>
 #include <string>
 
 #include "onnx/onnx_pb.h"
@@ -49,7 +48,11 @@ struct ImportOptions {
   bool allowSorting = true;
   bool useOutputNameAsLocation = false;
 
-  // Allow missing output types and use type inference to determine them.
+  // If true, type inference will be used to
+  // infer missing output types. This is done by copying the, potential
+  // inferred, output type of the node connected to the output. According to
+  // ONNX, all outputs MUST have types. Therefore this option has to be
+  // considered as a stretch best effort.
   bool allowMissingOutputTypes = false;
 
   // Custom shape information for the graph inputs.
@@ -90,7 +93,7 @@ struct ImportOptions {
 [[nodiscard]] std::error_code ImportFrontendModelArray(const void *onnxBuffer,
     int bufferSize, mlir::MLIRContext &context,
     mlir::OwningOpRef<mlir::ModuleOp> &module, std::string &errorMessage,
-    ImportOptions options = ImportOptions());
+    const ImportOptions &options = ImportOptions());
 
 /*!
  *  Import an ONNX model file into the ONNX Dialect.
@@ -100,7 +103,7 @@ struct ImportOptions {
 [[nodiscard]] std::error_code ImportFrontendModelFile(
     llvm::StringRef model_fname, mlir::MLIRContext &context,
     mlir::OwningOpRef<mlir::ModuleOp> &module, std::string &errorMessage,
-    ImportOptions options = ImportOptions());
+    const ImportOptions &options = ImportOptions());
 
 /*!
  *  Import an ONNX model proto into the ONNX Dialect.
@@ -109,7 +112,7 @@ struct ImportOptions {
  */
 [[nodiscard]] std::error_code ImportFrontendModel(const onnx::ModelProto &model,
     mlir::MLIRContext &context, mlir::OwningOpRef<mlir::ModuleOp> &module,
-    std::string &errorMessage, ImportOptions options = ImportOptions());
+    std::string &errorMessage, const ImportOptions &options = ImportOptions());
 
 /*!
  *  TODO: Import models into other extension dialects that cover the

test/mlir/onnx/parse/add_missing_output_types.json

@@ -15,6 +15,14 @@
 // INFERRED:           return [[VAR_0_]], [[VAR_1_]] : tensor<*xf32>, tensor<3x3xf32>
 // INFERRED:         }
 
+// RUN: onnx-mlir --EmitONNXIR --useOnnxModelTypes=true --allowMissingOutputTypes=true --printIR %s  | FileCheck --check-prefix=MODEL-TYPE %s
+// MODEL-TYPE-LABEL:  func.func @main_graph
+// MODEL-TYPE-SAME:   ([[PARAM_0_:%.+]]: tensor<3x1xf32> {onnx.name = "input_a"}, [[PARAM_1_:%.+]]: tensor<1x3xf32> {onnx.name = "input_b"}) -> (tensor<*xf32> {onnx.name = "output_c"}, tensor<3x3xf32> {onnx.name = "output_d"}) {
+// MODEL-TYPE-DAG:       [[VAR_0_:%.+]] = "onnx.Custom"([[PARAM_0_]], [[PARAM_1_]]) {domain_name = "test", function_name = "test.Add", onnx_node_name = "add_node_custom"} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<*xf32>
+// MODEL-TYPE-DAG:       [[VAR_1_:%.+]] = "onnx.Add"([[PARAM_0_]], [[PARAM_1_]]) {onnx_node_name = "add_node"} : (tensor<3x1xf32>, tensor<1x3xf32>) -> tensor<3x3xf32>
+// MODEL-TYPE:           return [[VAR_0_]], [[VAR_1_]] : tensor<*xf32>, tensor<3x3xf32>
+// MODEL-TYPE:         }
+
 {
   "irVersion": "10",
   "producerName": "onnx-example",