Get rid of NativePeer.java

examples/demo-apps/android/LlamaDemo/docs/delegates/xnnpack_README.md

@@ -172,7 +172,7 @@ This is running the shell script which configures the required core ExecuTorch,
 ```
 bash examples/demo-apps/android/LlamaDemo/download_prebuilt_lib.sh
 ```
-The prebuilt AAR library contains the Java library and the JNI binding for NativePeer.java and ExecuTorch native library, including core ExecuTorch runtime libraries, XNNPACK backend, Portable kernels, Optimized kernels, and Quantized kernels. It comes with two ABI variants, arm64-v8a and x86_64.
+The prebuilt AAR library contains the Java library and the JNI binding for Module.java and ExecuTorch native library, including core ExecuTorch runtime libraries, XNNPACK backend, Portable kernels, Optimized kernels, and Quantized kernels. It comes with two ABI variants, arm64-v8a and x86_64.
 If you need to use other dependencies (like tokenizer), please build from the local machine option.
 
 ## Run the Android Demo App

extension/android/BUCK

@@ -9,7 +9,6 @@ non_fbcode_target(_kind = fb_android_library,
         "executorch_android/src/main/java/org/pytorch/executorch/DType.java",
         "executorch_android/src/main/java/org/pytorch/executorch/EValue.java",
         "executorch_android/src/main/java/org/pytorch/executorch/Module.java",
-        "executorch_android/src/main/java/org/pytorch/executorch/NativePeer.java",
         "executorch_android/src/main/java/org/pytorch/executorch/Tensor.java",
         "executorch_android/src/main/java/org/pytorch/executorch/annotations/Experimental.java",
     ],

extension/android/executorch_android/src/main/java/org/pytorch/executorch/Module.java

@@ -9,6 +9,8 @@
 package org.pytorch.executorch;
 
 import android.util.Log;
+import com.facebook.jni.HybridData;
+import com.facebook.jni.annotations.DoNotStrip;
 import com.facebook.soloader.nativeloader.NativeLoader;
 import com.facebook.soloader.nativeloader.SystemDelegate;
 import java.io.File;
