Merge branch 'main' of github.com:onnx/onnx-mlir into words-matter

docker/onnx-mlir.py

@@ -1,40 +1,71 @@
 #!/usr/bin/env python3
 
-import os
-import re
-import shutil
-import stat
-import subprocess
-import sys
+# SPDX-License-Identifier: Apache-2.0
 
-DOCKER_SOCKET   = '/var/run/docker.sock'
-ONNX_MLIR_IMAGE = 'onnxmlirczar/onnx-mlir'
-WORK_DIR        = '/workdir'
-OUTPUT_DIR      = '/output'
-EMIT_IR_OPTS    = [ '--EmitONNXBasic',
-                    '--EmitONNXIR',
-                    '--EmitMLIR',
-                    '--EmitLLVMIR' ]
-EMIT_BIN_OPTS   = [ '--EmitLib',
-                    '--EmitObj',
-                    '--EmitJNI' ]
+########################## onnx-mlir.py ########################################
+#
+# Copyright 2022 The IBM Research Authors.
+#
+################################################################################
+#
+# This file scans for certain patterns (listed below) and generate an md table,
+# which list the operations supported, and optionally the unsupported operations.
+# Among the options, we can also list the TODOs in the table.
+# Invoke with the `-h` argument to list options.
+#
+# Limitation: currently handle at most one OP/LIMIT/TODO line per operation.
+# Script currently invoked by the `onnx_mlir_supported_ops` make target.
+#
+################################################################################
 
 # When running onnx-mlir inside a docker container, the directory
 # containing the input ONNX model file must be mounted into the
 # container for onnx-mlir to read the input and generate the output.
 #
 # This convenient script will do that automatically to make it
 # as if you are running onnx-mlir directly on the host.
+
+import os
+import re
+import shutil
+import stat
+import subprocess
+import sys
+
+DOCKER_SOCKET    = '/var/run/docker.sock'
+ONNX_MLIR_IMAGE  = 'onnxmlirczar/onnx-mlir'
+KNOWN_INPUT_TYPE = ( '.onnx', '.json', '.mlir' )
+
+mount_dirs       = []
+mount_args       = []
+onnx_mlir_args   = []
+
+# mount host path into container
+def mount_path(path):
+    global mount_dirs, mount_args, onnx_mlir_args
+
+    p = os.path.abspath(path)
+    d = os.path.dirname(p)
+    f = os.path.basename(p)
+
+    # Haven't seen this directory before
+    if not d in mount_dirs:
+        mount_dirs += [ d ]
+        mount_args += [ '-v', d + ':' + d ]
+        onnx_mlir_args += [ p ]
+    else:
+        onnx_mlir_args += [ path ]
+
 def main():
     # Make sure docker client is installed
     if not shutil.which('docker'):
         print('docker client not found')
-        return
+        sys.exit(1)
 
     # Make sure docker daemon is running
     if not stat.S_ISSOCK(os.stat(DOCKER_SOCKET).st_mode):
         print('docker daemon not running')
-        return
+        sys.exit(1)
 
     # Pull the latest onnxmlirczar/onnx-mlir image, if image
     # is already up-to-date, pull will do nothing.
@@ -48,112 +79,59 @@ def main():
         print(line if re.match('^([0-9a-f]{12})|Error', line) else '',
               end='', flush=True)
     proc.wait()
-    if (proc.returncode != 0):
+    if proc.returncode:
         print("docker pull failed")
-        return
-
-    # Go through the command line options and locate the
-    # input ONNX model file.
-    argi  = 0
-    ionnx = None
-    argo  = 0
-    obase = None
-    argv  = sys.argv
-    argc  = len(sys.argv)
-
-    for i, arg in enumerate(argv):
-        # File specified on the first argument, defaults to --EmitLib
-        if i == 1 and not argv[i].startswith('-'):
-            argi = i
-            ionnx = argv[i]
-        # If a file is not specified on the first argument, it must be
-        # specified after a valid --EmitXXX option.
-        elif (arg in EMIT_IR_OPTS+EMIT_BIN_OPTS and i < argc-1 and
-              not argv[i+1].startswith('-')):
-            # File specified more than once, treat as not specified
-            if ionnx:
-                sys.exit("Too many --EmitXXX options")
-            if (arg in EMIT_BIN_OPTS and sys.platform != 'linux'):
-                print(('Warning: host {} is not linux, ' +
-                       'output not directly usable').format(sys.platform))
-            argi = i + 1
-            ionnx = argv[argi]
-        elif (arg == "-o" and i < argc-1 and
-              not argv[i+1].startswith('-')):
-            if obase:
-                sys.exit("Too many -o options")
-            argo = i + 1
-            obase = argv[argo]
+        sys.exit(proc.returncode)
 
     # Prepare the arguments for docker run
     args = [ 'docker', 'run', '--rm', '-ti' ]
 
