Added Intel(R) initialization routine, adopted from Intel's patches to NumPy, that changes the interpretation of unset MKL_THREADING_LAYER to mean use Intel(R) OpenMP unless Gnu OMP was detected to be already loaded, in which case use Gnu OpenMP instead

Author: oleksandr-pavlyk

mkl/__init__.py

@@ -23,6 +23,33 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+import sys
+
+class RTLD_for_MKL():
+    def __init__(self):
+        self.saved_rtld = None
+
+    def __enter__(self):
+        import ctypes
+        try:
+            self.saved_rtld = sys.getdlopenflags()
+            # python loads libraries with RTLD_LOCAL, but MKL requires RTLD_GLOBAL
+            # pre-load MKL with RTLD_GLOBAL before loading the native extension
+            sys.setdlopenflags(self.saved_rtld | ctypes.RTLD_GLOBAL)
+        except AttributeError:
+            pass
+        del ctypes
+
+    def __exit__(self, *args):
+        if self.saved_rtld:
+            sys.setdlopenflags(self.saved_rtld)
+            self.saved_rtld = None
+
+with RTLD_for_MKL():
+    from . import _mklinit
+
+del RTLD_for_MKL
+del sys
 
 from ._py_mkl_service import *
 

mkl/_mklinitmodule.c

