Build and use a loader library to link in libomp and libomptarget

- Avoid platform specific issues in the order of loading those libraries
to ensure inter-dependent symbol resolution

Author: ggeorgakoudis

buildscripts/conda-recipes/pyomp/meta.yaml

@@ -71,6 +71,8 @@ test:
     - test -f $SP_DIR/numba/openmp/libs/openmp/lib/libomptarget-amdgpu.bc                                             # [linux]
     - test -f $SP_DIR/numba/openmp/libs/openmp/lib/libomptarget-nvptx.bc                                              # [linux]
     - test -f $SP_DIR/numba/openmp/libs/openmp/lib/libomptarget.so.{{ LLVM_VERSION_MAJOR }}.{{ LLVM_VERSION_MINOR }}  # [linux]
+    - test -f $SP_DIR/numba/openmp/libs/openmp/lib/libpyomp_loader.dylib                                              # [osx]
+    - test -f $SP_DIR/numba/openmp/libs/openmp/lib/libpyomp_loader.so                                                 # [linux]
 
 about:
   home: https://github.com/Python-for-HPC/PyOMP

setup.py

@@ -54,6 +54,10 @@ def run(self):
             else:
                 super().run()
 
+        # Compile the loader wrapper after building all extensions to link with
+        # libomp and libomptarget.
+        self._build_loader_wrapper()
+
     def finalize_options(self):
         super().finalize_options()
         # Create placeholder directories for package-data validation.
@@ -146,6 +150,81 @@ def _env_toolchain_args(self, ext):
             args.append(f"-DCMAKE_CXX_COMPILER={os.environ['CXX']}")
         return args
 
+    def _build_loader_wrapper(self):
+        """Compiles the pyomp_loader wrapper to lock in load order."""
+        # Find build directory.
+        if self.inplace:
+            lib_dir = Path(
+                self.get_finalized_command("build_py").get_package_dir(
+                    "numba.openmp.libs"
+                )
+            )
+        else:
+            lib_dir = Path(self.build_lib) / "numba/openmp/libs"
+
+        target_lib_dir = lib_dir / "openmp" / "lib"
+        target_lib_dir.mkdir(parents=True, exist_ok=True)
+
+        # Find libomp in build directory or system library paths.
+        omp_libs = list(target_lib_dir.glob("libomp.*"))
+        if omp_libs:
+            link_omp = str(omp_libs[0])
+            print(f"Wrapper linking local libomp: {link_omp}")
+        else:
+            link_omp = "-lomp"
+            print("Wrapper linking system libomp (-lomp)")
+
+        # Find libomptarget in build directory or system library paths (Linux
+        # only). For now we always build libomptarget but we may want to link
+        # with system libomptarget if it becomes available.
+        if sys.platform.startswith("linux"):
+            tgt_libs = list(target_lib_dir.glob("libomptarget.so*"))
+            if tgt_libs:
+                assert len(tgt_libs) == 1, "Expected single libomptarget library"
+                link_tgt = str(tgt_libs[0])
+                print(f"Wrapper linking local libomptarget: {link_tgt}")
+            else:
+                link_tgt = "-lomptarget"
+                print("Wrapper linking system libomptarget (-lomptarget)")
+        else:
+            link_tgt = ""
+
+        # Generate the C wrapper file.
+        loader_c = Path(self.build_temp) / "pyomp_loader.c"
+        loader_c.parent.mkdir(parents=True, exist_ok=True)
+        loader_c.write_text("void pyomp_loader_init(void) {}\n")
+
+        # Determine library extension and platform-specific linker flags.
+        if sys.platform.startswith("linux"):
+            lib_ext = ".so"
+            rpath_flag = "-Wl,-rpath=$ORIGIN"
+            extra_link_flags = ["-Wl,--no-as-needed"]
+        elif sys.platform == "darwin":
+            lib_ext = ".dylib"
+            rpath_flag = "-Wl,-rpath,@loader_path"
+            extra_link_flags = ["-Wl,-undefined,dynamic_lookup"]
+        else:
+            raise RuntimeError(f"Unsupported platform: {sys.platform}")
+
+        loader_so = target_lib_dir / f"libpyomp_loader{lib_ext}"
+        cc = os.environ.get("CC", "gcc")
+
+        # Compile the wrapper.
+        cmd = [
+            cc,
+            "-shared",
+            "-fPIC",
+            str(loader_c),
+            "-o",
+            str(loader_so),
+            link_omp,
+            link_tgt,
+            rpath_flag,
+        ] + extra_link_flags
+
+        print("Compiling pyomp_loader wrapper")
+        subprocess.run(cmd, check=True)
+
 
 class PrepareOpenMP:
     setup_done = False

