Improved installer and added a few arguments

hpvm/scripts/hpvm_installer.py

 #!/usr/bin/env python3
-from pathlib import Path
 from argparse import ArgumentParser, Namespace
+from os import chdir, environ, makedirs
+from pathlib import Path
 from subprocess import CalledProcessError, check_call
-from os import makedirs, chdir, environ
+from typing import List
 
 VERSION = "9.0.0"
 URL = "http://releases.llvm.org"
@@ -12,10 +13,9 @@ CLANG_TARBALL = f"{CLANG_DIR}.tar.xz"
 LLVM_DIR = f"llvm-{VERSION}.src"
 LLVM_TARBALL = f"{LLVM_DIR}.tar.xz"
 
-ROOT_DIR = (Path(__file__).parent / "..").absolute()
-BUILD_DIR = ROOT_DIR / "build"
+ROOT_DIR = (Path(__file__).parent / "..").absolute().resolve()
 TEST_DIR = ROOT_DIR / "test"
-LLVM_LIT = BUILD_DIR / "bin/llvm-lit"
+LLVM_LIT = "bin/llvm-lit"  # rel to build dir
 
 MODEL_PARAMS_DIR = TEST_DIR / "dnn_benchmarks/model_params"
 MODEL_PARAMS_TAR = ROOT_DIR / "model_params.tar.gz"
@@ -41,7 +41,8 @@ PY_PACKAGES = [
     "projects/keras",
 ]
 
-def parse_args():
+
+def parse_args(args=None):
     parser = ArgumentParser(
         "hpvm_installer", description="Script for automatic HPVM installation."
     )
@@ -60,6 +61,13 @@ def parse_args():
         default=2,
         help="How many threads to build with. This argument is relayed on to 'make'. Default: 2",
     )
