Add api to explicitly load a method (#2408)

Summary:
We lazy load methods during the first execute (forward). However, sometimes we want to explicitly load the method before the first inference, so that we get consistent performance between all runs.

Adding a method so that user can load method, if they wish. Otherwise, they can still depend on execute to lazy load for them.

Pull Request resolved: #2408

Test Plan:
CI

Explicitly call this:
```
 diff --git a/xplat/executorch/examples/demo-apps/android/ExecuTorchDemo/app/src/main/java/com/example/executorchdemo/MainActivity.java b/xplat/executorch/examples/demo-apps/android/ExecuTorchDemo/app/src/main/java/com/example/executorchdemo/MainActivity.java
 --- a/xplat/executorch/examples/demo-apps/android/ExecuTorchDemo/app/src/main/java/com/example/executorchdemo/MainActivity.java
+++ b/xplat/executorch/examples/demo-apps/android/ExecuTorchDemo/app/src/main/java/com/example/executorchdemo/MainActivity.java
@@ -101,6 +101,7 @@
     try {
       mModule =
           Module.load(MainActivity.assetFilePath(getApplicationContext(), "dl3_xnnpack_fp32.pte"));
+      mModule.loadMethod("forward");

     } catch (IOException e) {
       Log.e("ImageSegmentation", "Error reading assets", e);
@@ -136,6 +137,7 @@
               mModule =
                   Module.load(
                       MainActivity.assetFilePath(getApplicationContext(), "dl3_xnnpack_fp32.pte"));
+              mModule.loadMethod("forward");
             } catch (IOException e) {
               Log.e("ImageSegmentation", "Error reading assets", e);
               finish();
```
```
buck install fbandroid/mode/opt xplat/executorch/examples/demo-apps/android/ExecuTorchDemo/app/src/main:ExecuTorchDemo
```
and run app

Differential Revision: D54868635

Pulled By: kirklandsign

fbshipit-source-id: a680862020b41a73e1aa2c68a84cd68c22cb18c9

Author: kirklandsign

extension/android/jni/jni_layer.cpp

@@ -275,6 +275,10 @@ class ExecuTorchJni : public facebook::jni::HybridClass<ExecuTorchJni> {
     return execute_method(methodName->toStdString(), jinputs);
   }
 
+  jint load_method(facebook::jni::alias_ref<jstring> methodName) {
+    return static_cast<jint>(module_->load_method(methodName->toStdString()));
+  }
+
   facebook::jni::local_ref<facebook::jni::JArrayClass<JEValue>> execute_method(
       std::string method,
       facebook::jni::alias_ref<
@@ -343,6 +347,7 @@ class ExecuTorchJni : public facebook::jni::HybridClass<ExecuTorchJni> {
         makeNativeMethod("initHybrid", ExecuTorchJni::initHybrid),
         makeNativeMethod("forward", ExecuTorchJni::forward),
         makeNativeMethod("execute", ExecuTorchJni::execute),
+        makeNativeMethod("loadMethod", ExecuTorchJni::load_method),
     });
   }
 };

extension/android/src/main/java/org/pytorch/executorch/INativePeer.java

@@ -18,4 +18,11 @@ interface INativePeer {
 
   /** Run an arbitrary method on the module */
   EValue[] execute(String methodName, EValue... inputs);
+
+  /**
+   * Load a method on this module.
+   *
+   * @return the Error code if there was an error loading the method
+   */
+  int loadMethod(String methodName);
 }

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

@@ -67,6 +67,19 @@ public EValue[] execute(String methodName, EValue... inputs) {
     return mNativePeer.execute(methodName, 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
+   * synchronous, and will block until the method is loaded. Therefore, it is recommended to call
+   * this on a background thread. However, users need to make sure that they don't execute before
+   * this function returns.
+   *
+   * @return the Error code if there was an error loading the method
+   */
+  public int loadMethod(String methodName) {
+    return mNativePeer.loadMethod(methodName);
+  }
+
   /**
    * Explicitly destroys the native torch::jit::Module. Calling this method is not required, as the
    * native object will be destroyed when this object is garbage-collected. However, the timing of

extension/android/src/main/java/org/pytorch/executorch/NativePeer.java

@@ -38,4 +38,7 @@ public void resetNative() {
 
   @DoNotStrip
   public native EValue[] execute(String methodName, EValue... inputs);
+
+  @DoNotStrip
+  public native int loadMethod(String methodName);
 }