src/numba/openmp/__init__.py

@@ -60,48 +60,43 @@
 
 def _init():
     sys_platform = sys.platform
-    from ctypes.util import find_library
 
-    omplib = (
-        libpath
-        / "openmp"
-        / "lib"
-        / f"libomp{'.dylib' if sys_platform == 'darwin' else '.so'}"
-    )
+    # Find the wrapper library, which links with libomp and libomptarget, and
+    # loads them in the correct order with global visibility. Ensures RTLD_NEXT
+    # works correctly for libomp to find libomptarget symbols.
+    lib_ext = "dylib" if sys_platform == "darwin" else "so"
+    wrapper_lib = libpath / "openmp" / "lib" / f"libpyomp_loader.{lib_ext}"
 
-    # Prefer bundled libomp if it exists.
-    if omplib.exists():
-        if DEBUG_OPENMP >= 1:
-            print("Found bundled OpenMP runtime library at", omplib)
-        ll.load_library_permanently(str(omplib))
-    else:
-        # There is no bundled libomp, try to find it in standard library paths.
-        system_omplib = find_library("omp")
-        if system_omplib:
-            if DEBUG_OPENMP >= 1:
-                print(f"Found system OpenMP runtime library: {system_omplib}")
-            ll.load_library_permanently(system_omplib)
-        else:
-            raise RuntimeError(
-                f"OpenMP runtime not found. Bundled library missing at {omplib} "
-                "and no system libomp found via ctypes.util.find_library('omp'). "
-                "Ensure libomp is available in library paths."
-            )
-
-    # libomptarget is unavailable on apple, windows, so return.
-    if sys_platform.startswith("darwin") or sys_platform.startswith("win32"):
-        return
-
-    llvm_major, llvm_minor, _ = ll.llvm_version_info
-    omptargetlib = (
-        libpath / "openmp" / "lib" / f"libomptarget.so.{llvm_major}.{llvm_minor}"
-    )
-    if omptargetlib.exists():
+    if not wrapper_lib.exists():
+        raise RuntimeError(
+            f"OpenMP loader wrapper not found at {wrapper_lib}. "
+            "Ensure the package was built correctly."
+        )
+
+    if DEBUG_OPENMP >= 1:
+        print("Loading OpenMP runtimes via wrapper at", wrapper_lib)
+
+    # Load the wrapper.
+    ll.load_library_permanently(str(wrapper_lib))
+
+    # Initialize the OpenMP target runtime.
+    from ctypes import CFUNCTYPE
+
+    try:
+        addr = ll.address_of_symbol("__tgt_rtl_init")
+        if addr:
+            cfunctype = CFUNCTYPE(None)
+            cfunc = cfunctype(addr)
+            cfunc()
+
+        addr = ll.address_of_symbol("__tgt_init_all_rtls")
+        if addr:
+            cfunc = cfunctype(addr)
+            cfunc()
+
+    except Exception as e:
         if DEBUG_OPENMP >= 1:
-            print("Found OpenMP target runtime library at", omptargetlib)
-        ll.load_library_permanently(str(omptargetlib))
-    else:
-        raise RuntimeError(f"OpenMP target runtime not found at {omptargetlib}")
+            print(f"Warning: Failed to initialize OpenMP target runtime: {e}")
 
 
 _init()

src/numba/openmp/tests/test_openmp.py

@@ -5051,6 +5051,7 @@ def test_impl(lb, num_steps, pi_comp_func):
         njit_output = njit(test_impl)(0, 1024, test_pi_comp_njit)
         self.assert_outputs_equal(py_output, njit_output)
 
+
 class TestOpenmpRuntimeFunctions(TestOpenmpBase):
     def __init__(self, *args):
         TestOpenmpBase.__init__(self, *args)
@@ -5073,5 +5074,19 @@ def test_impl():
         python_num_procs = omp_get_num_procs()
         self.assertEqual(jit_num_procs, python_num_procs)
 
+    @linux_only
+    def test_omp_get_num_devices(self):
+        from numba.openmp import omp_get_num_devices
+
+        @njit
+        def test_impl():
+            return omp_get_num_devices()
+
+        jit_num_devices = test_impl()
+        python_num_devices = omp_get_num_devices()
+        self.assertEqual(jit_num_devices, python_num_devices)
+        self.assertGreater(jit_num_devices, 0)
+
+
 if __name__ == "__main__":
     unittest.main()