Add a new method metadata class, and allow users to get methods

Author: kirklandsign

extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/ModuleE2ETest.java

@@ -99,7 +99,7 @@ public void testXnnpackBackendRequired() throws IOException, URISyntaxException
 
         Module module = Module.load(getTestFilePath("/mv3_xnnpack_fp32.pte"));
         String[] expectedBackends = new String[] {"XnnpackBackend"};
-        assertArrayEquals(expectedBackends, module.getUsedBackends("forward"));
+        assertArrayEquals(expectedBackends, module.getMethodMetadata("forward").getBackends());
     }
 
     @Test

extension/android/executorch_android/src/androidTest/java/org/pytorch/executorch/ModuleInstrumentationTest.java

@@ -109,7 +109,7 @@ public void testModuleLoadMethodNonExistantMethod() throws IOException{
         assertEquals(loadMethod, INVALID_ARGUMENT);
     }
 
-    @Test
+    @Test(expected = RuntimeException.class)
     public void testNonPteFile() throws IOException{
         Module module = Module.load(getTestFilePath(NON_PTE_FILE_NAME));
 

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

@@ -8,13 +8,33 @@
 
 package org.pytorch.executorch;
 
-/**
- * Helper class to access the metadata for a method from a Module
- */
+/** Helper class to access the metadata for a method from a Module */
 public class MethodMetadata {
-    private String name;
+  private String mName;
+
+  private String[] mBackends;
+
+  MethodMetadata setName(String name) {
+    mName = name;
+    return this;
+  }
+
+  /**
+   * @return Method name
+   */
+  public String getName() {
+    return mName;
+  }
+
+  MethodMetadata setBackends(String[] backends) {
+    mBackends = backends;
+    return this;
+  }
 
-    public String getName() {
-        return name;
-    }
+  /**
+   * @return Backends used for this method
+   */
+  public String[] getBackends() {
+    return mBackends;
+  }
 }

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

@@ -9,12 +9,13 @@
 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;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
 import org.pytorch.executorch.annotations.Experimental;
@@ -49,19 +50,26 @@ public class Module {
 
   private final HybridData mHybridData;
 
-  private final MethodMetadata[] mMethodMetadata;
+  private final Map<String, MethodMetadata> mMethodMetadata;
 
   @DoNotStrip
   private static native HybridData initHybrid(String moduleAbsolutePath, int loadMode);
 
   private Module(String moduleAbsolutePath, int loadMode) {
     mHybridData = initHybrid(moduleAbsolutePath, loadMode);
 
-    populateMethodMeta();
+    mMethodMetadata = populateMethodMeta();
   }
 
-  void populateMethodMeta() {
-    mMethodMetadata = new MethodMetadata[1];
+  Map<String, MethodMetadata> populateMethodMeta() {
+    String[] methods = getMethods();
+    Map<String, MethodMetadata> metadata = new HashMap<String, MethodMetadata>();
+    for (int i = 0; i < methods.length; i++) {
+      String name = methods[i];
+      metadata.put(name, new MethodMetadata().setName(name).setBackends(getUsedBackends(name)));
+    }
+
+    return metadata;
   }
 
   /** Lock protecting the non-thread safe methods in mHybridData. */
@@ -153,13 +161,34 @@ public int loadMethod(String methodName) {
   private native int loadMethodNative(String methodName);
 
   /**
-   * Returns the names of the methods in a certain method.
+   * Returns the names of the backends in a certain method.
    *
    * @param methodName method name to query
    * @return an array of backend name
    */
   @DoNotStrip
-  public native String[] getUsedBackends(String methodName);
+  private native String[] getUsedBackends(String methodName);
+
+  /**
+   * Returns the names of methods.
+   *
+   * @return name of methods in this Module
+   */
+  @DoNotStrip
+  public native String[] getMethods();
+
+  /**
+   * Get the corresponding @MethodMetadata for a method
+   *
+   * @param name method name
+   * @return @MethodMetadata for this method
+   */
+  public MethodMetadata getMethodMetadata(String name) {
+    if (!mMethodMetadata.containsKey(name)) {
+      throw new RuntimeException("method " + name + "does not exist for this module");
+    }
+    return mMethodMetadata.get(name);
+  }
 
   /** Retrieve the in-memory log buffer, containing the most recent ExecuTorch log entries. */
   @DoNotStrip

extension/android/jni/jni_layer.cpp

@@ -429,6 +429,26 @@ class ExecuTorchJni : public facebook::jni::HybridClass<ExecuTorchJni> {
     return false;
   }
 
+  facebook::jni::local_ref<facebook::jni::JArrayClass<jstring>> getMethods() {
+    const auto& names_result = module_->method_names();
+    if (!names_result.ok()) {
+      facebook::jni::throwNewJavaException(
+          facebook::jni::gJavaLangIllegalArgumentException,
+          "Cannot get load module");
+    }
+    const auto& methods = names_result.get();
+    facebook::jni::local_ref<facebook::jni::JArrayClass<jstring>> ret =
+        facebook::jni::JArrayClass<jstring>::newArray(methods.size());
+    int i = 0;
+    for (auto s : methods) {
+      facebook::jni::local_ref<facebook::jni::JString> method_name =
+          facebook::jni::make_jstring(s.c_str());
+      (*ret)[i] = method_name;
+      i++;
+    }
+    return ret;
+  }
+
   facebook::jni::local_ref<facebook::jni::JArrayClass<jstring>> getUsedBackends(
       facebook::jni::alias_ref<jstring> methodName) {
     auto methodMeta = module_->method_meta(methodName->toStdString()).get();
@@ -456,6 +476,7 @@ class ExecuTorchJni : public facebook::jni::HybridClass<ExecuTorchJni> {
         makeNativeMethod("loadMethodNative", ExecuTorchJni::load_method),
         makeNativeMethod("readLogBuffer", ExecuTorchJni::readLogBuffer),
         makeNativeMethod("etdump", ExecuTorchJni::etdump),
+        makeNativeMethod("getMethods", ExecuTorchJni::getMethods),
         makeNativeMethod("getUsedBackends", ExecuTorchJni::getUsedBackends),
     });
   }