Allow for BLAS, LAPACK, and ARPACK symbols hijacking

This opens up support for NVBlas which "hijacks" symbols by defining `dgemm_` for example. It then expects the caller to call `dgemm_` and the loader to load NVBlas's symbol instead of OpenBLAS's or Intel MKL's.
luhenry · May 9, 2021 · 7644183 · 7644183
1 parent 96828db
commit 7644183
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 16 deletions.
diff --git a/arpack/src/main/native/jni.c b/arpack/src/main/native/jni.c
@@ -2124,11 +2124,9 @@ jboolean get_system_property(JNIEnv *env, jstring key, jstring def, jstring *res
   return TRUE;
 }
 
-static void *libhandle;
-
 jboolean load_symbols(void) {
 #define LOAD_SYMBOL(name) \
-  name = dlsym(libhandle, #name); \
+  name = dlsym(NULL, #name); \
   if (!name) { \
     return FALSE; \
   }
@@ -2196,6 +2194,8 @@ jboolean load_symbols(void) {
  return TRUE;
 }
 
+static void *libhandle;
+
 jint JNI_OnLoad(JavaVM *vm, UNUSED void *reserved) {
   JNIEnv *env;
   if ((*vm)->GetEnv(vm, (void**)&env, JNI_VERSION_1_6) != JNI_OK) {
@@ -2270,7 +2270,7 @@ jint JNI_OnLoad(JavaVM *vm, UNUSED void *reserved) {
     return -1;
   }
 
-  libhandle = dlopen(name, RTLD_LAZY);
+  libhandle = dlopen(name, RTLD_LAZY | RTLD_GLOBAL);
   if (!libhandle) {
     return -1;
   }

diff --git a/blas/src/main/native/jni.c b/blas/src/main/native/jni.c
@@ -1974,11 +1974,9 @@ jboolean get_system_property(JNIEnv *env, jstring key, jstring def, jstring *res
   return TRUE;
 }
 
-static void *libhandle;
-
 jboolean load_symbols(void) {
 #define LOAD_SYMBOL(name) \
-  name = dlsym(libhandle, #name); \
+  name = dlsym(NULL, #name); \
   if (!name) { \
     return FALSE; \
   }
@@ -2055,6 +2053,8 @@ jboolean load_symbols(void) {
  return TRUE;
 }
 
+static void *libhandle;
+
 jint JNI_OnLoad(JavaVM *vm, UNUSED void *reserved) {
   JNIEnv *env;
   if ((*vm)->GetEnv(vm, (void**)&env, JNI_VERSION_1_6) != JNI_OK) {
@@ -2129,7 +2129,7 @@ jint JNI_OnLoad(JavaVM *vm, UNUSED void *reserved) {
     return -1;
   }
 
-  libhandle = dlopen(name, RTLD_LAZY);
+  libhandle = dlopen(name, RTLD_LAZY | RTLD_GLOBAL);
   if (!libhandle) {
     return -1;
   }

diff --git a/generator.py b/generator.py
@@ -370,11 +370,9 @@ def __init__(self, pkg, libname, *routines):
     print("}")
     print()
     # Print symbols loading
-    print("static void *libhandle;")
-    print()
     print("jboolean load_symbols(void) {")
     print("#define LOAD_SYMBOL(name) \\")
-    print("  name = dlsym(libhandle, #name); \\")
+    print("  name = dlsym(NULL, #name); \\")
     print("  if (!name) { \\")
     print("    return FALSE; \\")
     print("  }")
@@ -387,6 +385,8 @@ def __init__(self, pkg, libname, *routines):
     print("}")
     print()
     # Print JNI entry functions
+    print("static void *libhandle;")
+    print()
     print("jint JNI_OnLoad(JavaVM *vm, UNUSED void *reserved) {")
     print("  JNIEnv *env;")
     print("  if ((*vm)->GetEnv(vm, (void**)&env, JNI_VERSION_1_6) != JNI_OK) {")
@@ -461,7 +461,7 @@ def __init__(self, pkg, libname, *routines):
     print("    return -1;")
     print("  }")
     print("")
-    print("  libhandle = dlopen(name, RTLD_LAZY);")
+    print("  libhandle = dlopen(name, RTLD_LAZY | RTLD_GLOBAL);")
     print("  if (!libhandle) {")
     print("    return -1;")
     print("  }")

diff --git a/lapack/src/main/native/jni.c b/lapack/src/main/native/jni.c
@@ -27507,11 +27507,9 @@ jboolean get_system_property(JNIEnv *env, jstring key, jstring def, jstring *res
   return TRUE;
 }
 
-static void *libhandle;
-
 jboolean load_symbols(void) {
 #define LOAD_SYMBOL(name) \
-  name = dlsym(libhandle, #name); \
+  name = dlsym(NULL, #name); \
   if (!name) { \
     return FALSE; \
   }
@@ -28244,6 +28242,8 @@ jboolean load_symbols(void) {
  return TRUE;
 }
 
+static void *libhandle;
+
 jint JNI_OnLoad(JavaVM *vm, UNUSED void *reserved) {
   JNIEnv *env;
   if ((*vm)->GetEnv(vm, (void**)&env, JNI_VERSION_1_6) != JNI_OK) {
@@ -28318,7 +28318,7 @@ jint JNI_OnLoad(JavaVM *vm, UNUSED void *reserved) {
     return -1;
   }
 
-  libhandle = dlopen(name, RTLD_LAZY);
+  libhandle = dlopen(name, RTLD_LAZY | RTLD_GLOBAL);
   if (!libhandle) {
     return -1;
   }