Skip to content

Change default blas_info dictionary in cmodule #444

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Nov 9, 2023
Merged

Change default blas_info dictionary in cmodule #444

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6078345
Make default_blas_ldflags to not rely on numpy blas_info
lucianopaz Sep 18, 2023
f0c1021
Unpin numpy
lucianopaz Sep 22, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion environment.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ channels:
dependencies:
- python
- compilers
- numpy>=1.17.0,<1.26.0
- numpy>=1.17.0
- scipy>=0.14
- filelock
- etuples
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
327 changes: 139 additions & 188 deletions pytensor/link/c/cmodule.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import importlib
import logging
import os
import pathlib
import pickle
import platform
import re
Expand Down Expand Up @@ -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():
Expand Down
Loading