+    parser.add_argument(
+        "-b",
+        "--build-dir",
+        type=Path,
+        default="build",
+        help=f"Where to create the build directory. Relative to ({ROOT_DIR}). Default: build",
+    )
     parser.add_argument(
         "-t",
         "--targets",
@@ -77,7 +85,18 @@ def parse_args():
     parser.add_argument(
         "--no-params", action="store_true", help="Don't download DNN model parameters"
     )
-    return parser.parse_args()
+    parser.add_argument(
+        "cmake_args",
+        type=str,
+        nargs="*",
+        default="",
+        help="Argument to relay on to CMake. Separate with space and do not include the dashes. "
+        "Example: DCMAKE_BUILD_TYPE=Release DCMAKE_INSTALL_PREFIX=install",
+    )
+    args = parser.parse_args(args)
+    args.build_dir = ROOT_DIR / args.build_dir
+    args.cmake_args = [f"-{arg}" for arg in args.cmake_args]
+    return args
 
 
 def prompt_args():
@@ -99,15 +118,20 @@ def prompt_args():
             return None
         return s
 
+    # Use this to get all default arguments
+    args = parse_args([])
+
     print("No Flags found. Using command line prompts.")
     print("Alternatively, please call this script with -h for all available options.")
+    print("CLI arguments cover more options than this interactive prompt.")
     auto_build = input_with_check(
         "Build and install HPVM automatically? [y/n]: ", parse_yn, "Please enter y or n"
     )
+    args.no_build = not auto_build
     if not auto_build:
-        # Just stuff something in the other fields. We won't need them.
-        return Namespace(no_build=True, parallel="", targets="", run_tests=False)
-    threads = input_with_check(
+        # We no longer need the following fields.
+        return args
+    args.parallel = input_with_check(
         "Number of threads: ", parse_int, "Please enter a positive integer"
     )
     print(
@@ -116,28 +140,36 @@ def prompt_args():
         "If building for multiple targets, seperate options with semicolon:\n"
         "e.g. X86;ARM"
     )
-    targets = input_with_check(
+    args.targets = input_with_check(
         "Build target: ", parse_targets, "Input shouldn't contain space"
     )
 
-    download_weights = input_with_check(
+    print(
+        """Additional arguments to CMake? Split by space and no dashes.
+Example: "DCMAKE_BUILD_TYPE=Release DCMAKE_INSTALL_PREFIX=install".
+Arguments: """
+    )
+    args.cmake_args = input()
+    args.cmake_args = [f"-{arg}" for arg in args.cmake_args.split(" ")]
+
+    args.no_params = not input_with_check(
         "Download DNN weights (recommended)? [y/n]: ", parse_yn, "Please enter y or n"
     )
-    run_tests = input_with_check(
+    args.run_tests = input_with_check(
         "Build and run tests? [y/n]: ", parse_yn, "Please enter y or n"
     )
-    return Namespace(
-        no_build=not auto_build, parallel=threads, targets=targets, run_tests=run_tests,
-        no_params=not download_weights
-    )
+    return args
 
 
 def print_args(args):
     print("Running with the following options:")
-    print(f"  Automated: {not args.no_build}")
+    print(f"  Automated build: {not args.no_build}")
+    print(f"  Build directory: {args.build_dir}")
     print(f"  Threads: {args.parallel}")
     print(f"  Targets: {args.targets}")
+    print(f"  Download DNN weights: {not args.no_params}")
     print(f"  Run tests: {args.run_tests}")
+    print(f"  CMake arguments: {args.cmake_args}")
 
 
 def check_download_llvm_clang():
@@ -183,13 +215,17 @@ def check_download_model_params():
         print(f"Downloading DNN model parameters: {MODEL_PARAMS_TAR}...")
         print(f"=============================")
         check_call([WGET, MODEL_PARAMS_LINK, "-O", MODEL_PARAMS_TAR])
-    print(f"Extracting DNN model parameters {MODEL_PARAMS_TAR} => {MODEL_PARAMS_DIR}...")
+    print(
+        f"Extracting DNN model parameters {MODEL_PARAMS_TAR} => {MODEL_PARAMS_DIR}..."
+    )
     # Decompression is pretty time-consuming so we try to show a progress bar:
     try:
         check_call(f"pv {MODEL_PARAMS_TAR} | tar xz", shell=True)
     except CalledProcessError:
         # Maybe `pv` is not installed. Fine, we'll run without progress bar.
-        print(">> 'pv' is not installed, no progress bar will be shown during decompression.")
+        print(
+            ">> 'pv' is not installed, no progress bar will be shown during decompression."
+        )
         print(">> Decompression ongoing...")
         check_call(["tar", "xzf", MODEL_PARAMS_TAR])
     check_call(["mv", "model_params", MODEL_PARAMS_DIR])
@@ -215,37 +251,30 @@ def link_and_patch():
     chdir("..")
 
 
-def maybe_build(build: bool, nthreads: int, targets: str, build_test_targets: bool):
-    if not build:
-        print(
-            """
-HPVM not installed.
-To complete installation, follow these instructions:
-- Create and navigate to a folder "./build" 
-- Run "cmake ../llvm [options]". Find potential options in README.md.
-- Run "make -j<number of threads> approxhpvm.py" and then "make install"
-For more details refer to README.md.
-"""
-        )
-        return
+def build(
+    build_dir: Path,
+    nthreads: int,
+    targets: str,
+    build_test_targets: bool,
+    cmake_additional_args: List[str],
+):
     print("Now building...")
     print(f"Using {nthreads} threads to build HPVM.")
-    makedirs(BUILD_DIR, exist_ok=True)
+    makedirs(build_dir, exist_ok=True)
 
-    chdir(BUILD_DIR)
+    chdir(build_dir)
     cmake_args = [
         "cmake",
         "../llvm",
         "-DCMAKE_C_COMPILER=gcc",
         "-DCMAKE_CXX_COMPILER=g++",
         f"-DLLVM_TARGETS_TO_BUILD={targets}",
+        *cmake_additional_args,
     ]
     print(f"CMake: {' '.join(cmake_args)}")
     print(f"=============================")
     check_call(cmake_args)
     make_args = ["make", f"-j{nthreads}", *MAKE_TARGETS]
-    if build_test_targets:
-        make_args += MAKE_TEST_TARGETS
     print(f"Make: {' '.join(make_args)}")
     print(f"=============================")
     check_call(make_args)
@@ -262,12 +291,13 @@ def install_py_packages():
         check_call([sys.executable, "-m", "pip", "install", str(package_home)])
 
 
-def run_tests():
-    chdir(BUILD_DIR)
-    # Run regression tests
-    check_call([LLVM_LIT, "-v", TEST_DIR / "regressionTests"])
-    # Run unit tests
-    check_call([LLVM_LIT, "-v", TEST_DIR / "unitTests"])
+def run_tests(build_dir: Path, nthreads: int):
+    chdir(build_dir)
+    make_args = ["make", f"-j{nthreads}", *MAKE_TEST_TARGETS]
+    print(f"Tests: {' '.join(make_args)}")
+    print(f"=============================")
+    check_call(make_args)
+    chdir(ROOT_DIR)
 
 
 def input_with_check(prompt: str, parse, prompt_when_invalid: str):
@@ -290,10 +320,29 @@ def main():
     link_and_patch()
     if not args.no_params:
         check_download_model_params()
-    maybe_build(not args.no_build, args.parallel, args.targets, args.run_tests)
     install_py_packages()
+    if args.no_build:
+        print(
+            """
+HPVM not installed.
+To complete installation, follow these instructions:
+- Create and navigate to a folder "./build" 
+- Run "cmake ../llvm [options]". Find potential options in README.md.
+- Run "make -j<number of threads> approxhpvm.py" and then "make install"
+For more details refer to README.md.
+"""
+        )
+        return
+    else:
+        build(
+            args.build_dir,
+            args.parallel,
+            args.targets,
+            args.run_tests,
+            args.cmake_args,
+        )
     if args.run_tests:
-        run_tests()
+        run_tests(args.build_dir, args.parallel)
     else:
         print("Skipping tests.")