Update IREE onnx import to be in sync with Torch-MLIR

compiler/bindings/python/iree/compiler/tools/import_onnx/__main__.py

@@ -15,8 +15,11 @@
   python -m iree.compiler.tools.import_onnx ...
 """
 import argparse
+import os
 from pathlib import Path
+import shutil
 import sys
+import tempfile
 
 try:
     import onnx
@@ -38,8 +41,8 @@
 )
 
 
-def main(args):
-    model_proto = load_onnx_model(args.input_file)
+def main(args: argparse.Namespace):
+    model_proto = load_onnx_model(args)
     context = Context()
     model_info = onnx_importer.ModelInfo(model_proto)
     m = model_info.create_module(context=context).operation
@@ -58,13 +61,64 @@ def main(args):
         print(m.get_asm(assume_verified=not args.no_verify))
 
 
-def load_onnx_model(file_path: Path) -> onnx.ModelProto:
-    raw_model = onnx.load(file_path)
-    inferred_model = onnx.shape_inference.infer_shapes(raw_model)
+def load_onnx_model(args: argparse.Namespace) -> onnx.ModelProto:
+    # Do shape inference two ways.  First, attempt in-memory to avoid redundant
+    # loading and the need for writing a temporary file somewhere.  If that
+    # fails, typically because of the 2 GB protobuf size limit, try again via
+    # files.  See
+    # https://github.com/onnx/onnx/blob/main/docs/PythonAPIOverview.md#shape-inference-a-large-onnx-model-2gb
+    # for details about the file-based technique.
+
+    # Make a temp dir for all the temp files we'll be generating as a side
+    # effect of infering shapes.  For now, the only file is a new .onnx holding
+    # the revised model with shapes.
+    #
+    # TODO: If the program temp_dir is None, we should be using an ephemeral
+    # temp directory instead of a hard-coded path in order to avoid data races
+    # by default.
+    input_dir = os.path.dirname(os.path.abspath(args.input_file))
+    temp_dir = tempfile.TemporaryDirectory(dir=input_dir)
+    temp_dir_path = Path(temp_dir.name)
+
+    # Load the model, with possible external data coming from the default
+    # location, or the location specified on the command line.
+    if args.data_dir is None:
+        raw_model = onnx.load(args.input_file)
+    else:
+        raw_model = onnx.load(args.input_file, load_external_data=False)
+        onnx.load_external_data_for_model(raw_model, args.data_dir)
+
+    # Run the checker to test whether the file is above the threshold for
+    # in-memory shape inference.  If not, go ahead and do the shape inference.
+    try:
+        onnx.checker.check_model(raw_model)
+        inferred_model = onnx.shape_inference.infer_shapes(
+            raw_model, data_prop=args.data_prop
+        )
+        return inferred_model
+    except ValueError:
+        pass
+
+    # Model is too big for in-memory inference: do file-based shape inference
+    # to a temp file.
+    temp_inferred_file = temp_dir_path / "temp-inferred.onnx"
+    onnx.shape_inference.infer_shapes_path(
+        args.input_file, temp_inferred_file, data_prop=args.data_prop
+    )
+
+    # Load the temp file and the external data.
+    inferred_model = onnx.load(temp_inferred_file, load_external_data=False)
+    data_dir = Path(input_dir if args.temp_dir is None else args.data_dir)
+    onnx.load_external_data_for_model(inferred_model, data_dir)
+
+    # Remove the inferred shape file unless asked to keep it
+    if not args.keep_temps:
+        temp_dir.cleanup()
+
     return inferred_model
 
 
-def parse_arguments(argv=None):
+def parse_arguments(argv=None) -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="IREE ONNX import tool")
     parser.add_argument("input_file", help="ONNX protobuf input", type=Path)
     parser.add_argument(
@@ -75,6 +129,30 @@ def parse_arguments(argv=None):
         action="store_true",
         help="Disable verification prior to printing",
     )
+    parser.add_argument(
+        "--data-prop",
+        dest="data_prop",
+        default=True,
+        action=argparse.BooleanOptionalAction,
+        help="Toggle data propogation for onnx shape inference",
+    )
+    parser.add_argument(
+        "--keep-temps", action="store_true", help="Keep intermediate files"
+    )
+    parser.add_argument(
+        "--temp-dir",
+        help="Pre-existing directory in which to create temporary files."
+        ' For example, to place temporaries under the directory "foo/bar",'
+        ' specify --temp-dir=foo/bar.  "foo/bar" must already exist.'
+        " Defaults to the directory of the input file.",
+        type=Path,
+    )
+    parser.add_argument(
+        "--data-dir",
+        help="Path to the base directory of the data."
+        " Defaults to the directory of the input file.",
+        type=Path,
+    )
     args = parser.parse_args(argv)
     return args