[js/rn] Supoort New Architecture

js/.eslintrc.js

@@ -175,7 +175,22 @@ module.exports = {
     }
   }, {
     files: ['react_native/lib/**/*.ts'], rules: {
-      '@typescript-eslint/naming-convention': 'off'
+      '@typescript-eslint/naming-convention': 'off',
+    }
+  }, {
+    files: ['react_native/lib/NativeOnnxruntime.ts'], rules: {
+      'unicorn/filename-case': 'off',
+      '@typescript-eslint/ban-types': [
+        'error',
+        {
+          types: {
+            // NOTE: We got issue like https://github.com/facebook/react-native/issues/36431
+            // So we have to use `{}` type here.
+            '{}': false,
+          },
+          extendDefaults: true,
+        }
+      ]
     }
   }, {
     files: ['react_native/scripts/**/*.ts'], rules: {

js/react_native/.gitignore

@@ -25,13 +25,15 @@ DerivedData
 *.ipa
 *.xcuserstate
 project.xcworkspace
+xcshareddata
 
 # Android/IJ
 #
 .idea
 .gradle
 local.properties
 android.iml
+.cxx
 
 # Cocoapods
 #

js/react_native/android/build.gradle

@@ -14,6 +14,14 @@ buildscript {
 
 apply plugin: 'com.android.library'
 
+def isNewArchitectureEnabled() {
+  return rootProject.hasProperty("newArchEnabled") && rootProject.getProperty("newArchEnabled") == "true"
+}
+
+if (isNewArchitectureEnabled()) {
+  apply plugin: "com.facebook.react"
+}
+
 def getExtOrDefault(name) {
   return rootProject.ext.has(name) ? rootProject.ext.get(name) : project.properties['OnnxruntimeModule_' + name]
 }
@@ -90,6 +98,7 @@ android {
         abiFilters (*reactNativeArchitectures())
       }
     }
+    buildConfigField "boolean", "IS_NEW_ARCHITECTURE_ENABLED", isNewArchitectureEnabled().toString()
   }
 
   if (rootProject.hasProperty("ndkPath")) {
@@ -115,6 +124,7 @@ android {
       "META-INF",
       "META-INF/**",
       "**/libjsi.so",
+      "**/libc++_shared.so",
     ]
   }
 
@@ -139,6 +149,12 @@ android {
       } else {
         java.exclude '**/OnnxruntimeExtensionsEnabled.java'
       }
+
+      if (isNewArchitectureEnabled()) {
+        java.srcDirs += ['src/newarch']
+      } else {
+        java.srcDirs += ['src/oldarch']
+      }
     }
   }
 }
@@ -238,3 +254,11 @@ dependencies {
     implementation "com.microsoft.onnxruntime:onnxruntime-extensions-android:latest.integration@aar"
   }
 }
+
+if (isNewArchitectureEnabled()) {
+  react {
+    jsRootDir = file("../lib/")
+    libraryName = "OnnxruntimeSpec"
+    codegenJavaPackageName = "ai.onnxruntime.reactnative"
+  }
+}

js/react_native/android/src/androidTest/java/ai/onnxruntime/reactnative/OnnxruntimeModuleTest.java

