1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
|
#!/usr/bin/python3
"""
Benchmarking script. See --help for details. Compiles benchmark, runs and outputs csv.
"""
import argparse
import os
import subprocess
import sys
import time
_POST_COMPILE_SLEEP = 0.5
_GEMMTYPE = {
'f32': 'SGEMM',
'f64': 'DGEMM',
'c32': 'CGEMM',
'c64': 'ZGEMM',
}
_COMPILE = "cargo rustc --example benchmark --release".split()
_DEFAULT_FEATURES = "constconf cgemm".split()
_EXEC = "./target/release/examples/benchmark"
def bench_loop(args, *, file):
extra_header = ",nc,kc,mc,threads"
print("m,k,n,layout,type,average_ns,minimum_ns,median_ns,samples,gflops", extra_header, sep="", file=file)
ncs = [None] if args.nc is None else args.nc
kcs = [None] if args.kc is None else args.kc
mcs = [None] if args.mc is None else args.mc
for threads in args.threads:
for ty in args.type:
for nc in ncs:
for kc in kcs:
for mc in mcs:
bench_iteration(args.size, ty, nc, kc, mc, threads=threads, file=file, sleep=args.sleep)
def bench_iteration(sizes, ty, nc, kc, mc, *, threads, file, sleep):
features = list(_DEFAULT_FEATURES)
if threads > 0:
features.append("threading")
compile_argv = list(_COMPILE)
compile_argv.append("--features=" + ",".join(features))
file.flush()
env = os.environ.copy()
for value, name in zip([nc, kc, mc], ["nc", "kc", "mc"]):
if value is not None:
env["MATMUL_" + _GEMMTYPE[ty] + "_" + name.upper()] = str(value)
print("Running", " ".join(compile_argv), file=sys.stderr)
subprocess.run(compile_argv, env=env)
time.sleep(_POST_COMPILE_SLEEP)
exec_env = os.environ.copy()
exec_env["MATMUL_NUM_THREADS"] = str(threads)
extra_column = ",".join((str(value) if value is not None else "") for value in [nc, kc, mc, threads])
for size in sizes:
argv = [_EXEC]
argv.extend(["--type", ty])
argv.extend(["--csv", "--layout", "fcc", "--extra-column", extra_column])
argv.extend([str(size)] * 3)
print("Running", " ".join(argv), file=sys.stderr)
subprocess.run(argv, env=exec_env, stdout=file)
time.sleep(sleep)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--type", "-t", type=str, default=["f64"], choices=["f64", "f32", "c64", "c32"], nargs="+")
parser.add_argument("--size", "-s", type=str, nargs="+", required=True,
help="Sizes or size ranges like 16 or 16:64:8")
parser.add_argument("--nc", type=str, nargs="+", help="Sizes or size ranges like 16 or 16:64:8")
parser.add_argument("--kc", type=str, nargs="+", help="Sizes or size ranges like 16 or 16:64:8")
parser.add_argument("--mc", type=str, nargs="+", help="Sizes or size ranges like 16 or 16:64:8")
parser.add_argument("--threads", type=int, default=[0], nargs="+",
help="Thread use. 0: not enabled; 1: enabled but one thread; n: enabled with n threads.")
parser.add_argument("--sleep", type=int, default=1, help="Time to wait between every run")
parser.add_argument("--output", type=str, default=None, help="Output file (csv format)")
args = parser.parse_args()
if len(args.type) > 1 and any([args.nc, args.kc, args.mc]):
print("Warning: Combining type loop with nc, mc, kc might not make sense", file=sys.stderr)
# postprocess nc, kc, mc to parse x:y:z into ranges
for var in ['size', 'nc', 'kc', 'mc']:
arg_values = getattr(args, var)
if arg_values is not None:
new_arg_values = []
for arg_value in arg_values:
if ":" in arg_value:
parts = list(int(v) for v in arg_value.split(":"))
values = list(range(*parts))
new_arg_values.extend(values)
else:
new_arg_values.append(arg_value)
setattr(args, var, new_arg_values)
if args.output is not None:
with open(args.output, "w") as f:
bench_loop(args, file=f)
else:
bench_loop(args, file=sys.stdout)
if __name__ == "__main__":
main()
|
