tune_ck.py 3.62 KB
Newer Older
1
import os, json, subprocess, tempfile, sys, argparse, contextlib, multiprocessing, multiprocessing.dummy
Paul's avatar
Paul committed
2

Paul's avatar
Format  
Paul committed
3

Paul's avatar
Fixes  
Paul committed
4
5
6
7
8
9
10
11
12
13
14
15
@contextlib.contextmanager
def tmp_file(dump=None):
    tmp_name = None
    try:
        with tempfile.NamedTemporaryFile(mode='w+', delete=False) as f:
            tmp_name = f.name
            if dump:
                dump(f)
        yield tmp_name
    finally:
        os.unlink(tmp_name)

Paul's avatar
Format  
Paul committed
16

Paul's avatar
Fixes  
Paul committed
17
def pretty_print(obj):
Paul's avatar
Paul committed
18
    print(json.dumps(obj, indent=2))
Paul's avatar
Paul committed
19

Paul's avatar
Format  
Paul committed
20

Paul's avatar
Paul committed
21
22
23
def run_driver(b):
    print(b)
    with tmp_file(lambda tf: json.dump(b, tf)) as tf:
Paul's avatar
Paul committed
24
25
26
        if not os.path.exists('./bin/gpu-driver'):
            print("./bin/gpu-driver not found")
            os.abort()
Paul's avatar
Paul committed
27
28
29
        cp = subprocess.run('./bin/gpu-driver {}'.format(tf),
                            capture_output=True,
                            shell=True)
Paul's avatar
Paul committed
30
31
        print(cp.stderr.decode())
        cp.check_returncode()
Paul's avatar
Paul committed
32
33
34
35
36
37
38
39
        for line in cp.stdout.decode().split("\n"):
            s = line.strip()
            if not s:
                continue
            if not ']: ' in s:
                continue
            yield s.split(']: ')[1].strip()

Paul's avatar
Format  
Paul committed
40

Paul's avatar
Paul committed
41
42
def convert_to_float(s):
    return s[:-2]
Paul's avatar
Format  
Paul committed
43

Paul's avatar
Format  
Paul committed
44

Paul's avatar
Paul committed
45
46
47
48
def get_device_time(s):
    fields = s.split(',')
    return convert_to_float(fields[-1].strip())

Paul's avatar
Format  
Paul committed
49

50
51
52
53
54
55
56
57
58
59
60
61
62
def run_driver_ck(config, tuning, iterations):
    b = {
        'settings': {
            'iterations': iterations
        },
        'compile_op': {
            'name': 'ck_gemm',
            'check': True,
            'tuning_val': tuning,
            'inputs': config
        }
    }
    return run_driver(b)
Paul's avatar
Format  
Paul committed
63

Paul's avatar
Format  
Paul committed
64

Paul's avatar
Paul committed
65
def benchmark_ck(config, tuning):
Paul's avatar
Paul committed
66
    try:
67
        for line in run_driver_ck(config, tuning, 100):
Paul's avatar
Paul committed
68
69
            dtime = get_device_time(line)
            print(dtime)
Paul's avatar
Paul committed
70
            return float(dtime)
Paul's avatar
Paul committed
71
72
        print("Failed")
        sys.exit(1)
Paul's avatar
Paul committed
73
    except:
Paul's avatar
Paul committed
74
        return sys.float_info.max
Paul's avatar
Paul committed
75

Paul's avatar
Format  
Paul committed
76

Paul's avatar
Paul committed
77
def benchmark(config, size):
Paul's avatar
Paul committed
78
    times = [benchmark_ck(config, i) for i in range(size)]
Paul's avatar
Use min  
Paul committed
79
    return times.index(min(times))
Paul's avatar
Paul committed
80

Paul's avatar
Format  
Paul committed
81

Paul's avatar
Paul committed
82
def parse_log(f):
Paul's avatar
Paul committed
83
84
85
86
87
88
    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)
Paul's avatar
Paul committed
89
90
        yield config

Paul's avatar
Format  
Paul committed
91

92
93
94
95
96
97
def precompile(x):
    try:
        list(run_driver_ck(x[0], x[1], 0))
    except:
        pass

Paul's avatar
Format  
Paul committed
98

99
100
101
102
def precompile_log(f, n):
    solutions = ((config, i) for config in parse_log(f) for i in range(n))
    with multiprocessing.Pool(24) as p:
        list(p.imap(precompile, solutions))
Paul's avatar
Format  
Paul committed
103

Paul's avatar
Format  
Paul committed
104

105
def benchmark_log(f, n):
Paul's avatar
Paul committed
106
107
    result = []
    for config in parse_log(f):
108
        tuned = benchmark(config, n)
Paul's avatar
Paul committed
109
        print("Tuned:", tuned)
Paul's avatar
Paul committed
110
111
112
113
114
        result.append([config, tuned])
    return result


def parse_args():
Paul's avatar
Format  
Paul committed
115
116
117
    parser = argparse.ArgumentParser(description="Simple tuner for CK gemms")
    parser.add_argument('--log',
                        '-l',
Paul's avatar
Paul committed
118
119
120
                        type=str,
                        metavar='file',
                        help='Path to logfile')
Paul's avatar
Format  
Paul committed
121
122
    parser.add_argument('--out',
                        '-o',
Paul's avatar
Paul committed
123
124
125
                        type=str,
                        metavar='file',
                        help='Output json file to save tunings')
126
127
128
129
    parser.add_argument('--precompile',
                        '-p',
                        action='store_true',
                        help='Precompile kernels first in parallel')
Paul's avatar
Format  
Paul committed
130
    parser.add_argument('-n', type=int, help='Number of instances to tune')
Paul's avatar
Paul committed
131
132
133
    args = parser.parse_args()
    return args

Paul's avatar
Format  
Paul committed
134

Paul's avatar
Paul committed
135
def run(args):
Paul's avatar
Format  
Paul committed
136
    if (args.precompile):
137
        precompile_log(args.log, args.n)
138
    tuned = benchmark_log(args.log, args.n)
Paul's avatar
Paul committed
139
140
    json.dump(tuned, open(args.out, 'w+'))

Paul's avatar
Format  
Paul committed
141

Paul's avatar
Paul committed
142
run(parse_args())