Skip to content

P
pytensor
提交 48f1deb9 authored 作者: lucianopaz's avatar lucianopaz 提交者： Thomas Wiecki
浏览文件
操作

Make default_blas_ldflags to not rely on numpy blas_info

上级 5ad11819
显示空白字符变更
内嵌 并排
正在显示 包含 200 行增加 和 196 行删除
pytensor/link/c/cmodule.py
浏览文件 @ 48f1deb9
...@@ -6,6 +6,7 @@ import atexit ...@@ -6,6 +6,7 @@ import atexit
import importlib import importlib
import logging import logging
import os import os
import pathlib
import pickle import pickle
import platform import platform
import re import re
...@@ -2715,203 +2716,153 @@ def default_blas_ldflags(): ...@@ -2715,203 +2716,153 @@ def default_blas_ldflags():
str str
""" """
warn_record = []
try:
blas_info = np.__config__.get_info("blas_opt")
# If we are in a EPD installation, mkl is available def check_required_file(paths, required_regexs):
if "EPD" in sys.version: libs = []
use_unix_epd = True for req in required_regexs:
if sys.platform == "win32": found = False
return " ".join( for path in paths:
['-L"%s"' % os.path.join(sys.prefix, "Scripts")] m = re.search(req, path.name)
+ if m:
# Why on Windows, the library used are not the libs.append((str(path.parent), m.string[slice(*m.span())]))
# same as what is in found = True
# blas_info['libraries']? break
[f"-l{l}" for l in ("mk2_core", "mk2_intel_thread", "mk2_rt")] if not found:
) raise RuntimeError(f"Required file {req} not found")
elif sys.platform == "darwin": return libs
# 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 def get_cxx_library_dirs():
if use_unix_epd: cmd = f"{config.cxx} -print-search-dirs"
return " ".join( p = subprocess_Popen(
["-L%s" % os.path.join(sys.prefix, "lib")] cmd,
+ ["-l%s" % l for l in blas_info["libraries"]] stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
stdin=subprocess.PIPE,
shell=True,
) )
(stdout, stderr) = p.communicate(input=b"")
# Canopy maybe_lib_dirs = [
if "Canopy" in sys.prefix: [pathlib.Path(p).resolve() for p in line[len("libraries: =") :].split(":")]
subsub = "lib" for line in stdout.decode(sys.stdout.encoding).splitlines()
if sys.platform == "win32": if line.startswith("libraries: =")
subsub = "Scripts" ][0]
lib_path = os.path.join(sys.base_prefix, subsub) return [str(d) for d in maybe_lib_dirs if d.exists() and d.is_dir()]
if not os.path.exists(lib_path):
# Old logic to find the path. I don't think we still def check_libs(
# need it, but I don't have the time to test all all_libs, required_libs, extra_compile_flags=None, cxx_library_dirs=None
# installation configuration. So I keep this as a fall ):
# back in case the current expectation don't work. if cxx_library_dirs is None:
cxx_library_dirs = []
# This old logic don't work when multiple version of if extra_compile_flags is None:
# Canopy is installed. extra_compile_flags = []
p = os.path.join(sys.base_prefix, "..", "..", "appdata") found_libs = check_required_file(
assert os.path.exists(p), "Canopy changed the location of MKL" all_libs,
lib_paths = os.listdir(p) required_libs,
# Try to remove subdir that can't contain MKL )
for sub in lib_paths: path_quote = '"' if sys.platform == "win32" else ""
if not os.path.exists(os.path.join(p, sub, subsub)): libdir_ldflags = list(
lib_paths.remove(sub) dict.fromkeys(
assert len(lib_paths) == 1, ( [
"Unexpected case when looking for Canopy MKL libraries", f"-L{path_quote}{lib_path}{path_quote}"
p, for lib_path, _ in found_libs
lib_paths, if lib_path not in cxx_library_dirs
[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")
] ]
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", [])]) flags = (
libdir_ldflags
+ [f"-l{lib_name}" for _, lib_name in found_libs]
+ extra_compile_flags
)
res = try_blas_flag(flags) res = try_blas_flag(flags)
if res: if res:
if any("mkl" in flag for flag in flags):
check_mkl_openmp() check_mkl_openmp()
maybe_add_to_os_environ_pathlist("PATH", lib_path[0])
return res return res
else:
raise RuntimeError(f"Supplied flags {flags} failed to compile")
# to support path that includes spaces, we need to wrap it with double quotes on Windows _std_lib_dirs = std_lib_dirs()
path_wrapper = '"' if os.name == "nt" else "" if len(_std_lib_dirs) > 0:
ret = ( rpath = _std_lib_dirs[0]
# TODO: the Gemm op below should separate the else:
# -L and -l arguments into the two callbacks rpath = None
# that CLinker uses for that stuff. for now,
# we just pass the whole ldflags as the -l cxx_library_dirs = get_cxx_library_dirs()
# options part. searched_library_dirs = cxx_library_dirs + _std_lib_dirs
[ all_libs = [
f"-L{path_wrapper}{l}{path_wrapper}" l
for l in blas_info.get("library_dirs", []) 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"}
] ]
+ [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)
if res:
if "mkl" in res:
check_mkl_openmp()
return res
# 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: 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 pass
try:
# Even if we could not detect what was used for numpy, or if these # 2. Try to use MKL with GNU OpenMP threading
# libraries are not found, most Linux systems have a libblas.so return check_libs(
# readily available. We try to see if that's the case, rather all_libs,
# than disable blas. To test it correctly, we must load a program. required_libs=["mkl_core", "mkl_rt", "mkl_gnu_thread", "gomp", "pthread"],
# Otherwise, there could be problem in the LD_LIBRARY_PATH. extra_compile_flags=[f"-Wl,-rpath,{rpath}"] if rpath is not None else [],
return try_blas_flag(["-lblas"]) 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(): def add_blas_configvars():
......
tests/link/c/test_cmodule.py
浏览文件 @ 48f1deb9
...@@ -4,11 +4,11 @@ We don't have real tests for the cache, but it would be great to make them! ...@@ -4,11 +4,11 @@ We don't have real tests for the cache, but it would be great to make them!
But this one tests a current behavior that isn't good: the c_code isn't 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. deterministic based on the input type and the op.
""" """
import logging
import multiprocessing import multiprocessing
import os import os
import sys
import tempfile import tempfile
from unittest.mock import patch from unittest.mock import MagicMock, patch
import numpy as np import numpy as np
import pytest import pytest
...@@ -161,16 +161,69 @@ def test_flag_detection(): ...@@ -161,16 +161,69 @@ def test_flag_detection():
assert isinstance(res, bool) assert isinstance(res, bool)
@patch("pytensor.link.c.cmodule.try_blas_flag", return_value=None) @pytest.fixture(
@patch("pytensor.link.c.cmodule.sys") scope="module",
def test_default_blas_ldflags(sys_mock, try_blas_flag_mock, caplog): params=["mkl_intel", "mkl_gnu", "openblas", "lapack", "blas", "no_blas"],
sys_mock.version = "3.8.0 | packaged by conda-forge | (default, Nov 22 2019, 19:11:38) \n[GCC 7.3.0]" )
def blas_libs(request):
with patch.dict("sys.modules", {"mkl": None}): key = request.param
with caplog.at_level(logging.WARNING): libs = {
default_blas_ldflags() "mkl_intel": ["mkl_core", "mkl_rt", "mkl_intel_thread", "iomp5", "pthread"],
"mkl_gnu": ["mkl_core", "mkl_rt", "mkl_gnu_thread", "gomp", "pthread"],
assert caplog.text == "" "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( @patch(
......
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论