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 distutils.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
|