Add support for fat jars to benchmark_cold_compile.py (#2232)

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

1 files changed, 19 insertions, 9 deletions

diff --git a/benchmark/scripts/benchmark_cold_compile.py b/benchmark/scripts/benchmark_cold_compile.py
index 2a0431be..37b41483 100755
--- a/benchmark/scripts/benchmark_cold_compile.py
+++ b/benchmark/scripts/benchmark_cold_compile.py
@@ -31,14 +31,17 @@ def run_firrtl(java, jar, design):
 
 def parseargs():
     parser = argparse.ArgumentParser("Benchmark FIRRTL")
-    parser.add_argument('--designs', type=str, nargs='+',
+    parser.add_argument('--designs', type=str, nargs='+', required=True,
                         help='FIRRTL input files to use as benchmarks')
-    parser.add_argument('--versions', type=str, nargs='+', default=['HEAD'],
-                        help='FIRRTL commit hashes to benchmark')
     parser.add_argument('--iterations', '-N', type=int, default=10,
                         help='Number of times to run each benchmark')
     parser.add_argument('--jvms', type=str, nargs='+', default=['java'],
                         help='JVMs to use')
+    group = parser.add_mutually_exclusive_group()
+    group.add_argument('--versions', type=str, nargs='+', default=['HEAD'],
+                        help='FIRRTL commit hashes to benchmark')
+    group.add_argument('--jars', type=str, nargs='+',
+                        help='FIRRTL fat jars to benchmark')
     return parser.parse_args()
 
 def get_version_hashes(versions):
@@ -83,17 +86,24 @@ def build_firrtl_jars(versions):
         jars[hashcode] = clone_and_build(hashcode)
     return jars
 
-def check_designs(designs):
-    for design in designs:
-        assert os.path.exists(design), '{} must be an existing file!'.format(design)
+
+def check_args_exist(args):
+    for arg in args:
+        assert os.path.exists(arg), '{} must be an existing file!'.format(arg)
 
 
 def main():
     args = parseargs()
     designs = args.designs
-    check_designs(designs)
-    hashes = get_version_hashes(args.versions)
-    jars = build_firrtl_jars(hashes)
+    check_args_exist(designs)
+    if args.jars is None:
+        hashes = get_version_hashes(args.versions)
+        jars = build_firrtl_jars(hashes)
+    else:
+        check_args_exist(args.jars)
+        jars = OrderedDict()
+        for jar in args.jars:
+            jars[os.path.basename(jar)] = jar
     jvms = args.jvms
     N = args.iterations
     info = [['java', 'revision', 'design', 'max heap (MiB)', 'SD', 'runtime (s)', 'SD']]