Tuning script

tools/tune_ck.py

+import os, json, subprocess, tempfile, sys, argparse
+
+
+def benchmark_one(config, tuning):
+    with tempfile.NamedTemporaryFile(mode="w+") as tf:
+        b = {
+            'settings': {'iterations': 100},
+            'compile_op': {
+                'name': 'ck_gemm',
+                'tuning_val': tuning,
+                'inputs': config
+            }
+        }
+        json.dump(b, tf)
+        cp = subprocess.run('./bin/gpu-driver {}'.format(tf.name), capture_output=True)
+        for line in cp.stdout.decode().split("\n"):
+            s = line.strip()
+            if not s:
+                continue
+            fields = s.split(',')
+            dtime = fields[-1].strip()
+            return float(dtime[:-2])
+    return sys.float_info.max
+
+def benchmark(config, size):
+    times = [benchmark_one(config, i) for i in range(size)]
+    return times.index(max(times))
+
+def benchmark_log(f):
+    result = []
+    for line in open(f).readlines():
+        line = line.strip()
+        if not line.startswith('ck_gemm:'):
+            continue
+        line = line[len('ck_gemm:'):].strip()
+        config = json.loads(line)
+        tuned = benchmark(config, 13)
+        result.append([config, tuned])
+    return result
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Simple tuner for CK gemms")
+    parser.add_argument('--log', '-l',
+                        type=str,
+                        metavar='file',
+                        help='Path to logfile')
+    parser.add_argument('--out', '-o',
+                        type=str,
+                        metavar='file',
+                        help='Output json file to save tunings')
+    args = parser.parse_args()
+    return args
+
+def run(args):
+    tuned = benchmark_log(args.log)
+    json.dump(tuned, open(args.out, 'w+'))
+
+run(parse_args())