forked from ANLAB-KAIST/NBA
/
compilelib.py
106 lines (98 loc) · 3.72 KB
/
compilelib.py
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
#! /usr/bin/env python3
import os
import re
from itertools import chain
from snakemake.io import dynamic
from snakemake.shell import shell
def joinpath(*args):
return os.path.normpath(os.path.join(*args))
def get_cuda_arch():
'''
Determine currently installed NVIDIA GPU cards by PCI device ID
and match them with the predefined GPU model lists.
It tries to detect all GPUs and include cubins suitable for all GPU
architectures detected.
If your GPU is not detected correctly, update *_DEVICES
files by referring https://pci-ids.ucw.cz/v2.2/pci.ids
and make a pull request!
'''
pci_list = str(shell('lspci -nn', read=True))
supported_archs = ['MAXWELL', 'KEPLER', 'FERMI']
devtypes_found = set()
for devtype in supported_archs:
fname = devtype + '_DEVICES'
with open(fname, 'r') as f:
for line in f:
line = line.strip()
if not line: continue
model, pciid = line.split('\t')
pciid = pciid.replace('0x', '')
if pciid in pci_list:
devtypes_found.add(devtype)
if len(devtypes_found) == 0:
return []
return list(sorted(devtypes_found, key=lambda k: supported_archs.index(k)))
def expand_subdirs(dirlist):
'''
Recursively expands subdirectories for first-level subdirectory lists,
except the root directory indicated as ".".
'''
for idx, dir_ in enumerate(dirlist[:]):
if dir_ in ('.', '..'):
continue
for root, subdirs, files in os.walk(dir_):
for subdir in subdirs:
dirlist.insert(idx + 1, joinpath(root, subdir))
return dirlist
def find_all(dirlist, filepattern):
'''
Retrieves the list of file paths in all the given directories that matches
with the given pattern.
'''
rx = re.compile(filepattern)
results = []
for root, dirs, files in chain(*map(lambda d: os.walk(d, topdown=False), dirlist)):
for fname in files:
if rx.search(fname):
results.append(joinpath(root, fname))
return results
_rx_included_local_header = re.compile(r'"(.+\.(h|hh))"')
def get_includes(srcfile, dynamic_inputs=None, visited=None):
'''
Gets a list of included local header files from the given source file.
(e.g., #include "xxxx.hh")
'''
results = set()
visited = visited if visited else set()
try:
with open(srcfile, 'r') as f:
for line in f:
if line.startswith('#include'):
m = _rx_included_local_header.search(line)
if not m: continue
p = joinpath(os.path.split(srcfile)[0], m.group(1))
if dynamic_inputs and any(di.endswith(p) for di in dynamic_inputs):
p = dynamic(p)
results.add(p)
for fname in results.copy():
if (fname.endswith('.h') or fname.endswith('.hh')) \
and not fname in visited:
visited.add(fname)
results.update(s for s in get_includes(fname, dynamic_inputs, visited))
except FileNotFoundError:
pass
return results
_rx_export_elem_decl = re.compile(r'^EXPORT_ELEMENT\(([a-zA-Z0-9_]+)\)')
def detect_element_def(header_file):
with open(header_file, 'r') as fin:
for line in fin:
m = _rx_export_elem_decl.search(line)
if not m: continue
return m.group(1)
_rx_export_lb_decl = re.compile(r'EXPORT_LOADBALANCER\(([a-zA-Z0-9_]+)\)')
def detect_loadbalancer_def(header_file):
with open(header_file, 'r') as fin:
for line in fin:
m = _rx_export_lb_decl.search(line)
if not m: continue
return m.group(1)