[CI] Added benchmarking to CI script (#65)

python/bench/run.py

@@ -0,0 +1,41 @@
+import argparse
+import sys
+import os
+import inspect
+import triton
+
+def run_all(result_dir, with_plots, names):
+    if not os.path.exists(result_dir):
+        os.makedirs(result_dir)
+    for mod in os.listdir(os.path.dirname(os.path.realpath(__file__))):
+        # skip non python files
+        if not mod.endswith('.py'):
+            continue
+        # skip file not in provided names
+        if names and names not in mod:
+            continue
+        # skip files that don't start with 'bench_'
+        if not mod.startswith('bench_'):
+            continue
+        print(f'running {mod}...')
+        mod = __import__(os.path.splitext(mod)[0])
+        benchmarks = inspect.getmembers(mod, lambda x: isinstance(x, triton.testing.Mark))
+        for name, bench in benchmarks:
+            curr_dir = os.path.join(result_dir, mod.__name__.replace('bench_', ''))
+            if len(benchmarks) > 1:
+                curr_dir = os.path.join(curr_dir, name.replace('bench_', ''))
+            if not os.path.exists(curr_dir):
+                os.makedirs(curr_dir)
+            bench.run(curr_dir, with_plots)
+
+def main(args):
+    parser = argparse.ArgumentParser(description="Run the benchmark suite.")
+    parser.add_argument("-r", "--result-dir", type=str, default='results', required=False)
+    parser.add_argument("-n", "--names", type=str, default='', required=False)
+    parser.add_argument("-p", "--with-plots", dest='with_plots', action='store_true')
+    parser.set_defaults(feature=False)
+    args = parser.parse_args(args)
+    run_all(args.result_dir, args.with_plots, args.names)
+
+if __name__ == '__main__':
+    main(sys.argv[1:])