@@ -0,0 +1,153 @@
+/* -*- c -*- */
+
+/*
+ * This is a dummy module whose purpose is to get distutils to generate the
+ * configuration files before the libraries are made.
+ */
+
+#if (defined(USING_MKL_RT) && defined(__linux__))
+#define FORCE_PRELOADING 1
+#define _GNU_SOURCE 1
+#include <dlfcn.h>
+#include <string.h>
+#undef _GNU_SOURCE
+#endif
+
+#include <Python.h>
+#if PY_MAJOR_VERSION >= 3
+#define IS_PY3K
+#endif
+#include "mkl.h"
+
+static struct PyMethodDef methods[] = {
+    {NULL, NULL, 0, NULL}
+};
+
+static inline void _set_mkl_ilp64() {
+#ifdef USING_MKL_RT
+    int i = mkl_set_interface_layer(MKL_INTERFACE_ILP64);
+#endif
+    return;
+}
+
+static inline void _set_mkl_lp64() {
+#ifdef USING_MKL_RT
+    int i = mkl_set_interface_layer(MKL_INTERFACE_LP64);
+#endif
+    return;
+}
+
+static void _preload_threading_layer() {
+#if FORCE_PRELOADING
+#define VERBOSE(...) if(verbose) printf("mkl-service + Intel(R) MKL: " __VA_ARGS__)
+#define SET_MTLAYER(L) do {                                      \
+            VERBOSE("setting Intel(R) MKL to use " #L " OpenMP runtime\n");  \
+            mkl_set_threading_layer(MKL_THREADING_##L);          \
+            setenv("MKL_THREADING_LAYER", #L, 0);                \
+        } while(0)
+#define PRELOAD(lib) do {                                        \
+            VERBOSE("preloading %s runtime\n", lib);             \
+            dlopen(lib, RTLD_LAZY|RTLD_GLOBAL);                  \
+        } while(0)
+    /*
+     * The following is the pseudo-code skeleton for reinterpreting unset MKL_THREADING_LAYER
+     *       
+     *       if MKL_THREADING_LAYER is empty
+     *            if kmp_calloc (or a suitable symbol identified by Terry) is loaded,
+     *                  we are using Intel (R) OpenMP, i.e. reinterpret as implicit value of INTEL
+     *            otherwise check if other Open MP is loaded by checking get_omp_num_threads symbol
+     *                  if not loaded: 
+     *                         assume INTEL, and force loading of IOMP5
+     *                  if loaded:
+     *                         if Gnu OMP, set MKL_THREADING_LAYER=GNU, and call set_mkl_threading_layer(MKL_THREADING_GNU)
+     *                         if other vendors?
+     *       if MKL_THREADING_LAYER is INTEL
+     *             force loading of iomp, to preempt possibility of other modules loading other OMP library before MKL is actually used
+     *
+     *       should we treat other possible values of MKL_THREADING_LAYER specially?
+     *
+     */
+
+    const char *libiomp = "libiomp5.so";
+    const char *verbose = getenv("MKL_VERBOSE");
+    const char *mtlayer = getenv("MKL_THREADING_LAYER");
+    void *omp = dlsym(RTLD_DEFAULT, "omp_get_num_threads");
+    const char *omp_name = "(unidentified)";
+    const char *iomp = NULL; /* non-zero indicates Intel OpenMP is loaded */
+    Dl_info omp_info;
+
+    if(verbose && (verbose[0] == 0 || atoi(verbose) == 0))
+        verbose = NULL;
+
+    VERBOSE("THREADING LAYER: %s\n", mtlayer);
+
+    if(omp) {
+        if(dladdr(omp, &omp_info)) {
+            omp_name = basename(omp_info.dli_fname); /* GNU version doesn't modify argument */
+            iomp = strstr(omp_name, libiomp);
+        }
+        VERBOSE("%s OpenMP runtime %s is already loaded\n", iomp?"Intel(R)":"Other vendor", omp_name);
+    }
+    if(!mtlayer || mtlayer[0] == 0) {                /* unset or empty */
+      if(omp) {                                      /* if OpenMP runtime is loaded */
+        if(iomp)                                     /* if Intel runtime is loaded */
+            SET_MTLAYER(INTEL);
+        else                                         /* otherwise, assume it is GNU OpenMP */
+            SET_MTLAYER(GNU);
+      } else {                                       /* nothing is loaded */
+          SET_MTLAYER(INTEL);
+          PRELOAD(libiomp);
+      }
+    } else if(strcasecmp(mtlayer, "intel") == 0) {   /* Intel runtime is requested */
+        if(omp && !iomp) {
+            fprintf(stderr, "Error: mkl-service + Intel(R) MKL: MKL_THREADING_LAYER=INTEL is incompatible with %s library."
+                            "\n\tTry to import numpy first or set the threading layer accordingly. "
+                            "Set MKL_SERVICE_FORCE_INTEL to force it.\n", omp_name);
+            if(!getenv("MKL_SERVICE_FORCE_INTEL"))
+                exit(1);
+        } else
+            PRELOAD(libiomp);
+    }
+#endif
+    return;
+}
+
+static inline void _set_mkl_interface() {
+    _set_mkl_lp64();
+    _preload_threading_layer();
+}
+
+#if defined(IS_PY3K)
+static struct PyModuleDef moduledef = {
+    PyModuleDef_HEAD_INIT,
+    "mklinit",
+    NULL,
+    -1,
+    methods,
+    NULL,
+    NULL,
+    NULL,
+    NULL
+};
+#endif
+
+/* Initialization function for the module */
+#if defined(IS_PY3K)
+PyMODINIT_FUNC PyInit__mklinit(void) {
+    PyObject *m;
+
+    _set_mkl_interface();
+    m = PyModule_Create(&moduledef);
+    if (!m) {
+        return NULL;
+    }
+
+    return m;
+}
+#else
+PyMODINIT_FUNC
+init_mklinit(void) {
+    _set_mkl_interface();
+    Py_InitModule("_mklinit", methods);
+}
+#endif

mkl/setup.py

@@ -39,6 +39,11 @@ def configuration(parent_package='', top_path=None):
     mkl_library_dirs = mkl_info.get('library_dirs', [])
     mkl_libraries = mkl_info.get('libraries', ['mkl_rt'])
 
+    defs = []
+    if any(['mkl_rt' in li for li in mkl_libraries]):
+        #libs += ['dl'] - by default on Linux
+        defs += [('USING_MKL_RT', None)]
+
     try:
         from Cython.Build import cythonize
         sources = [join(pdir, '_mkl_service.pyx')]
@@ -50,6 +55,19 @@ def configuration(parent_package='', top_path=None):
             raise ValueError(str(e) + '. ' +
                              'Cython is required to build the initial .c file.')
 
+    config.add_extension(
+        '_mklinit',
+        sources=['_mklinitmodule.c'],
+        define_macros=defs,
+        include_dirs=[mkl_include_dirs],
+        library_dirs=[mkl_library_dirs],
+        libraries=mkl_libraries,
+        extra_compile_args=[
+            '-DNDEBUG'
+            # '-g', '-O2', '-Wall',
+        ]
+    )
+
     config.add_extension(
         '_py_mkl_service',
         sources=sources,