1
# This module defines functions that can be used to check whether OpenMP is
2
# available and if so what flags to use. To use this, import the
3
# add_openmp_flags_if_available function in a setup_package.py file where you
4
# are defining your extensions:
5
#
6
#     from astropy_helpers.openmp_helpers import add_openmp_flags_if_available
7
#
8
# then call it with a single extension as the only argument:
9
#
10
#     add_openmp_flags_if_available(extension)
11
#
12
# this will add the OpenMP flags if available.
13

14 0
import os
15 0
import sys
16 0
import glob
17 0
import time
18 0
import datetime
19 0
import tempfile
20 0
import subprocess
21

22 0
from distutils import log
23 0
from distutils.ccompiler import new_compiler
24 0
from distutils.sysconfig import customize_compiler, get_config_var
25 0
from distutils.errors import CompileError, LinkError
26

27 0
from .distutils_helpers import get_compiler_option
28

29 0
__all__ = ['add_openmp_flags_if_available']
30

31 0
try:
32
    # Check if this has already been instantiated, only set the default once.
33 0
    _ASTROPY_DISABLE_SETUP_WITH_OPENMP_
34 0
except NameError:
35 0
    import builtins
36
    # It hasn't, so do so.
37 0
    builtins._ASTROPY_DISABLE_SETUP_WITH_OPENMP_ = False
38

39 0
CCODE = """
40
#include <omp.h>
41
#include <stdio.h>
42
int main(void) {
43
  #pragma omp parallel
44
  printf("nthreads=%d\\n", omp_get_num_threads());
45
  return 0;
46
}
47
"""
48

49

50 0
def _get_flag_value_from_var(flag, var, delim=' '):
51
    """
52
    Extract flags from an environment variable.
53

54
    Parameters
55
    ----------
56
    flag : str
57
        The flag to extract, for example '-I' or '-L'
58
    var : str
59
        The environment variable to extract the flag from, e.g. CFLAGS or LDFLAGS.
60
    delim : str, optional
61
        The delimiter separating flags inside the environment variable
62

63
    Examples
64
    --------
65
    Let's assume the LDFLAGS is set to '-L/usr/local/include -customflag'. This
66
    function will then return the following:
67

68
        >>> _get_flag_value_from_var('-L', 'LDFLAGS')
69
        '/usr/local/include'
70

71
    Notes
72
    -----
73
    Environment variables are first checked in ``os.environ[var]``, then in
74
    ``distutils.sysconfig.get_config_var(var)``.
75

76
    This function is not supported on Windows.
77
    """
78

79 0
    if sys.platform.startswith('win'):
80 0
        return None
81

82
    # Simple input validation
83 0
    if not var or not flag:
84 0
        return None
85 0
    flag_length = len(flag)
86 0
    if not flag_length:
87 0
        return None
88

89
    # Look for var in os.eviron then in get_config_var
90 0
    if var in os.environ:
91 0
        flags = os.environ[var]
92
    else:
93 0
        try:
94 0
            flags = get_config_var(var)
95 0
        except KeyError:
96 0
            return None
97

98
    # Extract flag from {var:value}
99 0
    if flags:
100 0
        for item in flags.split(delim):
101 0
            if item.startswith(flag):
102 0
                return item[flag_length:]
103

104

105 0
def get_openmp_flags():
106
    """
107
    Utility for returning compiler and linker flags possibly needed for
108
    OpenMP support.
109

110
    Returns
111
    -------
112
    result : `{'compiler_flags':<flags>, 'linker_flags':<flags>}`
113

114
    Notes
115
    -----
116
    The flags returned are not tested for validity, use
117
    `check_openmp_support(openmp_flags=get_openmp_flags())` to do so.
118
    """
119

120 0
    compile_flags = []
121 0
    link_flags = []
122

123 0
    if get_compiler_option() == 'msvc':
124 0
        compile_flags.append('-openmp')
125
    else:
126

127 0
        include_path = _get_flag_value_from_var('-I', 'CFLAGS')
128 0
        if include_path:
129 0
            compile_flags.append('-I' + include_path)
130

131 0
        lib_path = _get_flag_value_from_var('-L', 'LDFLAGS')
132 0
        if lib_path:
133 0
            link_flags.append('-L' + lib_path)
134 0
            link_flags.append('-Wl,-rpath,' + lib_path)
135

136 0
        compile_flags.append('-fopenmp')
137 0
        link_flags.append('-fopenmp')
138

139 0
    return {'compiler_flags': compile_flags, 'linker_flags': link_flags}
140

141

142 0
def check_openmp_support(openmp_flags=None):
143
    """
144
    Check whether OpenMP test code can be compiled and run.
145

146
    Parameters
147
    ----------
148
    openmp_flags : dict, optional
149
        This should be a dictionary with keys ``compiler_flags`` and
150
        ``linker_flags`` giving the compiliation and linking flags respectively.
151
        These are passed as `extra_postargs` to `compile()` and
152
        `link_executable()` respectively. If this is not set, the flags will
153
        be automatically determined using environment variables.
154

155
    Returns
156
    -------
157
    result : bool
158
        `True` if the test passed, `False` otherwise.
159
    """
160

161 0
    ccompiler = new_compiler()
162 0
    customize_compiler(ccompiler)
163

164 0
    if not openmp_flags:
165
        # customize_compiler() extracts info from os.environ. If certain keys
166
        # exist it uses these plus those from sysconfig.get_config_vars().
167
        # If the key is missing in os.environ it is not extracted from
168
        # sysconfig.get_config_var(). E.g. 'LDFLAGS' get left out, preventing
