imex-runner.py: Add --requires option to control conditional execution. (#493)

* imex-runner.py: Add --requires option to control conditional execution.
vulkan-runner, l0-runtime, sycl-runtime can used with --requires
to conditionally execute the script.
For example, if --requires=l0-runtime is given and l0 runtime is not
available, imex-runner.py will skip execution.
* Replace += with .append() to make code more consistent.
* Windows shared libraries (.dll) are placed under "bin" not "lib"

Author: silee2

CMakeLists.txt

@@ -109,7 +109,11 @@ set(LLVM_LIT_ARGS "-sv" CACHE STRING "lit default options")
 set(IMEX_SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR})
 # LLVM_EXTERNAL_PROJECTS build puts library, executables and tools in LLVM's CMAKE_BINARY_DIR
 set(IMEX_BINARY_DIR ${CMAKE_BINARY_DIR})
-set(IMEX_LIB_DIR ${IMEX_BINARY_DIR}/lib)
+if(WIN32)
+    set(IMEX_LIB_DIR ${IMEX_BINARY_DIR}/bin)
+else()
+    set(IMEX_LIB_DIR ${IMEX_BINARY_DIR}/lib)
+endif()
 
 include(sanitizers)
 
@@ -120,7 +124,7 @@ add_subdirectory(test)
 
 option(IMEX_INCLUDE_DOCS "Generate build targets for the IMEX docs." ON)
 if (IMEX_INCLUDE_DOCS)
