Make default_blas_ldflags to not rely on numpy blas_info

Author: lucianopaz

pytensor/link/c/cmodule.py

@@ -6,6 +6,7 @@
 import importlib
 import logging
 import os
+import pathlib
 import pickle
 import platform
 import re
@@ -2715,203 +2716,153 @@ def default_blas_ldflags():
     str
 
     """
-    warn_record = []
-    try:
-        blas_info = np.__config__.get_info("blas_opt")
-
-        # If we are in a EPD installation, mkl is available
-        if "EPD" in sys.version:
-            use_unix_epd = True
-            if sys.platform == "win32":
-                return " ".join(
-                    ['-L"%s"' % os.path.join(sys.prefix, "Scripts")]
-                    +
-                    # Why on Windows, the library used are not the
-                    # same as what is in
-                    # blas_info['libraries']?
-                    [f"-l{l}" for l in ("mk2_core", "mk2_intel_thread", "mk2_rt")]
-                )
-            elif sys.platform == "darwin":
-                # The env variable is needed to link with mkl
-                new_path = os.path.join(sys.prefix, "lib")
-                v = os.getenv("DYLD_FALLBACK_LIBRARY_PATH", None)
-                if v is not None:
-                    # Explicit version could be replaced by a symbolic
-                    # link called 'Current' created by EPD installer
-                    # This will resolve symbolic links
-                    v = os.path.realpath(v)
-
-                # The python __import__ don't seam to take into account
-                # the new env variable "DYLD_FALLBACK_LIBRARY_PATH"
-                # when we set with os.environ['...'] = X or os.putenv()
-                # So we warn the user and tell him what todo.
-                if v is None or new_path not in v.split(":"):
-                    _logger.warning(
-                        "The environment variable "
-                        "'DYLD_FALLBACK_LIBRARY_PATH' does not contain "
-                        "the '{new_path}' path in its value. This will make "
-                        "PyTensor use a slow version of BLAS. Update "
-                        "'DYLD_FALLBACK_LIBRARY_PATH' to contain the "
-                        "said value, this will disable this warning."
-                    )
-
-                    use_unix_epd = False
-            if use_unix_epd:
-                return " ".join(
-                    ["-L%s" % os.path.join(sys.prefix, "lib")]
-                    + ["-l%s" % l for l in blas_info["libraries"]]
-                )
-
-                # Canopy
-        if "Canopy" in sys.prefix:
-            subsub = "lib"
-            if sys.platform == "win32":
-                subsub = "Scripts"
-            lib_path = os.path.join(sys.base_prefix, subsub)
-            if not os.path.exists(lib_path):
-                # Old logic to find the path. I don't think we still
-                # need it, but I don't have the time to test all
-                # installation configuration. So I keep this as a fall
-                # back in case the current expectation don't work.
-
-                # This old logic don't work when multiple version of
-                # Canopy is installed.
-                p = os.path.join(sys.base_prefix, "..", "..", "appdata")
-                assert os.path.exists(p), "Canopy changed the location of MKL"
-                lib_paths = os.listdir(p)
-                # Try to remove subdir that can't contain MKL
-                for sub in lib_paths:
-                    if not os.path.exists(os.path.join(p, sub, subsub)):
-                        lib_paths.remove(sub)
-                assert len(lib_paths) == 1, (
-                    "Unexpected case when looking for Canopy MKL libraries",
-                    p,
-                    lib_paths,
-                    [os.listdir(os.path.join(p, sub)) for sub in lib_paths],
-                )
-                lib_path = os.path.join(p, lib_paths[0], subsub)
-                assert os.path.exists(lib_path), "Canopy changed the location of MKL"
-
-            if sys.platform == "linux2" or sys.platform == "darwin":
-                return " ".join(
-                    ["-L%s" % lib_path] + ["-l%s" % l for l in blas_info["libraries"]]
-                )
-            elif sys.platform == "win32":
-                return " ".join(
-                    ['-L"%s"' % lib_path]
-                    +
-                    # Why on Windows, the library used are not the
-                    # same as what is in blas_info['libraries']?
-                    [f"-l{l}" for l in ("mk2_core", "mk2_intel_thread", "mk2_rt")]
-                )
 
-        # MKL
-        # If mkl can be imported then use it. On conda:
-        # "conda install mkl-service" installs the Python wrapper and
-        # the low-level C libraries as well as optimised version of
-        # numpy and scipy.
-        try:
-            import mkl  # noqa
-        except ImportError:
-            pass
-        else:
-            # This branch is executed if no exception was raised
-            if sys.platform == "win32":
-                lib_path = os.path.join(sys.prefix, "Library", "bin")
-                flags = [f'-L"{lib_path}"']
-            else:
-                lib_path = blas_info.get("library_dirs", [])
-                flags = []
-                if lib_path:
-                    flags = [f"-L{lib_path[0]}"]
-            if "2018" in mkl.get_version_string():
-                thr = "mkl_gnu_thread"
-            else:
-                thr = "mkl_intel_thread"
-            base_flags = list(flags)
-            flags += [f"-l{l}" for l in ("mkl_core", thr, "mkl_rt")]
-            res = try_blas_flag(flags)
-
-            if not res and sys.platform == "win32" and thr == "mkl_gnu_thread":
-                # Check if it would work for intel OpenMP on windows
-                flags = base_flags + [
-                    f"-l{l}" for l in ("mkl_core", "mkl_intel_thread", "mkl_rt")
+    def check_required_file(paths, required_regexs):
+        libs = []
+        for req in required_regexs:
+            found = False
+            for path in paths:
+                m = re.search(req, path.name)
+                if m:
+                    libs.append((str(path.parent), m.string[slice(*m.span())]))
+                    found = True
+                    break
+            if not found:
+                raise RuntimeError(f"Required file {req} not found")
+        return libs
+
+    def get_cxx_library_dirs():
+        cmd = f"{config.cxx} -print-search-dirs"
+        p = subprocess_Popen(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            stdin=subprocess.PIPE,
+            shell=True,
+        )
+        (stdout, stderr) = p.communicate(input=b"")
+        maybe_lib_dirs = [
+            [pathlib.Path(p).resolve() for p in line[len("libraries: =") :].split(":")]
+            for line in stdout.decode(sys.stdout.encoding).splitlines()
+            if line.startswith("libraries: =")
+        ][0]
+        return [str(d) for d in maybe_lib_dirs if d.exists() and d.is_dir()]
+
+    def check_libs(
+        all_libs, required_libs, extra_compile_flags=None, cxx_library_dirs=None
+    ):
+        if cxx_library_dirs is None:
+            cxx_library_dirs = []
+        if extra_compile_flags is None:
+            extra_compile_flags = []
+        found_libs = check_required_file(
+            all_libs,
+            required_libs,
+        )
+        path_quote = '"' if sys.platform == "win32" else ""
+        libdir_ldflags = list(
+            dict.fromkeys(
+                [
+                    f"-L{path_quote}{lib_path}{path_quote}"
+                    for lib_path, _ in found_libs
+                    if lib_path not in cxx_library_dirs
                 ]
-                res = try_blas_flag(flags)
-
-            if res:
-                check_mkl_openmp()
-                return res
-
-            flags.extend(["-Wl,-rpath," + l for l in blas_info.get("library_dirs", [])])
-            res = try_blas_flag(flags)
-            if res:
-                check_mkl_openmp()
-                maybe_add_to_os_environ_pathlist("PATH", lib_path[0])
-                return res
-
-        # to support path that includes spaces, we need to wrap it with double quotes on Windows
-        path_wrapper = '"' if os.name == "nt" else ""
-        ret = (
-            # TODO: the Gemm op below should separate the
-            # -L and -l arguments into the two callbacks
-            # that CLinker uses for that stuff.  for now,
-            # we just pass the whole ldflags as the -l
-            # options part.
-            [
-                f"-L{path_wrapper}{l}{path_wrapper}"
-                for l in blas_info.get("library_dirs", [])
-            ]
-            + [f"-l{l}" for l in blas_info.get("libraries", [])]
-            + blas_info.get("extra_link_args", [])
+            )
         )
-        # For some very strange reason, we need to specify -lm twice
-        # to get mkl to link correctly.  I have no idea why.
-        if any("mkl" in fl for fl in ret):
-            ret.extend(["-lm", "-lm"])
-        res = try_blas_flag(ret)
-        if res:
-            if "mkl" in res:
-                check_mkl_openmp()
-            return res
 
-        # If we are using conda and can't reuse numpy blas, then doing
-        # the fallback and test -lblas could give slow computation, so
-        # warn about this.
-        for warn in warn_record:
-            _logger.warning(warn)
-        del warn_record
-
-        # Some environment don't have the lib dir in LD_LIBRARY_PATH.
-        # So add it.
-        ret.extend(["-Wl,-rpath," + l for l in blas_info.get("library_dirs", [])])
-        res = try_blas_flag(ret)
+        flags = (
+            libdir_ldflags
+            + [f"-l{lib_name}" for _, lib_name in found_libs]
+            + extra_compile_flags
+        )
+        res = try_blas_flag(flags)
         if res:
-            if "mkl" in res:
+            if any("mkl" in flag for flag in flags):
                 check_mkl_openmp()
             return res
+        else:
+            raise RuntimeError(f"Supplied flags {flags} failed to compile")
 
-        # Add sys.prefix/lib to the runtime search path. On
-        # non-system installations of Python that use the
-        # system linker, this is generally necessary.
-        if sys.platform in ("linux", "darwin"):
-            lib_path = os.path.join(sys.prefix, "lib")
-            ret.append("-Wl,-rpath," + lib_path)
-            res = try_blas_flag(ret)
-            if res:
-                if "mkl" in res:
-                    check_mkl_openmp()
-                return res
-
-    except KeyError:
+    _std_lib_dirs = std_lib_dirs()
+    if len(_std_lib_dirs) > 0:
+        rpath = _std_lib_dirs[0]
+    else:
+        rpath = None
+
+    cxx_library_dirs = get_cxx_library_dirs()
+    searched_library_dirs = cxx_library_dirs + _std_lib_dirs
+    all_libs = [
+        l
+        for path in [
+            pathlib.Path(library_dir)
+            for library_dir in searched_library_dirs
+            if pathlib.Path(library_dir).exists()
+        ]
+        for l in path.iterdir()
+        if l.suffix in {".so", ".dll", ".dylib"}
+    ]
+
+    if rpath is not None:
+        maybe_add_to_os_environ_pathlist("PATH", rpath)
+    try:
+        # 1. Try to use MKL with INTEL OpenMP threading
+        return check_libs(
+            all_libs,
+            required_libs=[
+                "mkl_core",
+                "mkl_rt",
+                "mkl_intel_thread",
+                "iomp5",
+                "pthread",
+            ],
+            extra_compile_flags=[f"-Wl,-rpath,{rpath}"] if rpath is not None else [],
+            cxx_library_dirs=cxx_library_dirs,
+        )
+    except Exception:
         pass
-
-    # Even if we could not detect what was used for numpy, or if these
-    # libraries are not found, most Linux systems have a libblas.so
-    # readily available. We try to see if that's the case, rather
-    # than disable blas. To test it correctly, we must load a program.
-    # Otherwise, there could be problem in the LD_LIBRARY_PATH.
-    return try_blas_flag(["-lblas"])
+    try:
+        # 2. Try to use MKL with GNU OpenMP threading
+        return check_libs(
+            all_libs,
+            required_libs=["mkl_core", "mkl_rt", "mkl_gnu_thread", "gomp", "pthread"],
+            extra_compile_flags=[f"-Wl,-rpath,{rpath}"] if rpath is not None else [],
+            cxx_library_dirs=cxx_library_dirs,
+        )
+    except Exception:
+        pass
+    try:
+        # 3. Try to use LAPACK + BLAS
+        return check_libs(
+            all_libs,
+            required_libs=["lapack", "blas", "cblas", "m"],
+            extra_compile_flags=[f"-Wl,-rpath,{rpath}"] if rpath is not None else [],
+            cxx_library_dirs=cxx_library_dirs,
+        )
+    except Exception:
+        pass
+    try:
+        # 4. Try to use BLAS alone
+        return check_libs(
+            all_libs,
+            required_libs=["blas", "cblas"],
+            extra_compile_flags=[f"-Wl,-rpath,{rpath}"] if rpath is not None else [],
+            cxx_library_dirs=cxx_library_dirs,
+        )
+    except Exception:
+        pass
+    try:
+        # 5. Try to use openblas
+        return check_libs(
+            all_libs,
+            required_libs=["openblas", "gfortran", "gomp", "m"],
+            extra_compile_flags=["-fopenmp", f"-Wl,-rpath,{rpath}"]
+            if rpath is not None
+            else ["-fopenmp"],
+            cxx_library_dirs=cxx_library_dirs,
+        )
+    except Exception:
+        pass
+    return ""
 
 
 def add_blas_configvars():