Linking OpenMP libraries when compiling with setuptools (on macOS and linux)

[sodiumnitrate]

(Preface: this is extra complicated because I want this to work both on macOS and linux, and macOS+OpenMP is just pure pain. But, I think my question is valid for any platform, because I'm assuming I'd have to link to some omp library.)

I have the following C++ code:

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
// specific to macOS with libomp installed -- how to generalize?
#include "/usr/local/opt/libomp/include/omp.h"

namespace py = pybind11;

// OpenMP test
int sum_thread_ids(){
    int sum = 0;
    #pragma omp parallel shared(sum)
    {
        sleep(3);
        #pragma omp critical
        sum += omp_get_thread_num();
    }
    return sum;
}

PYBIND11_MODULE(openmp_test, m){
    m.def("get_max_threads", &omp_get_max_threads, "Returns max number of threads");
    m.def("set_num_threads", &omp_set_num_threads, "Set number of threads");
    m.def("sum_thread_ids", &sum_thread_ids, "Adds the id of threads");
}

And in setup.py I have:

import os
from glob import glob
from pybind11.setup_helpers import Pybind11Extension
from setuptools import setup, find_packages
import module_name

cxx_std = int(os.environ.get("CMAKE_CXX_STANDARD", "20"))

ext_modules = [
    Pybind11Extension("openmp_test", sorted(glob("src/*.cpp")), cxx_std=cxx_std)
]

setup(
    ...
    packages=find_packages(),
    install_requires=['pybind11'],
    python_requires='>=3.10',
    ext_modules=ext_modules
)

I install within the module directory that contains setup.py by doing:

pip install .

It compiles without errors or warnings. However, when I try to use it in python, it gives an import error:

>>>import openmp_test
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ImportError: dlopen(/Users/.../.pyenv/versions/3.11.2/lib/python3.11/site-packages/openmp_test.cpython-311-darwin.so, 0x0002): symbol not found in flat namespace '_omp_get_max_threads'

How do I link libomp to my package? (Preferably in both macOS and linux)

[cadop]

Did you figure this out?

I get a similar error on Ubuntu, but same code is working on Windows.

I also have this in setup.py

compiler_type = ccompiler.get_default_compiler()
openmp_flag = '/openmp' if compiler_type == 'msvc' else '-fopenmp'