diff --git a/environment.yml b/environment.yml index 06ba81250b..98be2b7856 100644 --- a/environment.yml +++ b/environment.yml @@ -9,7 +9,7 @@ channels: dependencies: - python - compilers - - numpy>=1.17.0,<1.26.0 + - numpy>=1.17.0 - scipy>=0.14 - filelock - etuples diff --git a/pyproject.toml b/pyproject.toml index cb6e351763..baa642628a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,7 +52,7 @@ keywords = [ dependencies = [ "setuptools>=48.0.0", "scipy>=0.14", - "numpy>=1.17.0,<1.26", + "numpy>=1.17.0", "filelock", "etuples", "logical-unification", diff --git a/pytensor/link/c/cmodule.py b/pytensor/link/c/cmodule.py index 895215a4ce..69363570bb 100644 --- a/pytensor/link/c/cmodule.py +++ b/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(): diff --git a/tests/link/c/test_cmodule.py b/tests/link/c/test_cmodule.py index ab8ae3f817..b718fd406b 100644 --- a/tests/link/c/test_cmodule.py +++ b/tests/link/c/test_cmodule.py @@ -4,11 +4,11 @@ But this one tests a current behavior that isn't good: the c_code isn't deterministic based on the input type and the op. """ -import logging import multiprocessing import os +import sys import tempfile -from unittest.mock import patch +from unittest.mock import MagicMock, patch import numpy as np import pytest @@ -161,16 +161,69 @@ def test_flag_detection(): assert isinstance(res, bool) -@patch("pytensor.link.c.cmodule.try_blas_flag", return_value=None) -@patch("pytensor.link.c.cmodule.sys") -def test_default_blas_ldflags(sys_mock, try_blas_flag_mock, caplog): - sys_mock.version = "3.8.0 | packaged by conda-forge | (default, Nov 22 2019, 19:11:38) \n[GCC 7.3.0]" - - with patch.dict("sys.modules", {"mkl": None}): - with caplog.at_level(logging.WARNING): - default_blas_ldflags() - - assert caplog.text == "" +@pytest.fixture( + scope="module", + params=["mkl_intel", "mkl_gnu", "openblas", "lapack", "blas", "no_blas"], +) +def blas_libs(request): + key = request.param + libs = { + "mkl_intel": ["mkl_core", "mkl_rt", "mkl_intel_thread", "iomp5", "pthread"], + "mkl_gnu": ["mkl_core", "mkl_rt", "mkl_gnu_thread", "gomp", "pthread"], + "openblas": ["openblas", "gfortran", "gomp", "m"], + "lapack": ["lapack", "blas", "cblas", "m"], + "blas": ["blas", "cblas"], + "no_blas": [], + } + return libs[key] + + +@pytest.fixture(scope="function", params=["Linux", "Windows", "Darwin"]) +def mock_system(request): + with patch("platform.system", return_value=request.param): + yield request.param + + +@pytest.fixture() +def cxx_search_dirs(blas_libs, mock_system): + libext = {"Linux": "so", "Windows": "dll", "Darwin": "dylib"} + libtemplate = f"{{lib}}.{libext[mock_system]}" + libraries = [] + with tempfile.TemporaryDirectory() as d: + flags = None + for lib in blas_libs: + lib_path = os.path.join(d, libtemplate.format(lib=lib)) + with open(lib_path, "wb") as f: + f.write(b"1") + libraries.append(lib_path) + if flags is None: + flags = f"-l{lib}" + else: + flags += f" -l{lib}" + if "gomp" in blas_libs and "mkl_gnu_thread" not in blas_libs: + flags += " -fopenmp" + if len(blas_libs) == 0: + flags = "" + yield f"libraries: ={d}".encode(sys.stdout.encoding), flags + + +@patch("pytensor.link.c.cmodule.std_lib_dirs", return_value=[]) +@patch("pytensor.link.c.cmodule.check_mkl_openmp", return_value=None) +def test_default_blas_ldflags( + mock_std_lib_dirs, mock_check_mkl_openmp, cxx_search_dirs +): + cxx_search_dirs, expected_blas_ldflags = cxx_search_dirs + mock_process = MagicMock() + mock_process.communicate = lambda *args, **kwargs: (cxx_search_dirs, None) + with patch("pytensor.link.c.cmodule.subprocess_Popen", return_value=mock_process): + with patch.object( + pytensor.link.c.cmodule.GCC_compiler, + "try_compile_tmp", + return_value=(True, True), + ): + assert set(default_blas_ldflags().split(" ")) == set( + expected_blas_ldflags.split(" ") + ) @patch(