1010#define _GNU_SOURCE 1
1111#include <dlfcn.h>
1212#include <string.h>
13+ #include <pthread.h>
1314#undef _GNU_SOURCE
1415#endif
1516
@@ -29,30 +30,47 @@ static struct PyMethodDef methods[] = {
2930#define MKL_SERVICE_INLINE inline
3031#endif
3132
32- static MKL_SERVICE_INLINE void _set_mkl_ilp64 ();
33- static MKL_SERVICE_INLINE void _set_mkl_lp64 ();
34- static MKL_SERVICE_INLINE void _set_mkl_interface ();
33+ static MKL_SERVICE_INLINE void _set_mkl_ilp64 (void );
34+ static MKL_SERVICE_INLINE void _set_mkl_lp64 (void );
35+ static MKL_SERVICE_INLINE void _set_mkl_interface (void );
36+
37+ static const char * mtlayer ;
38+ static const char * verbose ;
3539
36- static void _preload_threading_layer () {
3740#if FORCE_PRELOADING
3841#define VERBOSE (...) if(verbose) printf("mkl-service + Intel(R) MKL: " __VA_ARGS__)
39- #define SET_MTLAYER (L ) do { \
42+
43+ static void restore_mtlayer (void ) {
44+ if (mtlayer ) {
45+ VERBOSE ("Re-setting Intel(R) MKL_THREADING_LAYER=%s for the forked process\n" , mtlayer );
46+ setenv ("MKL_THREADING_LAYER" , mtlayer , 1 );
47+ } else {
48+ VERBOSE ("Unsetting Intel(R) MKL_THREADING_LAYER variable for the forked process \n" );
49+ unsetenv ("MKL_THREADING_LAYER" );
50+ }
51+ }
52+ #endif
53+
54+ static void _preload_threading_layer (void ) {
55+ #if FORCE_PRELOADING
56+ #define SET_MTLAYER (L ) do { \
4057 VERBOSE("setting Intel(R) MKL to use " #L " OpenMP runtime\n"); \
41- mkl_set_threading_layer(MKL_THREADING_##L); \
42- setenv("MKL_THREADING_LAYER", #L, 0); \
58+ mkl_set_threading_layer(MKL_THREADING_##L); \
59+ setenv("MKL_THREADING_LAYER", #L, 0); \
60+ pthread_atfork(NULL, NULL, &restore_mtlayer); \
4361 } while(0)
44- #define PRELOAD (lib ) do { \
45- VERBOSE("preloading %s runtime\n", lib); \
46- dlopen(lib, RTLD_LAZY|RTLD_GLOBAL); \
62+ #define PRELOAD (lib ) do { \
63+ VERBOSE("preloading %s runtime\n", lib); \
64+ dlopen(lib, RTLD_LAZY|RTLD_GLOBAL); \
4765 } while(0)
4866 /*
4967 * The following is the pseudo-code skeleton for reinterpreting unset MKL_THREADING_LAYER
50- *
68+ *
5169 * if MKL_THREADING_LAYER is empty
5270 * if kmp_calloc (or a suitable symbol identified by Terry) is loaded,
5371 * we are using Intel(R) OpenMP, i.e. reinterpret as implicit value of INTEL
5472 * otherwise check if other Open MP is loaded by checking get_omp_num_threads symbol
55- * if not loaded:
73+ * if not loaded:
5674 * assume INTEL, and force loading of IOMP5
5775 * if loaded:
5876 * if Gnu OMP, set MKL_THREADING_LAYER=GNU, and call set_mkl_threading_layer(MKL_THREADING_GNU)
@@ -65,8 +83,14 @@ static void _preload_threading_layer() {
6583 */
6684
6785 const char * libiomp = "libiomp5.so" ;
68- const char * verbose = getenv ("MKL_VERBOSE" );
69- const char * mtlayer = getenv ("MKL_THREADING_LAYER" );
86+ verbose = getenv ("MKL_VERBOSE" );
87+ mtlayer = getenv ("MKL_THREADING_LAYER" );
88+
89+ /* Use of RTLD_DEFAULT handler is to indicate that symbol is being lookup-up among symbols
90+ * presently known to this process.
91+ *
92+ * See: https://pubs.opengroup.org/onlinepubs/9699919799/functions/dlsym.html
93+ */
7094 void * omp = dlsym (RTLD_DEFAULT , "omp_get_num_threads" );
7195 const char * omp_name = "(unidentified)" ;
7296 const char * iomp = NULL ; /* non-zero indicates Intel(R) OpenMP is loaded */
@@ -108,21 +132,21 @@ static void _preload_threading_layer() {
108132 return ;
109133}
110134
111- static MKL_SERVICE_INLINE void _set_mkl_ilp64 () {
135+ static MKL_SERVICE_INLINE void _set_mkl_ilp64 (void ) {
112136#ifdef USING_MKL_RT
113- int i = mkl_set_interface_layer (MKL_INTERFACE_ILP64 );
137+ mkl_set_interface_layer (MKL_INTERFACE_ILP64 );
114138#endif
115139 return ;
116140}
117141
118- static MKL_SERVICE_INLINE void _set_mkl_lp64 () {
142+ static MKL_SERVICE_INLINE void _set_mkl_lp64 (void ) {
119143#ifdef USING_MKL_RT
120- int i = mkl_set_interface_layer (MKL_INTERFACE_LP64 );
144+ mkl_set_interface_layer (MKL_INTERFACE_LP64 );
121145#endif
122146 return ;
123147}
124148
125- static MKL_SERVICE_INLINE void _set_mkl_interface () {
149+ static MKL_SERVICE_INLINE void _set_mkl_interface (void ) {
126150 _set_mkl_lp64 ();
127151 _preload_threading_layer ();
128152}
0 commit comments