Skip to content

Commit 6078345

Browse files
lucianopazlucianopaz
committed
Make default_blas_ldflags to not rely on numpy blas_info
1 parent c9159b2 commit 6078345

File tree

2 files changed

+204
-200
lines changed

2 files changed

+204
-200
lines changed

‎pytensor/link/c/cmodule.py

+139-188
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import importlib
77
import logging
88
import os
9+
import pathlib
910
import pickle
1011
import platform
1112
import re
@@ -2715,203 +2716,153 @@ def default_blas_ldflags():
27152716
str
27162717
27172718
"""
2718-
warn_record = []
2719-
try:
2720-
blas_info = np.__config__.get_info("blas_opt")
2721-
2722-
# If we are in a EPD installation, mkl is available
2723-
if "EPD" in sys.version:
2724-
use_unix_epd = True
2725-
if sys.platform == "win32":
2726-
return " ".join(
2727-
['-L"%s"' % os.path.join(sys.prefix, "Scripts")]
2728-
+
2729-
# Why on Windows, the library used are not the
2730-
# same as what is in
2731-
# blas_info['libraries']?
2732-
[f"-l{l}" for l in ("mk2_core", "mk2_intel_thread", "mk2_rt")]
2733-
)
2734-
elif sys.platform == "darwin":
2735-
# The env variable is needed to link with mkl
2736-
new_path = os.path.join(sys.prefix, "lib")
2737-
v = os.getenv("DYLD_FALLBACK_LIBRARY_PATH", None)
2738-
if v is not None:
2739-
# Explicit version could be replaced by a symbolic
2740-
# link called 'Current' created by EPD installer
2741-
# This will resolve symbolic links
2742-
v = os.path.realpath(v)
2743-
2744-
# The python __import__ don't seam to take into account
2745-
# the new env variable "DYLD_FALLBACK_LIBRARY_PATH"
2746-
# when we set with os.environ['...'] = X or os.putenv()
2747-
# So we warn the user and tell him what todo.
2748-
if v is None or new_path not in v.split(":"):
2749-
_logger.warning(
2750-
"The environment variable "
2751-
"'DYLD_FALLBACK_LIBRARY_PATH' does not contain "
2752-
"the '{new_path}' path in its value. This will make "
2753-
"PyTensor use a slow version of BLAS. Update "
2754-
"'DYLD_FALLBACK_LIBRARY_PATH' to contain the "
2755-
"said value, this will disable this warning."
2756-
)
2757-
2758-
use_unix_epd = False
2759-
if use_unix_epd:
2760-
return " ".join(
2761-
["-L%s" % os.path.join(sys.prefix, "lib")]
2762-
+ ["-l%s" % l for l in blas_info["libraries"]]
2763-
)
2764-
2765-
# Canopy
2766-
if "Canopy" in sys.prefix:
2767-
subsub = "lib"
2768-
if sys.platform == "win32":
2769-
subsub = "Scripts"
2770-
lib_path = os.path.join(sys.base_prefix, subsub)
2771-
if not os.path.exists(lib_path):
2772-
# Old logic to find the path. I don't think we still
2773-
# need it, but I don't have the time to test all
2774-
# installation configuration. So I keep this as a fall
2775-
# back in case the current expectation don't work.
2776-
2777-
# This old logic don't work when multiple version of
2778-
# Canopy is installed.
2779-
p = os.path.join(sys.base_prefix, "..", "..", "appdata")
2780-
assert os.path.exists(p), "Canopy changed the location of MKL"
2781-
lib_paths = os.listdir(p)
2782-
# Try to remove subdir that can't contain MKL
2783-
for sub in lib_paths:
2784-
if not os.path.exists(os.path.join(p, sub, subsub)):
2785-
lib_paths.remove(sub)
2786-
assert len(lib_paths) == 1, (
2787-
"Unexpected case when looking for Canopy MKL libraries",
2788-
p,
2789-
lib_paths,
2790-
[os.listdir(os.path.join(p, sub)) for sub in lib_paths],
2791-
)
2792-
lib_path = os.path.join(p, lib_paths[0], subsub)
2793-
assert os.path.exists(lib_path), "Canopy changed the location of MKL"
2794-
2795-
if sys.platform == "linux2" or sys.platform == "darwin":
2796-
return " ".join(
2797-
["-L%s" % lib_path] + ["-l%s" % l for l in blas_info["libraries"]]
2798-
)
2799-
elif sys.platform == "win32":
2800-
return " ".join(
2801-
['-L"%s"' % lib_path]
2802-
+
2803-
# Why on Windows, the library used are not the
2804-
# same as what is in blas_info['libraries']?
2805-
[f"-l{l}" for l in ("mk2_core", "mk2_intel_thread", "mk2_rt")]
2806-
)
28072719

2808-
# MKL
2809-
# If mkl can be imported then use it. On conda:
2810-
# "conda install mkl-service" installs the Python wrapper and
2811-
# the low-level C libraries as well as optimised version of
2812-
# numpy and scipy.
2813-
try:
2814-
import mkl # noqa
2815-
except ImportError:
2816-
pass
2817-
else:
2818-
# This branch is executed if no exception was raised
2819-
if sys.platform == "win32":
2820-
lib_path = os.path.join(sys.prefix, "Library", "bin")
2821-
flags = [f'-L"{lib_path}"']
2822-
else:
2823-
lib_path = blas_info.get("library_dirs", [])
2824-
flags = []
2825-
if lib_path:
2826-
flags = [f"-L{lib_path[0]}"]
2827-
if "2018" in mkl.get_version_string():
2828-
thr = "mkl_gnu_thread"
2829-
else:
2830-
thr = "mkl_intel_thread"
2831-
base_flags = list(flags)
2832-
flags += [f"-l{l}" for l in ("mkl_core", thr, "mkl_rt")]
2833-
res = try_blas_flag(flags)
2834-
2835-
if not res and sys.platform == "win32" and thr == "mkl_gnu_thread":
2836-
# Check if it would work for intel OpenMP on windows
2837-
flags = base_flags + [
2838-
f"-l{l}" for l in ("mkl_core", "mkl_intel_thread", "mkl_rt")
2720+
def check_required_file(paths, required_regexs):
2721+
libs = []
2722+
for req in required_regexs:
2723+
found = False
2724+
for path in paths:
2725+
m = re.search(req, path.name)
2726+
if m:
2727+
libs.append((str(path.parent), m.string[slice(*m.span())]))
2728+
found = True
2729+
break
2730+
if not found:
2731+
raise RuntimeError(f"Required file {req} not found")
2732+
return libs
2733+
2734+
def get_cxx_library_dirs():
2735+
cmd = f"{config.cxx} -print-search-dirs"
2736+
p = subprocess_Popen(
2737+
cmd,
2738+
stdout=subprocess.PIPE,
2739+
stderr=subprocess.PIPE,
2740+
stdin=subprocess.PIPE,
2741+
shell=True,
2742+
)
2743+
(stdout, stderr) = p.communicate(input=b"")
2744+
maybe_lib_dirs = [
2745+
[pathlib.Path(p).resolve() for p in line[len("libraries: =") :].split(":")]
2746+
for line in stdout.decode(sys.stdout.encoding).splitlines()
2747+
if line.startswith("libraries: =")
2748+
][0]
2749+
return [str(d) for d in maybe_lib_dirs if d.exists() and d.is_dir()]
2750+
2751+
def check_libs(
2752+
all_libs, required_libs, extra_compile_flags=None, cxx_library_dirs=None
2753+
):
2754+
if cxx_library_dirs is None:
2755+
cxx_library_dirs = []
2756+
if extra_compile_flags is None:
2757+
extra_compile_flags = []
2758+
found_libs = check_required_file(
2759+
all_libs,
2760+
required_libs,
2761+
)
2762+
path_quote = '"' if sys.platform == "win32" else ""
2763+
libdir_ldflags = list(
2764+
dict.fromkeys(
2765+
[
2766+
f"-L{path_quote}{lib_path}{path_quote}"
2767+
for lib_path, _ in found_libs
2768+
if lib_path not in cxx_library_dirs
28392769
]
2840-
res = try_blas_flag(flags)
2841-
2842-
if res:
2843-
check_mkl_openmp()
2844-
return res
2845-
2846-
flags.extend(["-Wl,-rpath," + l for l in blas_info.get("library_dirs", [])])
2847-
res = try_blas_flag(flags)
2848-
if res:
2849-
check_mkl_openmp()
2850-
maybe_add_to_os_environ_pathlist("PATH", lib_path[0])
2851-
return res
2852-
2853-
# to support path that includes spaces, we need to wrap it with double quotes on Windows
2854-
path_wrapper = '"' if os.name == "nt" else ""
2855-
ret = (
2856-
# TODO: the Gemm op below should separate the
2857-
# -L and -l arguments into the two callbacks
2858-
# that CLinker uses for that stuff. for now,
2859-
# we just pass the whole ldflags as the -l
2860-
# options part.
2861-
[
2862-
f"-L{path_wrapper}{l}{path_wrapper}"
2863-
for l in blas_info.get("library_dirs", [])
2864-
]
2865-
+ [f"-l{l}" for l in blas_info.get("libraries", [])]
2866-
+ blas_info.get("extra_link_args", [])
2770+
)
28672771
)
2868-
# For some very strange reason, we need to specify -lm twice
2869-
# to get mkl to link correctly. I have no idea why.
2870-
if any("mkl" in fl for fl in ret):
2871-
ret.extend(["-lm", "-lm"])
2872-
res = try_blas_flag(ret)
2873-
if res:
2874-
if "mkl" in res:
2875-
check_mkl_openmp()
2876-
return res
28772772

2878-
# If we are using conda and can't reuse numpy blas, then doing
2879-
# the fallback and test -lblas could give slow computation, so
2880-
# warn about this.
2881-
for warn in warn_record:
2882-
_logger.warning(warn)
2883-
del warn_record
2884-
2885-
# Some environment don't have the lib dir in LD_LIBRARY_PATH.
2886-
# So add it.
2887-
ret.extend(["-Wl,-rpath," + l for l in blas_info.get("library_dirs", [])])
2888-
res = try_blas_flag(ret)
2773+
flags = (
2774+
libdir_ldflags
2775+
+ [f"-l{lib_name}" for _, lib_name in found_libs]
2776+
+ extra_compile_flags
2777+
)
2778+
res = try_blas_flag(flags)
28892779
if res:
2890-
if "mkl" in res:
2780+
if any("mkl" in flag for flag in flags):
28912781
check_mkl_openmp()
28922782
return res
2783+
else:
2784+
raise RuntimeError(f"Supplied flags {flags} failed to compile")
28932785

2894-
# Add sys.prefix/lib to the runtime search path. On
2895-
# non-system installations of Python that use the
2896-
# system linker, this is generally necessary.
2897-
if sys.platform in ("linux", "darwin"):
2898-
lib_path = os.path.join(sys.prefix, "lib")
2899-
ret.append("-Wl,-rpath," + lib_path)
2900-
res = try_blas_flag(ret)
2901-
if res:
2902-
if "mkl" in res:
2903-
check_mkl_openmp()
2904-
return res
2905-
2906-
except KeyError:
2786+
_std_lib_dirs = std_lib_dirs()
2787+
if len(_std_lib_dirs) > 0:
2788+
rpath = _std_lib_dirs[0]
2789+
else:
2790+
rpath = None
2791+
2792+
cxx_library_dirs = get_cxx_library_dirs()
2793+
searched_library_dirs = cxx_library_dirs + _std_lib_dirs
2794+
all_libs = [
2795+
l
2796+
for path in [
2797+
pathlib.Path(library_dir)
2798+
for library_dir in searched_library_dirs
2799+
if pathlib.Path(library_dir).exists()
2800+
]
2801+
for l in path.iterdir()
2802+
if l.suffix in {".so", ".dll", ".dylib"}
2803+
]
2804+
2805+
if rpath is not None:
2806+
maybe_add_to_os_environ_pathlist("PATH", rpath)
2807+
try:
2808+
# 1. Try to use MKL with INTEL OpenMP threading
2809+
return check_libs(
2810+
all_libs,
2811+
required_libs=[
2812+
"mkl_core",
2813+
"mkl_rt",
2814+
"mkl_intel_thread",
2815+
"iomp5",
2816+
"pthread",
2817+
],
2818+
extra_compile_flags=[f"-Wl,-rpath,{rpath}"] if rpath is not None else [],
2819+
cxx_library_dirs=cxx_library_dirs,
2820+
)
2821+
except Exception:
29072822
pass
2908-
2909-
# Even if we could not detect what was used for numpy, or if these
2910-
# libraries are not found, most Linux systems have a libblas.so
2911-
# readily available. We try to see if that's the case, rather
2912-
# than disable blas. To test it correctly, we must load a program.
2913-
# Otherwise, there could be problem in the LD_LIBRARY_PATH.
2914-
return try_blas_flag(["-lblas"])
2823+
try:
2824+
# 2. Try to use MKL with GNU OpenMP threading
2825+
return check_libs(
2826+
all_libs,
2827+
required_libs=["mkl_core", "mkl_rt", "mkl_gnu_thread", "gomp", "pthread"],
2828+
extra_compile_flags=[f"-Wl,-rpath,{rpath}"] if rpath is not None else [],
2829+
cxx_library_dirs=cxx_library_dirs,
2830+
)
2831+
except Exception:
2832+
pass
2833+
try:
2834+
# 3. Try to use LAPACK + BLAS
2835+
return check_libs(
2836+
all_libs,
2837+
required_libs=["lapack", "blas", "cblas", "m"],
2838+
extra_compile_flags=[f"-Wl,-rpath,{rpath}"] if rpath is not None else [],
2839+
cxx_library_dirs=cxx_library_dirs,
2840+
)
2841+
except Exception:
2842+
pass
2843+
try:
2844+
# 4. Try to use BLAS alone
2845+
return check_libs(
2846+
all_libs,
2847+
required_libs=["blas", "cblas"],
2848+
extra_compile_flags=[f"-Wl,-rpath,{rpath}"] if rpath is not None else [],
2849+
cxx_library_dirs=cxx_library_dirs,
2850+
)
2851+
except Exception:
2852+
pass
2853+
try:
2854+
# 5. Try to use openblas
2855+
return check_libs(
2856+
all_libs,
2857+
required_libs=["openblas", "gfortran", "gomp", "m"],
2858+
extra_compile_flags=["-fopenmp", f"-Wl,-rpath,{rpath}"]
2859+
if rpath is not None
2860+
else ["-fopenmp"],
2861+
cxx_library_dirs=cxx_library_dirs,
2862+
)
2863+
except Exception:
2864+
pass
2865+
return ""
29152866

29162867

29172868
def add_blas_configvars():

0 commit comments

Comments
 (0)