Merge pull request #404 from Xilinx/jrickert.more_robust_onnx_model_types

jorickert · web-flow · commit d32fa8722079 · 2025-07-28T18:32:43.000+02:00
Make parsing of models more robust if useOnnxModelTypes is enabled.
diff --git a/src/Builder/FrontendDialectTransformer.cpp b/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;
   }
diff --git a/src/Builder/FrontendDialectTransformer.hpp b/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
diff --git a/test/mlir/onnx/parse/add_missing_output_types.json b/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",