-  add_subdirectory(docs)
+    add_subdirectory(docs)
 endif()
 
 install(DIRECTORY ${PROJECT_SOURCE_DIR}/include/mlir

test/Integration/Dialect/Linalg/Vulkan/linalg_addf.mlir

@@ -1,4 +1,4 @@
-// RUN: %python_executable %imex_runner -i %s --pass-pipeline-file=%p/linalg-to-gpu-vulkan.pp --runner mlir-vulkan-runner --shared-libs=%vulkan_runtime_wrappers,%mlir_runner_utils --entry-point-result=void | FileCheck %s --check-prefix=VULKAN
+// RUN: %python_executable %imex_runner -i %s --pass-pipeline-file=%p/linalg-to-gpu-vulkan.pp --runner mlir-vulkan-runner --shared-libs=%vulkan_runtime_wrappers,%mlir_runner_utils --entry-point-result=void --filecheck --check-prefix=VULKAN
 
 #map = affine_map<(d0, d1) -> (d0, d1)>
 module {

tools/imex-runner/imex-runner.py.in

@@ -13,8 +13,11 @@
 
 """
 Run a mlir file through imex-opt using a user provided pass pipeline
-and run resulting output through mlir-runner.
-The pass pipeline is provided either in a file or as a string; expected
+and optionally(On by default and turned off by option --no-mlir-runner)
+run resulting output through an mlir runner
+(Default is mlir-cpu-runner and can be overridden by option --runner).
+After that, optionally(Option --filecheck), output can be passed to FileCheck utility.
+The pass pipeline for imex-opt is provided either in a file or as a string; expected
 syntax is similar to mlir-opt's syntax.
 
 If provided in a file passes can be separted by newline or ',' and
@@ -32,10 +35,14 @@ func.func(scf-bufferize,
 Pass pipelines privded as strings will not be modified.
 
 All unknown arguments will be forwarded to mlir-runner. Currently there is no
-option to forward user-provided args to imex-opt.
+option to forward user-provided args directly to imex-opt.
+Options --imex-print-before-all and --imex-print-after-all gets converted
+into --mlir-print-ir-before-all, --mlir-print-ir-after-all and forwarded to
+imex-opt to enable printing IR before and after running passes.
+Option --check-prefix gets forwarded to FileCheck.
 
 To use this in a lit test, you can do something similar to
-`// RUN: %python_executable %imex-runner -i %s --pass-pipeline-file=%p/ptensor.pp -e main -entry-point-result=void --shared-libs=%mlir_c_runner_utils --shared-libs=%mlir_runner_utils | FileCheck %s`
+`// RUN: %python_executable %imex-runner -i %s --pass-pipeline-file=%p/ptensor.pp -e main -entry-point-result=void --shared-libs=%mlir_c_runner_utils,%mlir_runner_utils --filecheck`
 """
 
 import os, sys, re
@@ -44,23 +51,52 @@ import argparse
 import subprocess
 
 imex_enable_vulkan_runner = @IMEX_ENABLE_VULKAN_RUNNER@
+imex_enable_l0_runtime = @IMEX_ENABLE_L0_RUNTIME@
+imex_enable_sycl_runtime = @IMEX_ENABLE_SYCL_RUNTIME@
+
 runner_choices = ['mlir-cpu-runner']
+enabled_features = []
+all_features = ['vulkan-runner', 'l0-runtime', 'sycl-runtime']
 if imex_enable_vulkan_runner:
-    runner_choices += ['mlir-vulkan-runner']
+    runner_choices.append('mlir-vulkan-runner')
+    enabled_features.append('vulkan-runner')
+if imex_enable_l0_runtime:
+    enabled_features.append('l0-runtime')
+if imex_enable_sycl_runtime:
+    enabled_features.append('sycl-runtime')
+
+class SplitArgs(argparse.Action):
+    def __call__(self, parser, namespace, values, option_string=None):
+        setattr(namespace, self.dest, values.split(','))
 
 parser = argparse.ArgumentParser(
-    description="run imex-opt and pipe into mlir-runner"
+    description="Run imex-opt, optionally pipe result into selected mlir runner (default: mlir-cpu-runner) and then optionally pipe output into FileCheck"
 )
 parser.add_argument("--input-file", "-i", default=None, help="input MLIR file")
 parser.add_argument("--pass-pipeline-file", "-f", default=None, help="file defining pass pipeline")
 parser.add_argument("--pass-pipeline", "-p", default=None, help="pass pipeline (string)")
 parser.add_argument("--imex-print-before-all", "-b", action='store_true', dest='before', help="print ir before all passes")
 parser.add_argument("--imex-print-after-all", "-a", action='store_true', dest='after', help="print ir after all passes")
 parser.add_argument("--no-mlir-runner", "-n", action='store_true', dest='no_mlir_runner', help="skip mlir runner")
+parser.add_argument("--filecheck", action='store_true', dest='run_filecheck', help="run FileCheck")
+parser.add_argument("--check-prefix", default=None, help="change check prefix (default: CHECK) used by FileCheck")
 parser.add_argument("--runner", "-r", default="mlir-cpu-runner", choices=runner_choices, help="mlir runner name")
+parser.add_argument('--requires', default=enabled_features, action=SplitArgs, help="skip if any of the required in the comma separated list is missing.")
 
 args, unknown = parser.parse_known_args()
 
+# Check if requirements are valid
+required_features = args.requires
+if not set(required_features) <= set(all_features):
+    print('Invalid option passed to --requires')
+    print('your options were: ', *required_features)
+    print('valid options are: ', *all_features)
+    exit(1)
+
+# Skip runner if requirements are not met be enabled features
+if not set(required_features) <= set(enabled_features):
+    exit(0)
+
 ppipeline = None
 if args.pass_pipeline_file:
     with open(args.pass_pipeline_file, 'r') as infile:
@@ -82,6 +118,36 @@ if args.pass_pipeline_file:
 imex_binary_dir = '@IMEX_BINARY_DIR@'
 llvm_binary_dir = '@LLVM_BINARY_DIR@'
 
+def run_pipeline(cmds):
+    """
+    Run commands in PIPE, return the last process in chain.
+    stdout of intermediate commands are sent to PIPE.
+    stdout of last command is not sent to PIPE.
+    """
+    first_cmd, *rest_cmds = cmds
+    if len(cmds) == 1:
+        procs = [subprocess.Popen(first_cmd)]
+    else:
+        procs = [subprocess.Popen(first_cmd, stdout=subprocess.PIPE)]
+        if len(cmds) == 2:
+            last_stdout = procs[-1].stdout
+            cmd = rest_cmds[-1]
+            proc = subprocess.Popen(cmd, stdin=last_stdout)
+            procs.append(proc)
+        else:
+            for cmd in rest_cmds[:-1]:
+                last_stdout = procs[-1].stdout
+                proc = subprocess.Popen(cmd, stdin=last_stdout, stdout=subprocess.PIPE)
+                procs.append(proc)
+            last_stdout = procs[-1].stdout
+            cmd = rest_cmds[-1]
+            proc = subprocess.Popen(cmd, stdin=last_stdout)
+            procs.append(proc)
+    return procs[-1]
+
+# All commands to create a pipeline
+cmds = []
+
 # build imex-opt command
 cmd = [os.path.normpath(os.path.join(imex_binary_dir, 'bin', 'imex-opt'))]
 if ppipeline:
@@ -94,14 +160,28 @@ if args.before:
     cmd.append(f'--mlir-print-ir-before-all')
 if args.after:
     cmd.append(f'--mlir-print-ir-after-all')
-# run and feed into pipe
-if args.no_mlir_runner:
-    subprocess.run(cmd)
-else:
-    p1 = subprocess.Popen(cmd, stdout=subprocess.PIPE)
-    # build mlir-opt command: all unknown args will be passed to imex-opt
+cmds.append(cmd)
+
+# build runner command
+if not args.no_mlir_runner:
+    # build runner command: all unknown args will be passed to the runner
     cmd = [os.path.normpath(os.path.join(llvm_binary_dir, 'bin', args.runner))] + unknown
-    # get stdout from imex-opt and pipe into mlir-opt
-    p2 = subprocess.Popen(cmd, stdin=p1.stdout)
-    p1.wait()
-    p2.wait()
+    cmds.append(cmd)
+
+# build FileCheck command
+if args.run_filecheck:
+    cmd = [os.path.normpath(os.path.join(llvm_binary_dir, 'bin', 'FileCheck'))]
+    if args.check_prefix:
+        cmd.append(f'--check-prefix={args.check_prefix}')
+    if not args.input_file:
+        print('Error: --input-file is required for filecheck')
+        exit(1)
+    cmd.append(args.input_file)
+    cmds.append(cmd)
+
+# Run command pipeline
+last_proc = run_pipeline(cmds)
+
+# Wait for last command completion and exit with returncode
+last_proc.wait()
+exit(last_proc.returncode)