@@ -24,6 +26,14 @@
 @Experimental
 public class Module {
 
+  static {
+    if (!NativeLoader.isInitialized()) {
+      NativeLoader.init(new SystemDelegate());
+    }
+    // Loads libexecutorch.so from jniLibs
+    NativeLoader.loadLibrary("executorch");
+  }
+
   /** Load mode for the module. Load the whole file as a buffer. */
   public static final int LOAD_MODE_FILE = 0;
 
@@ -36,10 +46,17 @@ public class Module {
   /** Load mode for the module. Use memory locking and ignore errors. */
   public static final int LOAD_MODE_MMAP_USE_MLOCK_IGNORE_ERRORS = 3;
 
-  /** Reference to the NativePeer object of this module. */
-  private NativePeer mNativePeer;
+  private final HybridData mHybridData;
+
+  @DoNotStrip
+  private static native HybridData initHybrid(
+      String moduleAbsolutePath, int loadMode, int initHybrid);
 
-  /** Lock protecting the non-thread safe methods in NativePeer. */
+  private Module(String moduleAbsolutePath, int loadMode, int numThreads) {
+    mHybridData = initHybrid(moduleAbsolutePath, loadMode, numThreads);
+  }
+
+  /** Lock protecting the non-thread safe methods in mHybridData. */
   private Lock mLock = new ReentrantLock();
 
   /**
@@ -63,14 +80,11 @@ public static Module load(final String modelPath, int loadMode) {
    * @return new {@link org.pytorch.executorch.Module} object which owns the model module.
    */
   public static Module load(final String modelPath, int loadMode, int numThreads) {
-    if (!NativeLoader.isInitialized()) {
-      NativeLoader.init(new SystemDelegate());
-    }
     File modelFile = new File(modelPath);
     if (!modelFile.canRead() || !modelFile.isFile()) {
       throw new RuntimeException("Cannot load model path " + modelPath);
     }
-    return new Module(new NativePeer(modelPath, loadMode, numThreads));
+    return new Module(modelPath, loadMode, numThreads);
   }
 
   /**
@@ -83,10 +97,6 @@ public static Module load(final String modelPath) {
     return load(modelPath, LOAD_MODE_FILE);
   }
 
-  Module(NativePeer nativePeer) {
-    this.mNativePeer = nativePeer;
-  }
-
   /**
    * Runs the 'forward' method of this module with the specified arguments.
    *
@@ -96,16 +106,7 @@ public static Module load(final String modelPath) {
    * @return return value from the 'forward' method.
    */
   public EValue[] forward(EValue... inputs) {
-    try {
-      mLock.lock();
-      if (mNativePeer == null) {
-        Log.e("ExecuTorch", "Attempt to use a destroyed module");
-        return new EValue[0];
-      }
-      return mNativePeer.forward(inputs);
-    } finally {
-      mLock.unlock();
-    }
+    return execute("forward", inputs);
   }
 
   /**
@@ -118,16 +119,19 @@ public EValue[] forward(EValue... inputs) {
   public EValue[] execute(String methodName, EValue... inputs) {
     try {
       mLock.lock();
-      if (mNativePeer == null) {
+      if (!mHybridData.isValid()) {
         Log.e("ExecuTorch", "Attempt to use a destroyed module");
         return new EValue[0];
       }
-      return mNativePeer.execute(methodName, inputs);
+      return executeNative(methodName, inputs);
     } finally {
       mLock.unlock();
     }
   }
 
+  @DoNotStrip
+  private native EValue[] executeNative(String methodName, EValue... inputs);
+
   /**
    * Load a method on this module. This might help with the first time inference performance,
    * because otherwise the method is loaded lazily when it's execute. Note: this function is
@@ -140,31 +144,36 @@ public EValue[] execute(String methodName, EValue... inputs) {
   public int loadMethod(String methodName) {
     try {
       mLock.lock();
-      if (mNativePeer == null) {
+      if (!mHybridData.isValid()) {
         Log.e("ExecuTorch", "Attempt to use a destroyed module");
         return 0x2; // InvalidState
       }
-      return mNativePeer.loadMethod(methodName);
+      return loadMethodNative(methodName);
     } finally {
       mLock.unlock();
     }
   }
 
+  @DoNotStrip
+  private native int loadMethodNative(String methodName);
+
   /**
    * Returns the names of the methods in a certain method.
    *
    * @param methodName method name to query
    * @return an array of backend name
    */
-  public String[] getUsedBackends(String methodName) {
-    return mNativePeer.getUsedBackends(methodName);
-  }
+  @DoNotStrip
+  public native String[] getUsedBackends(String methodName);
 
   /** Retrieve the in-memory log buffer, containing the most recent ExecuTorch log entries. */
   public String[] readLogBuffer() {
-    return mNativePeer.readLogBuffer();
+    return readLogBufferNative();
   }
 
+  @DoNotStrip
+  private native String[] readLogBufferNative();
+
   /**
    * Dump the ExecuTorch ETRecord file to /data/local/tmp/result.etdump.
    *
@@ -173,9 +182,8 @@ public String[] readLogBuffer() {
    * @return true if the etdump was successfully written, false otherwise.
    */
   @Experimental
-  public boolean etdump() {
-    return mNativePeer.etdump();
-  }
+  @DoNotStrip
+  public native boolean etdump();
 
   /**
    * Explicitly destroys the native Module object. Calling this method is not required, as the
@@ -186,13 +194,11 @@ public boolean etdump() {
   public void destroy() {
     if (mLock.tryLock()) {
       try {
-        mNativePeer.resetNative();
+        mHybridData.resetNative();
       } finally {
-        mNativePeer = null;
         mLock.unlock();
       }
     } else {
-      mNativePeer = null;
       Log.w(
           "ExecuTorch",
           "Destroy was called while the module was in use. Resources will not be immediately"

extension/android/jni/jni_layer.cpp

@@ -222,7 +222,7 @@ class ExecuTorchJni : public facebook::jni::HybridClass<ExecuTorchJni> {
   std::unique_ptr<Module> module_;
 
  public:
-  constexpr static auto kJavaDescriptor = "Lorg/pytorch/executorch/NativePeer;";
+  constexpr static auto kJavaDescriptor = "Lorg/pytorch/executorch/Module;";
 
   static facebook::jni::local_ref<jhybriddata> initHybrid(
       facebook::jni::alias_ref<jclass>,
@@ -273,13 +273,6 @@ class ExecuTorchJni : public facebook::jni::HybridClass<ExecuTorchJni> {
 #endif
   }
 
-  facebook::jni::local_ref<facebook::jni::JArrayClass<JEValue>> forward(
-      facebook::jni::alias_ref<
-          facebook::jni::JArrayClass<JEValue::javaobject>::javaobject>
-          jinputs) {
-    return execute_method("forward", jinputs);
-  }
-
   facebook::jni::local_ref<facebook::jni::JArrayClass<JEValue>> execute(
       facebook::jni::alias_ref<jstring> methodName,
       facebook::jni::alias_ref<
@@ -461,10 +454,9 @@ class ExecuTorchJni : public facebook::jni::HybridClass<ExecuTorchJni> {
   static void registerNatives() {
     registerHybrid({
         makeNativeMethod("initHybrid", ExecuTorchJni::initHybrid),
-        makeNativeMethod("forward", ExecuTorchJni::forward),
-        makeNativeMethod("execute", ExecuTorchJni::execute),
-        makeNativeMethod("loadMethod", ExecuTorchJni::load_method),
-        makeNativeMethod("readLogBuffer", ExecuTorchJni::readLogBuffer),
+        makeNativeMethod("executeNative", ExecuTorchJni::execute),
+        makeNativeMethod("loadMethodNative", ExecuTorchJni::load_method),
+        makeNativeMethod("readLogBufferNative", ExecuTorchJni::readLogBuffer),
         makeNativeMethod("etdump", ExecuTorchJni::etdump),
         makeNativeMethod("getUsedBackends", ExecuTorchJni::getUsedBackends),
     });