-    # Construct the mount option if an input ONNX model file is found
-    if ionnx:
-        p = os.path.abspath(ionnx)
-        d = os.path.dirname(p)
-        f = os.path.basename(p)
-
-        # Add the mount option, directory containing the input
-        # ONNX model file will be mounted under /workdir inside
-        # the container. If /workdir doesn't exist, it will be
-        # created.
-        args.append('-v')
-        args.append(d + ':' + WORK_DIR + d)
-
-        # Change directory into /workdir
-        #args.append('-w')
-        #args.append(WORK_DIR)
-
-        # Rewrite the original input ONNX model file, which will
-        # reside under /workdir inside the container.
-        argv[argi] = WORK_DIR + p
-
-    # Construct the mount option if -o is specified
-    if obase:
-        # Check invalid -o values such as ".", "..", "/.", "./", etc.
-        if re.match('(.*/)*\.*$', obase):
-            sys.exit("Invalid value for -o option")
-
-        p = os.path.abspath(obase)
-        d = os.path.dirname(p)
-        f = os.path.basename(p)
-
-        # Add the mount option, directory containing the output
-        # files will be mounted under /output inside the container.
-        # If /output/... doesn't exist, it will be
-        # created.
-        args.append('-v')
-        args.append(d + ':' + OUTPUT_DIR + d)
-
-        # Rewrite the original output basename, which will
-        # reside under /output inside the container.
-        argv[argo] = OUTPUT_DIR + p
+    # Go through the command line options and locate the known
+    # file types. For each file located, construct a docker mount
+    # option that mounts the host directory into the container.
+    #
+    # Also do the same for the output path specified by the -o
+    # option.
+    argv  = sys.argv
+    argc  = len(sys.argv)
+
+    global mount_dirs, mount_args, onnx_mlir_args
+
+    verbose = False
+    for i in range(1, argc):
+        if argv[i].endswith(KNOWN_INPUT_TYPE):
+            mount_path(argv[i])
+        elif argv[i-1] == '-o' and not argv[i].startswith('-'):
+            mount_path(argv[i])
+        elif argv[i] == '-v':
+            verbose = True
+            onnx_mlir_args += [ argv[i] ]
+        else:
+            onnx_mlir_args += [ argv[i] ]
 
     # Add effective uid and gid
-    args.append('-u')
-    args.append(str(os.geteuid()) + ':' + str(os.getegid()))
+    args += [ '-u', str(os.geteuid()) + ':' + str(os.getegid()) ]
+
+    # Add mount options
+    args += mount_args
 
     # Add image name
-    args.append(ONNX_MLIR_IMAGE)
+    args += [ ONNX_MLIR_IMAGE ]
 
     # Pass in all the original arguments
-    argv.remove(argv[0])
-    args.extend(argv)
-    # print(args) #debug only
+    args += onnx_mlir_args
+
+    if verbose:
+        print(args)
 
     # Run onnx-mlir in the container
     proc = subprocess.Popen(args,
                             stdout=subprocess.PIPE,
                             stderr=subprocess.STDOUT)
     for line in proc.stdout:
-        # Remove first occurrence of /workdir or /output before printing
-        print(re.sub(WORK_DIR + '|' + OUTPUT_DIR, '',
-                     line.decode('utf-8'), 1),
-              end='', flush=True)
+        print(line.decode('utf-8'), end='', flush=True)
     proc.wait()
-    if (proc.returncode != 0):
-        print(os.strerror(proc.returncode))
+    sys.exit(proc.returncode)
 
 if __name__ == "__main__":
     main()

src/Builder/CMakeLists.txt

@@ -7,6 +7,7 @@
 add_onnx_mlir_library(OMBuilder
   FrontendDialectHelper.cpp
   FrontendDialectTransformer.cpp
+  ModelInputShaper.cpp
 
   LINK_LIBS PUBLIC
   OMHasOnnxSubgraphOpInterface