Merge pull request #723 from evalf/support-multiple-mkl-versions

support multiple versions of MKL

.github/workflows/test.yaml

@@ -49,8 +49,10 @@ jobs:
           - {name: "python 3.8", os: ubuntu-latest, python-version: "3.8", matrix-backend: numpy, nprocs: 1}
           - {name: "python 3.9", os: ubuntu-latest, python-version: "3.9", matrix-backend: numpy, nprocs: 1}
           - {name: "scipy matrix", os: ubuntu-latest, python-version: "3.10", matrix-backend: scipy, nprocs: 1}
-          - {name: "mkl matrix", os: ubuntu-latest, python-version: "3.10", matrix-backend: mkl, nprocs: 1}
-          - {name: "mkl matrix parallel", os: ubuntu-latest, python-version: "3.10", matrix-backend: mkl, nprocs: 2}
+          - {name: "mkl linux", os: ubuntu-latest, python-version: "3.10", matrix-backend: mkl, nprocs: 1}
+          - {name: "mkl linux parallel", os: ubuntu-latest, python-version: "3.10", matrix-backend: mkl, nprocs: 2}
+          - {name: "mkl windows", os: windows-latest, python-version: "3.10", matrix-backend: mkl, nprocs: 1}
+          - {name: "mkl macos", os: macos-latest, python-version: "3.10", matrix-backend: mkl, nprocs: 1}
           - {name: "parallel", os: ubuntu-latest, python-version: "3.10", matrix-backend: numpy, nprocs: 2}
           - {name: "numpy 1.17", os: ubuntu-latest, python-version: "3.7", matrix-backend: numpy, nprocs: 1, numpy-version: ==1.17}
           - {name: "tensorial", os: ubuntu-latest, python-version: "3.10", matrix-backend: numpy, nprocs: 1, tensorial: test}

devtools/gha/configure_mkl.py

@@ -1,25 +1,38 @@
 from .. import log
+import re
 import sys
 import site
 import os
+from pathlib import Path
 
-libsubdir = 'lib'
-prefixes = list(site.PREFIXES)
-if hasattr(site, 'getuserbase'):
-    prefixes.append(site.getuserbase())
-
-candidates = [os.path.join(prefix, libsubdir, 'libmkl_rt.so' + ext) for prefix in prefixes for ext in ('.1', '.2')]
-for path in candidates:
-    if os.path.exists(path):
-        break
+if sys.platform == 'linux':
+    libsubdir = 'lib',
+    re_libmkl = re.compile('libmkl_rt[.]so[.][0-9]+')
+elif sys.platform == 'darwin':
+    libsubdir = 'lib',
+    re_libmkl = re.compile('libmkl_rt[.][0-9]+[.]dylib')
+elif sys.platform == 'win32':
+    libsubdir = 'Library', 'bin'
+    re_libmkl = re.compile('mkl_rt[.][0-9]+[.]dll')
 else:
-    log.error('cannot find any of {}'.format(' '.join(candidates)))
+    log.error(f'unsupported platform: {sys.platform}')
     raise SystemExit(1)
 
-lib = os.path.splitext(path)[0]
-if not os.path.exists(lib):
-    os.symlink(path, lib)
+prefixes = list(map(Path, site.PREFIXES))
+if hasattr(site, 'getuserbase'):
+    prefixes.append(Path(site.getuserbase()))
+
+libdirs = {libdir := prefix.joinpath(*libsubdir).resolve() for prefix in prefixes}
+libs = {file for libdir in libdirs if libdir.is_dir() for file in libdir.iterdir() if re_libmkl.match(file.name)}
+if len(libs) == 0:
+    log.error('cannot find MKL in any of {}'.format(', '.join(map(str, libdirs))))
+    raise SystemExit(1)
+elif len(libs) != 1:
+    log.error('found MKL at more than one location: {}'.format(', '.join(map(str, libs))))
+    raise SystemExit(1)
+else:
+    lib, = libs
+    log.info(f'using MKL at {lib}')
 
-ld_library_path = os.pathsep.join(filter(None, (os.environ.get('LD_LIBRARY_PATH', ''), os.path.dirname(path))))
 with open(os.environ['GITHUB_ENV'], 'a') as f:
-    print('LD_LIBRARY_PATH={}'.format(ld_library_path), file=f)
+    print(f'NUTILS_MATRIX_MKL_LIB={lib}', file=f)

nutils/matrix/_mkl.py

@@ -1,14 +1,21 @@
 from ._base import Matrix, MatrixError, BackendNotAvailable
 from .. import numeric, _util as util, warnings
 from contextlib import contextmanager
-from ctypes import c_int, byref
+from ctypes import c_int, byref, CDLL
 import treelog as log
 import os
 import numpy
 
-libmkl = util.loadlib(linux='libmkl_rt.so', darwin='libmkl_rt.dylib', win32='mkl_rt.dll')
-if not libmkl:
-    raise BackendNotAvailable('the Intel MKL matrix backend requires libmkl to be installed (try: pip install mkl)')
+libmkl_path = os.environ.get('NUTILS_MATRIX_MKL_LIB', None)
+if libmkl_path:
+    libmkl = CDLL(libmkl_path)
+else:
+    for v in '.2', '.1', '':
+        libmkl = util.loadlib(linux=f'libmkl_rt.so{v}', darwin=f'libmkl_rt{v}.dylib', win32=f'mkl_rt{v}.dll')
+        if libmkl:
+            break
+    else:
+        raise BackendNotAvailable('the Intel MKL matrix backend requires libmkl to be installed (try: pip install mkl)')
 
 
 def assemble(data, index, shape):