Check and set MKL THREADS on the fly (fix issue pyscf#1102)

Author: sunqm

conda/build.sh

@@ -11,7 +11,10 @@ set -x -e
 # find pyscf/lib/deps -name "*cint*" -exec rm {} \+
 # rm pyscf-2.0-depsa-openblas.tar.gz
 
+# C extensions must be installed with sequential BLAS library
+# https://pyscf.org/install.html#using-optimized-blas
+export CMAKE_CONFIGURE_ARGS="-DWITH_F12=OFF -DBLA_VENDOR=Intel10_64lp_seq"
+
 # env PYTHON not defined in certain conda-build version
 # $PYTHON -m pip install . -vv
-export CMAKE_CONFIGURE_ARGS="-DWITH_F12=OFF"
 pip install -v --prefix=$PREFIX .

pyscf/__init__.py

@@ -39,12 +39,6 @@
 
 import os
 import sys
-# Avoid too many threads being created in OMP loops.
-# See issue https://github.com/pyscf/pyscf/issues/317
-if 'OPENBLAS_NUM_THREADS' not in os.environ:
-    os.environ['OPENBLAS_NUM_THREADS'] = '1'
-if 'MKL_NUM_THREADS' not in os.environ:
-    os.environ['MKL_NUM_THREADS'] = '1'
 
 # Load modules which are developed as plugins of the namespace package
 PYSCF_EXT_PATH = os.getenv('PYSCF_EXT_PATH')

pyscf/lib/misc.py

@@ -37,6 +37,41 @@
 except ImportError:
     ThreadPoolExecutor = None
 
+if sys.platform.startswith('linux'):
+    # Avoid too many threads being created in OMP loops.
+    # See issue https://github.com/pyscf/pyscf/issues/317
+    try:
+        from elftools.elf.elffile import ELFFile
+    except ImportError:
+        pass
+    else:
+        def _ldd(so_file):
+            libs = []
+            with open(so_file, 'rb') as f:
+                elf = ELFFile(f)
+                for seg in elf.iter_segments():
+                    if seg.header.p_type != 'PT_DYNAMIC':
+                        continue
+                    for t in seg.iter_tags():
+                        if t.entry.d_tag == 'DT_NEEDED':
+                            libs.append(t.needed)
+                    break
+            return libs
+
+        so_file = os.path.abspath(os.path.join(__file__, '..', 'libnp_helper.so'))
+        for p in _ldd(so_file):
+            if 'mkl' in p and 'thread' in p:
+                warnings.warn(f'PySCF C exteions are incompatible with {p}. '
+                              'MKL_NUM_THREADS is set to 1')
+                os.environ['MKL_NUM_THREADS'] = '1'
+                break
+            elif 'openblasp' in p or 'openblaso' in p:
+                warnings.warn(f'PySCF C exteions are incompatible with {p}. '
+                              'OPENBLAS_NUM_THREADS is set to 1')
+                os.environ['OPENBLAS_NUM_THREADS'] = '1'
+                break
+        del p, so_file, _ldd
+
 from pyscf.lib import param
 from pyscf import __config__