-
Notifications
You must be signed in to change notification settings - Fork 0
/
setup.py
71 lines (62 loc) · 2.11 KB
/
setup.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
#!/usr/bin/env python3
from setuptools import setup, Extension
import numpy as np
from os import environ
def main():
"""
Compile time options (set as environment variables):
MKL: use intel math kernel library
CBLAS: use C bindings of BLAS
LAPACK_C: use C bindings of LAPACK
PARALLEL: use some parallelization (not always helpful)
PARALLEL_EXTRAPOLATION: use parallel matrix extrapolation when mitigating
Floquet matrix truncation effects
PARALLEL_EXTRA_DIMS: parallelize calculation for voltage-shifted replicas
DEBUG: enable debugging options
ANALYZE: enable benchmark/analysis options
"""
compiler_args = ["-O3", "-Wall", "-Wextra", "-std=c11"]
linker_args = []
include_dirs = [np.get_include()]
library_dirs = []
libraries = []
if "MKL" in environ:
libraries += ["mkl_rt"]
include_dirs += ["/opt/intel/mkl/include"]
library_dirs += ["/opt/intel/mkl/lib/intel64"]
compiler_args += ["-DCBLAS"]
compiler_args += ["-DMKL"]
else:
libraries += ["lapack"]
if "CBLAS" in environ:
compiler_args += ["-DCBLAS"]
libraries += ["cblas"]
else:
libraries += ["blas"]
if "LAPACK_C" in environ:
compiler_args += ["-DLAPACK_C"]
parallel_modifiers = ("PARALLEL", "PARALLEL_EXTRAPOLATION", "PARALLEL_EXTRA_DIMS")
need_omp = False
for modifier in parallel_modifiers:
if modifier in environ:
compiler_args += ["-D" + modifier]
need_omp = True
if need_omp:
compiler_args += ["-fopenmp"]
linker_args += ["-fopenmp"]
if "DEBUG" in environ:
compiler_args += ["-DDEBUG"]
if "ANALYZE" in environ:
compiler_args += ["-DANALYZE"]
module = Extension(
"frtrg.rtrg_c",
sources=["src/frtrg/rtrg_c.c"],
include_dirs=include_dirs,
library_dirs=library_dirs,
libraries=libraries,
extra_compile_args=compiler_args,
extra_link_args=linker_args,
)
setup(ext_modules=[module])
if __name__ == "__main__":
main()