Update MKL init.

fwk/wutils.py

@@ -133,34 +133,29 @@ def initDLL():
 
 # ------------------------------------------------------------------------------
 
-def initMKL(nthreads, verb=False):
+def initMKL(verb=False):
+    """Initilize the environment for MKL
+       Since we use the SDL for MKL, we need to configure it at runtime.
+       This is MUST be performed before any calls to MKL.
+       Could be improved by performing the init from the C++ and actually checking that the threading layer is the one we want, instead of using env var.
+    """
     import os
-    # we try to have full control over the threading environment
-    # => set MKL/OMP variables if (and only if) they are not set!
-    # number of "default" threads for MKL (available MKL domains: ALL, BLAS, VML, PARDISO)
-    if not 'MKL_NUM_THREADS' in os.environ:
-        os.environ['MKL_NUM_THREADS'] = '%d' % nthreads
-        os.environ['MKL_DOMAIN_NUM_THREADS'] = 'MKL_DOMAIN_ALL=%d' % nthreads
-    # threading interface
-    if not 'MKL_THREADING_LAYER' in os.environ:
-        if nthreads == 1:
-            os.environ['MKL_THREADING_LAYER'] = 'SEQUENTIAL'
-        else:
-            os.environ['MKL_THREADING_LAYER'] = 'TBB'
-    # interface (32 or 64 bits)
+    # Define threading interface
     if not 'MKL_INTERFACE_LAYER' in os.environ:
         os.environ['MKL_INTERFACE_LAYER'] = 'LP64' # 32 bits (default), ILP64 for 64 bits
-    # we do not want to use OpenMP
-    if not 'OMP_NUM_THREADS' in os.environ:
-        os.environ['OMP_NUM_THREADS'] = '1'
+    # Define threading layer
+    if not 'MKL_THREADING_LAYER' in os.environ:
+        os.environ['MKL_THREADING_LAYER'] = 'TBB' # sequential computation will be performed with 1 TBB thread
+    # Force number of OMP threads to 1, TBB threads are controlled from inside the C++ code
+    if not 'MKL_NUM_THREADS' in os.environ:
+        os.environ['MKL_NUM_THREADS'] = '1'
+        os.environ['MKL_DOMAIN_NUM_THREADS'] = 'MKL_DOMAIN_ALL=1'
+    # Force number of OMP threads to remain unchanged
     if not 'MKL_DYNAMIC' in os.environ:
         os.environ['MKL_DYNAMIC'] = 'FALSE'
-    if not 'OMP_DYNAMIC' in os.environ:
-        os.environ['OMP_DYNAMIC'] = 'FALSE'
-    # display variables
-    print(f'Envrionment initialized for MKL with {nthreads} threads.')
+    # Display variables
     if verb:
-        for s in ['OMP_', 'MKL', 'KMP']:
+        for s in ['OMP', 'MKL', 'KMP']:
             print('* %s environment:' % s)        
             for key, value in os.environ.items():
                 if s in key:
@@ -283,7 +278,7 @@ def run():
 
     # parse args
     args = parseargs()
-    initMKL(args.k, verb=True) # initialize threading layer and number of threads
+    initMKL(verb=True) # initialize MKL environment (threading interface, layer and number of threads)
     if args.fpe:
         fwk.enableFpe() # enables floating point exceptions at runtime
         print(ccolors.ANSI_YELLOW + 'Floating point exceptions will cause numpy to overflow!' + ccolors.ANSI_RESET)

init.py

@@ -2,3 +2,6 @@
 # Add this dir to the python path so that sub-dir can be imported as sub-module
 import os.path, sys
 sys.path.append(os.path.dirname(os.path.realpath(__file__)))
+# Initialize MKL environment
+from .fwk import wutils
+wutils.initMKL()