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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
|
import contextlib
import glob
import os
import subprocess
import sys
import tempfile
import warnings
from distutils.ccompiler import new_compiler
from distutils.errors import CompileError, LinkError
from distutils.sysconfig import customize_compiler
CCODE = """
#include <omp.h>
#include <stdio.h>
int main() {
omp_set_num_threads(2);
#pragma omp parallel
printf("nthreads=%d\\n", omp_get_num_threads());
return 0;
}
"""
@contextlib.contextmanager
def stdchannel_redirected(stdchannel, dest_filename):
"""
A context manager to temporarily redirect stdout or stderr
e.g.:
with stdchannel_redirected(sys.stderr, os.devnull):
if compiler.has_function('clock_gettime', libraries=['rt']):
libraries.append('rt')
Code adapted from https://stackoverflow.com/a/17752455/1382869
"""
try:
oldstdchannel = os.dup(stdchannel.fileno())
dest_file = open(dest_filename, "w")
os.dup2(dest_file.fileno(), stdchannel.fileno())
yield
finally:
if oldstdchannel is not None:
os.dup2(oldstdchannel, stdchannel.fileno())
if dest_file is not None:
dest_file.close()
def check_for_openmp():
"""Returns True if local setup supports OpenMP, False otherwise
Code adapted from astropy_helpers, originally written by Tom
Robitaille and Curtis McCully.
"""
# Create a temporary directory
ccompiler = new_compiler()
customize_compiler(ccompiler)
tmp_dir = tempfile.mkdtemp()
start_dir = os.path.abspath(".")
if os.name == "nt":
# TODO: make this work with mingw
# AFAICS there's no easy way to get the compiler distutils
# will be using until compilation actually happens
compile_flag = "-openmp"
link_flag = ""
else:
compile_flag = "-fopenmp"
link_flag = "-fopenmp"
try:
os.chdir(tmp_dir)
with open("test_openmp.c", "w") as f:
f.write(CCODE)
os.mkdir("objects")
# Compile, link, and run test program
with stdchannel_redirected(sys.stderr, os.devnull):
ccompiler.compile(
["test_openmp.c"],
output_dir="objects",
extra_postargs=[compile_flag],
)
ccompiler.link_executable(
glob.glob(os.path.join("objects", "*")),
"test_openmp",
extra_postargs=[link_flag],
)
output = (
subprocess.check_output("./test_openmp")
.decode(sys.stdout.encoding or "utf-8")
.splitlines()
)
if "nthreads=" in output[0]:
nthreads = int(output[0].strip().split("=")[1])
if len(output) == nthreads:
using_openmp = True
else:
warnings.warn(
"Unexpected number of lines from output of test "
f"OpenMP program (output was {output})",
stacklevel=1,
)
using_openmp = False
else:
warnings.warn(
f"Unexpected output from test OpenMP program (output was {output})",
stacklevel=1,
)
using_openmp = False
except (CompileError, LinkError):
using_openmp = False
finally:
os.chdir(start_dir)
if using_openmp:
warnings.warn("Using OpenMP to compile parallel extensions", stacklevel=1)
else:
warnings.warn(
"Unable to compile OpenMP test program so Cython\n"
"extensions will be compiled without parallel support",
stacklevel=1,
)
return using_openmp
|
