(tvm) branch main updated: [JVM] Automatic Compatibility of JVM AttachCurrentThread (#16987)

[ruihangl]

This is an automated email from the ASF dual-hosted git repository.

ruihangl pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 4403379e39 [JVM] Automatic Compatibility of JVM AttachCurrentThread 
(#16987)
4403379e39 is described below

commit 4403379e3949e3339958ee01a41b9ece9c48ea8d
Author: Tianqi Chen <tqc...@users.noreply.github.com>
AuthorDate: Sun May 12 13:10:12 2024 -0400

    [JVM] Automatic Compatibility of JVM AttachCurrentThread (#16987)
    
    Different JDK may have different signature for AttachCurrentThread.
    This can cause issues for example between code for android and normal java.
    This PR uses a helper class to enable compact with both.
---
 .../src/main/native/org_apache_tvm_native_c_api.cc | 29 ++++++++++++++--------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/jvm/native/src/main/native/org_apache_tvm_native_c_api.cc 
b/jvm/native/src/main/native/org_apache_tvm_native_c_api.cc
index f86191d45b..09522381f1 100644
--- a/jvm/native/src/main/native/org_apache_tvm_native_c_api.cc
+++ b/jvm/native/src/main/native/org_apache_tvm_native_c_api.cc
@@ -222,17 +222,30 @@ JNIEXPORT jint JNICALL 
Java_org_apache_tvm_LibInfo_tvmFuncCall(JNIEnv* env, jobj
   return ret;
 }
 
+// A helper object to take in JNIEnv ptr
+// and allow automatic casting to both JNIEnv** and void**
+// Background: different version of JDK may choose to have one signature
+// or another for the case of AttachCurrentThread
+// we use this universal helper object to enable compatibility with both
+class JNIEnvPtrHelper {
+ public:
+  explicit JNIEnvPtrHelper(JNIEnv** penv) : penv_(penv) {}
+
+  operator JNIEnv**() { return penv_; }
+
+  operator void**() { return reinterpret_cast<void**>(penv_); }
+
+ private:
+  JNIEnv** penv_;
+};
+
 // Callback function
 extern "C" int funcInvokeCallback(TVMValue* args, int* typeCodes, int numArgs,
                                   TVMRetValueHandle ret, void* resourceHandle) 
{
   JNIEnv* env;
   int jniStatus = _jvm->GetEnv(reinterpret_cast<void**>(&env), 
JNI_VERSION_1_6);
   if (jniStatus == JNI_EDETACHED) {
-#ifdef TVM4J_ANDROID
-    _jvm->AttachCurrentThread(&env, nullptr);
-#else
-    _jvm->AttachCurrentThread(reinterpret_cast<void**>(&env), nullptr);
-#endif
+    _jvm->AttachCurrentThread(JNIEnvPtrHelper(&env), nullptr);
   } else {
     CHECK(jniStatus == JNI_OK);
   }
@@ -305,11 +318,7 @@ extern "C" void funcFreeCallback(void* resourceHandle) {
   JNIEnv* env;
   int jniStatus = _jvm->GetEnv(reinterpret_cast<void**>(&env), 
JNI_VERSION_1_6);
   if (jniStatus == JNI_EDETACHED) {
-#ifdef TVM4J_ANDROID
-    _jvm->AttachCurrentThread(&env, nullptr);
-#else
-    _jvm->AttachCurrentThread(reinterpret_cast<void**>(&env), nullptr);
-#endif
+    _jvm->AttachCurrentThread(JNIEnvPtrHelper(&env), nullptr);
   } else {
     CHECK(jniStatus == JNI_OK);
   }