Support Apple clang specific compilation/linking flags to enable OpenMP

extension_helpers/_openmp_helpers.py

@@ -19,10 +19,11 @@
 import datetime
 import tempfile
 import subprocess
+import textwrap
 
 from setuptools.command.build_ext import customize_compiler, get_config_var, new_compiler
 
-from ._setup_helpers import get_compiler
+from ._setup_helpers import check_apple_clang, get_compiler
 
 __all__ = ['add_openmp_flags_if_available']
 
@@ -60,14 +61,20 @@ def _get_flag_value_from_var(flag, var, delim=' '):
         The environment variable to extract the flag from, e.g. CFLAGS or LDFLAGS.
     delim : str, optional
         The delimiter separating flags inside the environment variable
+
+    Returns
+    -------
+    extracted_flags : None|list
+        List of flags starting with flag extracted from var 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'
+        ['/usr/local/include']
 
     Notes
     -----
@@ -97,10 +104,15 @@ def _get_flag_value_from_var(flag, var, delim=' '):
             return None
 
     # Extract flag from {var:value}
+    extracted_flags = []
     if flags:
         for item in flags.split(delim):
             if item.startswith(flag):
-                return item[flag_length:]
+                extracted_flags.append(item[flag_length:])
+    if len(extracted_flags) > 0:
+        return extracted_flags
+    else:
+        return None
 
 
 def get_openmp_flags():
@@ -116,6 +128,9 @@ def get_openmp_flags():
     -----
     The flags returned are not tested for validity, use
     `check_openmp_support(openmp_flags=get_openmp_flags())` to do so.
+
+    On MacOS, it may require that you install `libomp` (e.g. with
+    `brew install libomp`).
     """
 
     compile_flags = []
@@ -127,15 +142,76 @@ def get_openmp_flags():
 
         include_path = _get_flag_value_from_var('-I', 'CFLAGS')
         if include_path:
-            compile_flags.append('-I' + include_path)
+            for _ in include_path:
+                compile_flags.append('-I' + _)
+
+        include_path = _get_flag_value_from_var('-I', 'CXXFLAGS')
+        if include_path:
+            for _ in include_path:
+                compile_flags.append('-I' + _)
 
         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')
+            for _ in lib_path:
+                link_flags.append('-L' + _)
+                link_flags.append('-Wl,-rpath,' + _)
+
+        if not check_apple_clang():
+            compile_flags.append('-fopenmp')
+            link_flags.append('-fopenmp')
+        else:
+            msg = textwrap.dedent(
+                """\
+                You are using Apple Clang compiler.
+
+                Your system should be prepared:
+                1. You should have specfically installed OpenMP,
+                   for instance by running `brew install libomp`.
+                2. OpenMP source and library should be findable by the compiler.
+
+                By default, `brew` will be use to find OpenMP source and
+                library. If not available, they will be looked for in
+                standard system directories.
+
+                To override this behavior and use specific OpenMP source and
+                library paths, you can setup the following environment
+                variables `CFLAGS` (or `CXXFLAGS`) and `LDFLAGS` before any
+                compilation/installation, e.g.
+                ```
+                export CFLAGS="-I/usr/local/opt/libomp/include"
+                export LDFLAGS="-L/usr/local/opt/libomp/lib"
+                ```
+                """
+            )
+            log.warn(msg)
+            # try to find path to libomp
+            try:
+                brew_check = subprocess.run(
+                    ["brew", "--prefix", "libomp"], capture_output=True
+                )
+                libomp = brew_check.stdout.decode('utf-8').strip()
+            except Exception:
+                if os.path.isdir('/usr/local/opt/libomp'):
+                    libomp = '/usr/local/opt/libomp'
+                elif os.path.isfile('/opt/homebrew/include/omp.h') \
+                    and os.path.isfile('/opt/homebrew/lib/libomp.a'):
+                    libomp = '/opt/homebrew'
+                else:
+                    libomp = None
+            # compile flags
+            compile_flags.append('-Xpreprocessor -fopenmp')
+            # additional include flag
+            if not 'CFLAGS' in os.environ and not 'CXXFLAGS' in os.environ \
+                and libomp is not None \
+                and os.path.isdir(os.path.join(libomp, 'include')):
+                    compile_flags.append('-I' + os.path.join(libomp, 'include'))
+            # link flag
+            link_flags.append('-lomp')
+            # additional link flag
+            if not 'LDFLAGS' in os.environ \
+                and libomp is not None \
+                and os.path.isdir(os.path.join(libomp, 'lib')):
+                link_flags.append('-L' + os.path.join(libomp, 'lib'))
 
     return {'compiler_flags': compile_flags, 'linker_flags': link_flags}
 
@@ -189,12 +265,12 @@ def check_openmp_support(openmp_flags=None):
 
             # Compile, test program
             ccompiler.compile(['test_openmp.c'], output_dir='objects',
-                              extra_postargs=compile_flags)
+                              extra_preargs=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)
+                                      extra_preargs=link_flags)
 
             # Run test program
             output = subprocess.check_output('./test_openmp')

extension_helpers/_setup_helpers.py

@@ -9,10 +9,11 @@
 import shutil
 import logging
 import subprocess
+import sys
 from collections import defaultdict
 
 from setuptools import Extension, find_packages
-from setuptools.command.build_ext import new_compiler
+from setuptools.command.build_ext import customize_compiler,new_compiler
 
 from ._utils import import_file, walk_skip_hidden
 
@@ -36,6 +37,34 @@ def get_compiler():
     return new_compiler().compiler_type
 
 
+def check_apple_clang():
+    """
+    Detemines if compiler that will be used to build extension modules is
+    'Apple Clang' (which requires a specific management of OpenMP compilation
+    and linking flags).
+
+    Note: it first checks that the OS is indeed MacOS.
+
+    Returns
+    -------
+    apple_clang : bool
+        Indicator whether current compiler is 'Apple Clang'.
+    """
+    if sys.platform != "darwin":
+        return False
+    else:
+        try:
+            ccompiler = new_compiler()
+            customize_compiler(ccompiler)
+            compiler_version = subprocess.run(
+                [ccompiler.compiler[0], "--version"], capture_output=True
+            )
+            apple_clang = "Apple clang" in compiler_version.stdout.decode('utf-8')
+        except Exception:
+            apple_clang = False
+        return apple_clang
+
+
 def get_extensions(srcdir='.'):
     """
     Collect all extensions from Cython files and ``setup_package.py`` files.

