openmp_helpers.py at 7522f8c ⋅ astropy/astropy-helpers

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

import os
import sys
import glob
import time
import datetime
import tempfile
import subprocess

from distutils import log
from distutils.ccompiler import new_compiler
from distutils.sysconfig import customize_compiler, get_config_var
from distutils.errors import CompileError, LinkError

from .distutils_helpers import get_compiler_option

__all__ = ['add_openmp_flags_if_available']

try:
    # Check if this has already been instantiated, only set the default once.
    _ASTROPY_DISABLE_SETUP_WITH_OPENMP_
except NameError:
    import builtins
    # It hasn't, so do so.
    builtins._ASTROPY_DISABLE_SETUP_WITH_OPENMP_ = False

CCODE = """
#include <omp.h>
#include <stdio.h>
int main(void) {
  #pragma omp parallel
  printf("nthreads=%d\\n", omp_get_num_threads());
  return 0;
}
"""


def _get_flag_value_from_var(flag, var, delim=' '):
    """
    Extract flags from an environment variable.

    Parameters
    ----------
    flag : str
        The flag to extract, for example '-I' or '-L'
    var : str
        The environment variable to extract the flag from, e.g. CFLAGS or LDFLAGS.
    delim : str, optional
        The delimiter separating flags inside the environment variable

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

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

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

    This function is not supported on Windows.
    """

    if sys.platform.startswith('win'):
        return None

    # Simple input validation
    if not var or not flag:
        return None
    flag_length = len(flag)
    if not flag_length:
        return None

    # Look for var in os.eviron then in get_config_var
    if var in os.environ:
        flags = os.environ[var]
    else:
        try:
            flags = get_config_var(var)
        except KeyError:
            return None

    # Extract flag from {var:value}
    if flags:
        for item in flags.split(delim):
            if item.startswith(flag):
                return item[flag_length:]


def get_openmp_flags():
    """
    Utility for returning compiler and linker flags possibly needed for
    OpenMP support.

    Returns
    -------
    result : `{'compiler_flags':<flags>, 'linker_flags':<flags>}`

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

    compile_flags = []
    link_flags = []

    if get_compiler_option() == 'msvc':
        compile_flags.append('-openmp')
    else:

        include_path = _get_flag_value_from_var('-I', 'CFLAGS')
        if include_path:
            compile_flags.append('-I' + include_path)

        lib_path = _get_flag_value_from_var('-L', 'LDFLAGS')
        if lib_path:
            link_flags.append('-L' + lib_path)
            link_flags.append('-Wl,-rpath,' + lib_path)

        compile_flags.append('-fopenmp')
        link_flags.append('-fopenmp')

    return {'compiler_flags': compile_flags, 'linker_flags': link_flags}


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

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

    Returns
    -------
    result : bool
        `True` if the test passed, `False` otherwise.
    """

    ccompiler = new_compiler()
    customize_compiler(ccompiler)

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

    compile_flags = openmp_flags.get('compiler_flags')
    link_flags = openmp_flags.get('linker_flags')

    tmp_dir = tempfile.mkdtemp()
    start_dir = os.path.abspath('.')

    try:
        os.chdir(tmp_dir)

        # Write test program
        with open('test_openmp.c', 'w') as f:
            f.write(CCODE)

        os.mkdir('objects')

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

        # Link test program
        objects = glob.glob(os.path.join('objects', '*' + ccompiler.obj_extension))
        ccompiler.link_executable(objects, 'test_openmp',
                                  extra_postargs=link_flags)

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

        if 'nthreads=' in output[0]:
            nthreads = int(output[0].strip().split('=')[1])
            if len(output) == nthreads:
                is_openmp_supported = True
            else:
                log.warn("Unexpected number of lines from output of test OpenMP "
                         "program (output was {0})".format(output))
                is_openmp_supported = False
        else:
            log.warn("Unexpected output from test OpenMP "
                     "program (output was {0})".format(output))
            is_openmp_supported = False
    except (CompileError, LinkError, subprocess.CalledProcessError):
        is_openmp_supported = False

    finally:
        os.chdir(start_dir)

    return is_openmp_supported


def is_openmp_supported():
    """
    Determine whether the build compiler has OpenMP support.
    """
    log_threshold = log.set_threshold(log.FATAL)
    ret = check_openmp_support()
    log.set_threshold(log_threshold)
    return ret


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

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

    if _ASTROPY_DISABLE_SETUP_WITH_OPENMP_:
        log.info("OpenMP support has been explicitly disabled.")
        return False

    openmp_flags = get_openmp_flags()
    using_openmp = check_openmp_support(openmp_flags=openmp_flags)

    if using_openmp:
        compile_flags = openmp_flags.get('compiler_flags')
        link_flags = openmp_flags.get('linker_flags')
        log.info("Compiling Cython/C/C++ extension with OpenMP support")
        extension.extra_compile_args.extend(compile_flags)
        extension.extra_link_args.extend(link_flags)
    else:
        log.warn("Cannot compile Cython/C/C++ extension with OpenMP, reverting "
                 "to non-parallel code")

    return using_openmp


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

def is_openmp_enabled():
    \"\"\"
    Determine whether this package was built with OpenMP support.
    \"\"\"
    return {return_bool}
"""[1:]


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

    if packagename.lower() == 'astropy':
        packagetitle = 'Astropy'
    else:
        packagetitle = packagename

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

    if disable_openmp is not None:
        import builtins
        builtins._ASTROPY_DISABLE_SETUP_WITH_OPENMP_ = disable_openmp
    if _ASTROPY_DISABLE_SETUP_WITH_OPENMP_:
        log.info("OpenMP support has been explicitly disabled.")
    openmp_support = False if _ASTROPY_DISABLE_SETUP_WITH_OPENMP_ else is_openmp_supported()

    src = _IS_OPENMP_ENABLED_SRC.format(packagetitle=packagetitle,
                                        timestamp=timestamp,
                                        return_bool=openmp_support)

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

Read our documentation on viewing source code .

Navigation	Overlay
`t` Navigate files	`h` Toggle hits
`y` Change url to tip of branch	`m` Toggle misses
`b / v` Jump to prev/next hit line	`p` Toggle partial
`z / x` Jump to prev/next missed or partial line	`1..9` Toggle flags
`shift + o` Open current page in GitHub	`a` Toggle all on
`/ or ?` Show keyboard shortcuts dialog	`c` Toggle context lines or commits

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)