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
|
"""Helpers for OpenMP support during the build."""
# This code is adapted for a large part from the astropy openmp helpers, which
# can be found at: https://github.com/astropy/extension-helpers/blob/master/extension_helpers/_openmp_helpers.py # noqa
import os
import sys
import textwrap
import subprocess
from setuptools.errors import CompileError, LinkError
from pre_build_helpers import compile_test_program
def get_openmp_flag(compiler):
"""Get openmp flags for a given compiler"""
if hasattr(compiler, "compiler"):
compiler = compiler.compiler[0]
else:
compiler = compiler.__class__.__name__
if sys.platform == "win32" and ("icc" in compiler or "icl" in compiler):
omp_flag = ["/Qopenmp"]
elif sys.platform == "win32":
omp_flag = ["/openmp"]
elif sys.platform in ("darwin", "linux") and "icc" in compiler:
omp_flag = ["-qopenmp"]
elif sys.platform == "darwin" and "openmp" in os.getenv("CPPFLAGS", ""):
omp_flag = []
else:
# Default flag for GCC and clang:
omp_flag = ["-fopenmp"]
if sys.platform.startswith("darwin"):
omp_flag += ["-Xpreprocessor", "-lomp"]
return omp_flag
def check_openmp_support():
"""Check whether OpenMP test code can be compiled and run"""
code = textwrap.dedent(
"""\
#include <omp.h>
#include <stdio.h>
int main(void) {
#pragma omp parallel
printf("nthreads=%d\\n", omp_get_num_threads());
return 0;
}
"""
)
extra_preargs = os.getenv("LDFLAGS", None)
if extra_preargs is not None:
extra_preargs = extra_preargs.strip().split(" ")
extra_preargs = [
flag
for flag in extra_preargs
if flag.startswith(("-L", "-Wl,-rpath", "-l"))
]
extra_postargs = get_openmp_flag
try:
output, compile_flags = compile_test_program(
code, extra_preargs=extra_preargs, extra_postargs=extra_postargs
)
if output and "nthreads=" in output[0]:
nthreads = int(output[0].strip().split("=")[1])
openmp_supported = len(output) == nthreads
elif "PYTHON_CROSSENV" in os.environ:
# Since we can't run the test program when cross-compiling
# assume that openmp is supported if the program can be
# compiled.
openmp_supported = True
else:
openmp_supported = False
except (CompileError, LinkError, subprocess.CalledProcessError):
openmp_supported = False
compile_flags = []
return openmp_supported, compile_flags
|