extension_helpers/tests/test_openmp_helpers.py

@@ -8,7 +8,7 @@
 import pytest
 from setuptools import Extension
 
-from .._openmp_helpers import add_openmp_flags_if_available, generate_openmp_enabled_py
+from .._openmp_helpers import _get_flag_value_from_var, add_openmp_flags_if_available, generate_openmp_enabled_py
 
 
 @pytest.fixture
@@ -51,3 +51,18 @@ def test_generate_openmp_enabled_py(openmp_expected):
 
     if openmp_expected is not None:
         assert openmp_expected is is_openmp_enabled
+
+
+def test_get_flag_value_from_var():
+    # define input
+    var = 'EXTTESTFLAGS'
+    flag = '-I'
+    # non existing var (at least should not)
+    assert _get_flag_value_from_var(flag, var) is None
+    # setup env varN
+    os.environ[var] = '-I/path/to/file1 -I/path/to/file2 -custom_option1 -custom_option2'
+    # non existing flag
+    assert _get_flag_value_from_var('-L', var) is None
+    # existing flag
+    assert _get_flag_value_from_var(flag, var) == ['/path/to/file1', '/path/to/file2']
+

extension_helpers/tests/test_setup_helpers.py

@@ -7,7 +7,7 @@
 
 import pytest
 
-from .._setup_helpers import get_compiler, get_extensions
+from .._setup_helpers import check_apple_clang, get_compiler, get_extensions
 from . import cleanup_import, run_setup
 
 extension_helpers_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))  # noqa
@@ -27,6 +27,9 @@ def teardown_module(module):
 def test_get_compiler():
     assert get_compiler() in POSSIBLE_COMPILERS
 
+def test_check_apple_clang():
+    assert check_apple_clang() in [True, False]
+
 
 def _extension_test_package(tmpdir, request, extension_type='c',
                             include_numpy=False):