diff options
author | kguerda-idris <84066930+kguerda-idris@users.noreply.github.com> | 2021-09-29 15:29:31 +0200 |
---|---|---|
committer | GitHub <noreply@github.com> | 2021-09-29 15:29:31 +0200 |
commit | 1c7e7ce2da8bb362c184fb6eae71fe7e36356494 (patch) | |
tree | 92fdc31870b6d5384c8ba83ff72d85a2d5a1eee6 /ot/helpers/openmp_helpers.py | |
parent | 7dde9e8e4b6aae756e103d49198caaa4f24150e3 (diff) |
[MRG] OpenMP support (#260)
* Added : OpenMP support
Restored : Epsilon and Debug mode
Replaced : parmap => multiprocessing is now replace by multithreading
* Commit clean up
* Number of CPUs correctly calculated on SLURM clusters
* Corrected number of processes for cluster slurm
* Mistake corrected
* parmap is now deprecated
* Now a different solver is used depending on the requested number of threads
* Tiny mistake corrected
* Folders are now in the ot library instead of at the root
* Helpers is now correctly placed
* Attempt to make compilation work smoothly
* OS compatible path
* NumThreads now defaults to 1
* Better flags
* Mistake corrected in case of OpenMP unavailability
* Revert OpenMP flags modification, which do not compile on Windows
* Test helper functions
* Helpers comments
* Documentation update
* File title corrected
* Warning no longer using print
* Last attempt for macos compilation
* pls work
* atempt
* solving a type error
* TypeError OpenMP
* Compilation finally working on Windows
* Bug solve, number of threads now correctly selected
* 64 bits solver to avoid overflows for bigger problems
* 64 bits EMD corrected
Co-authored-by: kguerda-idris <ssos023@jean-zay3.idris.fr>
Co-authored-by: ncassereau-idris <84033440+ncassereau-idris@users.noreply.github.com>
Co-authored-by: ncassereau <nathan.cassereau@idris.fr>
Co-authored-by: RĂ©mi Flamary <remi.flamary@gmail.com>
Diffstat (limited to 'ot/helpers/openmp_helpers.py')
-rw-r--r-- | ot/helpers/openmp_helpers.py | 85 |
1 files changed, 85 insertions, 0 deletions
diff --git a/ot/helpers/openmp_helpers.py b/ot/helpers/openmp_helpers.py new file mode 100644 index 0000000..a6ad38b --- /dev/null +++ b/ot/helpers/openmp_helpers.py @@ -0,0 +1,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 |