diff --git a/environment.yaml b/environment.yaml index 8733562..006581f 100644 --- a/environment.yaml +++ b/environment.yaml @@ -8,3 +8,4 @@ dependencies: - numpy<=2.0 - horton - h5py + - llvm-openmp diff --git a/setup.py b/setup.py index 38734dd..4600e55 100644 --- a/setup.py +++ b/setup.py @@ -1,13 +1,35 @@ +import os +import platform + import pybind11 from setuptools import Extension, setup + +def get_openmp_flags(): + if platform.system() == "Darwin": + # Apple clang does not support -fopenmp directly. + conda_prefix = os.environ.get("CONDA_PREFIX", "") + if not conda_prefix: + raise RuntimeError( + "CONDA_PREFIX is not set. Activate your conda environment first." + ) + compile_args = ["-Xpreprocessor", "-fopenmp", f"-I{conda_prefix}/include"] + link_args = [f"-L{conda_prefix}/lib", "-lomp"] + else: + compile_args = ["-fopenmp"] + link_args = ["-fopenmp"] + return compile_args, link_args + + +_omp_compile, _omp_link = get_openmp_flags() + ext_modules = [ Extension( "pyxdm.core.exchange_hole_cpp", ["pyxdm/cpp/exchange_hole.cpp"], include_dirs=[pybind11.get_include()], - extra_compile_args=["-O3", "-fopenmp"], - extra_link_args=["-fopenmp"], + extra_compile_args=["-O3"] + _omp_compile, + extra_link_args=_omp_link, language="c++", ) ]