169
        # clang from finding libomp.dylib because -L<path> is not passed to
170
        # linker. Call get_openmp_flags() to get flags missed by
171
        # customize_compiler().
172 0
        openmp_flags = get_openmp_flags()
173

174 0
    compile_flags = openmp_flags.get('compiler_flags')
175 0
    link_flags = openmp_flags.get('linker_flags')
176

177 0
    tmp_dir = tempfile.mkdtemp()
178 0
    start_dir = os.path.abspath('.')
179

180 0
    try:
181 0
        os.chdir(tmp_dir)
182

183
        # Write test program
184 0
        with open('test_openmp.c', 'w') as f:
185 0
            f.write(CCODE)
186

187 0
        os.mkdir('objects')
188

189
        # Compile, test program
190 0
        ccompiler.compile(['test_openmp.c'], output_dir='objects',
191
                          extra_postargs=compile_flags)
192

193
        # Link test program
194 0
        objects = glob.glob(os.path.join('objects', '*' + ccompiler.obj_extension))
195 0
        ccompiler.link_executable(objects, 'test_openmp',
196
                                  extra_postargs=link_flags)
197

198
        # Run test program
199 0
        output = subprocess.check_output('./test_openmp')
200 0
        output = output.decode(sys.stdout.encoding or 'utf-8').splitlines()
201

202 0
        if 'nthreads=' in output[0]:
203 0
            nthreads = int(output[0].strip().split('=')[1])
204 0
            if len(output) == nthreads:
205 0
                is_openmp_supported = True
206
            else:
207 0
                log.warn("Unexpected number of lines from output of test OpenMP "
208
                         "program (output was {0})".format(output))
209 0
                is_openmp_supported = False
210
        else:
211 0
            log.warn("Unexpected output from test OpenMP "
212
                     "program (output was {0})".format(output))
213 0
            is_openmp_supported = False
214 0
    except (CompileError, LinkError, subprocess.CalledProcessError):
215 0
        is_openmp_supported = False
216

217
    finally:
218 0
        os.chdir(start_dir)
219

220 0
    return is_openmp_supported
221

222

223 0
def is_openmp_supported():
224
    """
225
    Determine whether the build compiler has OpenMP support.
226
    """
227 0
    log_threshold = log.set_threshold(log.FATAL)
228 0
    ret = check_openmp_support()
229 0
    log.set_threshold(log_threshold)
230 0
    return ret
231

232

233 0
def add_openmp_flags_if_available(extension):
234
    """
235
    Add OpenMP compilation flags, if supported (if not a warning will be
236
    printed to the console and no flags will be added.)
237

238
    Returns `True` if the flags were added, `False` otherwise.
239
    """
240

241 0
    if _ASTROPY_DISABLE_SETUP_WITH_OPENMP_:
242 0
        log.info("OpenMP support has been explicitly disabled.")
243 0
        return False
244

245 0
    openmp_flags = get_openmp_flags()
246 0
    using_openmp = check_openmp_support(openmp_flags=openmp_flags)
247

248 0
    if using_openmp:
249 0
        compile_flags = openmp_flags.get('compiler_flags')
250 0
        link_flags = openmp_flags.get('linker_flags')
251 0
        log.info("Compiling Cython/C/C++ extension with OpenMP support")
252 0
        extension.extra_compile_args.extend(compile_flags)
253 0
        extension.extra_link_args.extend(link_flags)
254
    else:
255 0
        log.warn("Cannot compile Cython/C/C++ extension with OpenMP, reverting "
256
                 "to non-parallel code")
257

258 0
    return using_openmp
259

260

261 0
_IS_OPENMP_ENABLED_SRC = """
262
# Autogenerated by {packagetitle}'s setup.py on {timestamp!s}
263

264
def is_openmp_enabled():
265
    \"\"\"
266
    Determine whether this package was built with OpenMP support.
267
    \"\"\"
268
    return {return_bool}
269
"""[1:]
270

271

272 0
def generate_openmp_enabled_py(packagename, srcdir='.', disable_openmp=None):
273
    """
274
    Generate ``package.openmp_enabled.is_openmp_enabled``, which can then be used
275
    to determine, post build, whether the package was built with or without
276
    OpenMP support.
277
    """
278

279 0
    if packagename.lower() == 'astropy':
280 0
        packagetitle = 'Astropy'
281
    else:
282 0
        packagetitle = packagename
283

284 0
    epoch = int(os.environ.get('SOURCE_DATE_EPOCH', time.time()))
285 0
    timestamp = datetime.datetime.utcfromtimestamp(epoch)
286

287 0
    if disable_openmp is not None:
288 0
        import builtins
289 0
        builtins._ASTROPY_DISABLE_SETUP_WITH_OPENMP_ = disable_openmp
290 0
    if _ASTROPY_DISABLE_SETUP_WITH_OPENMP_:
291 0
        log.info("OpenMP support has been explicitly disabled.")
292 0
    openmp_support = False if _ASTROPY_DISABLE_SETUP_WITH_OPENMP_ else is_openmp_supported()
293

294 0
    src = _IS_OPENMP_ENABLED_SRC.format(packagetitle=packagetitle,
295
                                        timestamp=timestamp,
296
                                        return_bool=openmp_support)
297

298 0
    package_srcdir = os.path.join(srcdir, *packagename.split('.'))
299 0
    is_openmp_enabled_py = os.path.join(package_srcdir, 'openmp_enabled.py')
300 0
    with open(is_openmp_enabled_py, 'w') as f:
301 0
        f.write(src)

Read our documentation on viewing source code .

Loading