@@ -72,6 +72,7 @@ public void onnxruntime_module() throws Exception {
 
       OnnxruntimeModule ortModule = new OnnxruntimeModule(reactContext);
       ortModule.blobModule = blobModule;
+      ortModule.checkBlobModule();
       String sessionKey = "";
 
       // test loadModel()
@@ -82,7 +83,7 @@ public void onnxruntime_module() throws Exception {
 
           JavaOnlyMap options = new JavaOnlyMap();
           try {
-            ReadableMap resultMap = ortModule.loadModel(modelBuffer, options);
+            ReadableMap resultMap = ortModule.getOnnxruntime().loadModel(modelBuffer, options);
             sessionKey = resultMap.getString("key");
             ReadableArray inputNames = resultMap.getArray("inputNames");
             ReadableArray outputNames = resultMap.getArray("outputNames");
@@ -132,7 +133,7 @@ public void onnxruntime_module() throws Exception {
         options.putBoolean("encodeTensorData", true);
 
         try {
-          ReadableMap resultMap = ortModule.run(sessionKey, inputDataMap, outputNames, options);
+          ReadableMap resultMap = ortModule.getOnnxruntime().run(sessionKey, inputDataMap, outputNames, options);
 
           ReadableMap outputMap = resultMap.getMap("output");
           for (int i = 0; i < 2; ++i) {
@@ -151,7 +152,7 @@ public void onnxruntime_module() throws Exception {
       }
 
       // test dispose
-      ortModule.dispose(sessionKey);
+      ortModule.getOnnxruntime().dispose(sessionKey);
     } finally {
       mockSession.finishMocking();
     }
@@ -166,6 +167,7 @@ public void onnxruntime_module_append_nnapi() throws Exception {
 
       OnnxruntimeModule ortModule = new OnnxruntimeModule(reactContext);
       ortModule.blobModule = blobModule;
+      ortModule.checkBlobModule();
       String sessionKey = "";
 
       // test loadModel() with nnapi ep options
@@ -182,7 +184,7 @@ public void onnxruntime_module_append_nnapi() throws Exception {
         options.putArray("executionProviders", epArray);
 
         try {
-          ReadableMap resultMap = ortModule.loadModel(modelBuffer, options);
+          ReadableMap resultMap = ortModule.getOnnxruntime().loadModel(modelBuffer, options);
           sessionKey = resultMap.getString("key");
           ReadableArray inputNames = resultMap.getArray("inputNames");
           ReadableArray outputNames = resultMap.getArray("outputNames");
@@ -195,7 +197,7 @@ public void onnxruntime_module_append_nnapi() throws Exception {
           Assert.fail(e.getMessage());
         }
       }
-      ortModule.dispose(sessionKey);
+      ortModule.getOnnxruntime().dispose(sessionKey);
     } finally {
       mockSession.finishMocking();
     }

js/react_native/android/src/main/java/ai/onnxruntime/reactnative/OnnxruntimeModule.java → js/react_native/android/src/main/java/ai/onnxruntime/reactnative/Onnxruntime.java

@@ -15,22 +15,16 @@
 import android.net.Uri;
 import android.os.Build;
 import android.util.Log;
-import androidx.annotation.NonNull;
 import androidx.annotation.RequiresApi;
 import com.facebook.react.bridge.Arguments;
-import com.facebook.react.bridge.LifecycleEventListener;
-import com.facebook.react.bridge.Promise;
 import com.facebook.react.bridge.ReactApplicationContext;
-import com.facebook.react.bridge.ReactContextBaseJavaModule;
-import com.facebook.react.bridge.ReactMethod;
 import com.facebook.react.bridge.ReadableArray;
 import com.facebook.react.bridge.ReadableMap;
 import com.facebook.react.bridge.ReadableType;
 import com.facebook.react.bridge.WritableArray;
 import com.facebook.react.bridge.WritableMap;
 import com.facebook.react.modules.blob.BlobModule;
 import java.io.BufferedReader;
-import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.io.Reader;
@@ -46,7 +40,7 @@
 import java.util.stream.Stream;
 
 @RequiresApi(api = Build.VERSION_CODES.N)
-public class OnnxruntimeModule extends ReactContextBaseJavaModule implements LifecycleEventListener {
+public class Onnxruntime {
   private static ReactApplicationContext reactContext;
 
   private static OrtEnvironment ortEnvironment = OrtEnvironment.getEnvironment();
@@ -59,104 +53,13 @@ private static String getNextSessionKey() {
     return key;
   }
 
-  protected BlobModule blobModule;
+  private BlobModule blobModule;
 
-  public OnnxruntimeModule(ReactApplicationContext context) {
-    super(context);
-    reactContext = context;
-  }
+  public Onnxruntime(ReactApplicationContext context) { reactContext = context; }
 
-  @NonNull
-  @Override
-  public String getName() {
-    return "Onnxruntime";
-  }
+  protected void setBlobModule(BlobModule blobModule) { this.blobModule = blobModule; }
 
-  public void checkBlobModule() {
-    if (blobModule == null) {
-      blobModule = getReactApplicationContext().getNativeModule(BlobModule.class);
-      if (blobModule == null) {
-        throw new RuntimeException("BlobModule is not initialized");
-      }
-    }
-  }
-
-  /**
-   * React native binding API to load a model using given uri.
-   *
-   * @param uri a model file location
-   * @param options onnxruntime session options
-   * @param promise output returning back to react native js
-   * @note the value provided to `promise` includes a key representing the session.
-   *       when run() is called, the key must be passed into the first parameter.
-   */
-  @ReactMethod
-  public void loadModel(String uri, ReadableMap options, Promise promise) {
-    try {
-      WritableMap resultMap = loadModel(uri, options);
-      promise.resolve(resultMap);
-    } catch (Exception e) {
-      promise.reject("Failed to load model \"" + uri + "\": " + e.getMessage(), e);
-    }
-  }
-
-  /**
-   * React native binding API to load a model using blob object that data stored in BlobModule.
-   *
-   * @param data the blob object
-   * @param options onnxruntime session options
-   * @param promise output returning back to react native js
-   * @note the value provided to `promise` includes a key representing the session.
-   *       when run() is called, the key must be passed into the first parameter.
-   */
-  @ReactMethod
-  public void loadModelFromBlob(ReadableMap data, ReadableMap options, Promise promise) {
-    try {
-      checkBlobModule();
-      String blobId = data.getString("blobId");
-      byte[] bytes = blobModule.resolve(blobId, data.getInt("offset"), data.getInt("size"));
-      blobModule.remove(blobId);
-      WritableMap resultMap = loadModel(bytes, options);
-      promise.resolve(resultMap);
-    } catch (Exception e) {
-      promise.reject("Failed to load model from buffer: " + e.getMessage(), e);
-    }
-  }
-
-  /**
-   * React native binding API to dispose a session.
-   *
-   * @param key session key representing a session given at loadModel()
-   * @param promise output returning back to react native js
-   */
-  @ReactMethod
-  public void dispose(String key, Promise promise) {
-    try {
-      dispose(key);
-      promise.resolve(null);
-    } catch (OrtException e) {
-      promise.reject("Failed to dispose session: " + e.getMessage(), e);
-    }
-  }
-
-  /**
-   * React native binding API to run a model using given uri.
-   *
-   * @param key session key representing a session given at loadModel()
-   * @param input an input tensor
-   * @param output an output names to be returned
-   * @param options onnxruntime run options
-   * @param promise output returning back to react native js
-   */
-  @ReactMethod
-  public void run(String key, ReadableMap input, ReadableArray output, ReadableMap options, Promise promise) {
-    try {
-      WritableMap resultMap = run(key, input, output, options);
-      promise.resolve(resultMap);
-    } catch (Exception e) {
-      promise.reject("Fail to inference: " + e.getMessage(), e);
-    }
-  }
+  protected Map<String, OrtSession> getSessionMap() { return sessionMap; }
 
   /**
    * Load a model from raw resource directory.
@@ -259,8 +162,6 @@ public WritableMap run(String key, ReadableMap input, ReadableArray output, Read
 
     RunOptions runOptions = parseRunOptions(options);
 
-    checkBlobModule();
-
     long startTime = System.currentTimeMillis();
     Map<String, OnnxTensor> feed = new HashMap<>();
     Iterator<String> iterator = ortSession.getInputNames().iterator();
@@ -436,22 +337,4 @@ private RunOptions parseRunOptions(ReadableMap options) throws OrtException {
 
     return runOptions;
   }
-
-  @Override
-  public void onHostResume() {}
-
-  @Override
-  public void onHostPause() {}
-
-  @Override
-  public void onHostDestroy() {
-    for (String key : sessionMap.keySet()) {
-      try {
-        dispose(key);
-      } catch (Exception e) {
-        Log.e("onHostDestroy", "Failed to dispose session: " + key, e);
-      }
-    }
-    sessionMap.clear();
-  }
 }