input and output tag API

Author: kirklandsign

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

@@ -68,9 +68,8 @@ class ModuleE2ETest {
         inputStream.close()
 
         val module = Module.load(getTestFilePath("/mv3_xnnpack_fp32.pte"))
-        val expectedBackends = arrayOf("XnnpackBackend")
         Assert.assertArrayEquals(
-            expectedBackends,
+            arrayOf("XnnpackBackend"),
             module.getMethodMetadata("forward").getBackends(),
         )
     }

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

@@ -62,6 +62,8 @@ class ModuleInstrumentationTest {
 
         Assert.assertArrayEquals(arrayOf("forward"), module.getMethods())
         Assert.assertTrue(module.getMethodMetadata("forward").backends.isEmpty())
+        Assert.assertArrayEquals(intArrayOf(1, 1, 3), module.getMethodMetadata("forward").inputTags)
+        Assert.assertArrayEquals(intArrayOf(1), module.getMethodMetadata("forward").outputTags)
     }
 
     @Test

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

@@ -11,8 +11,9 @@
 /** Helper class to access the metadata for a method from a Module */
 public class MethodMetadata {
   private String mName;
-
   private String[] mBackends;
+  private int[] mInputTags;
+  private int[] mOutputTags;
 
   MethodMetadata setName(String name) {
     mName = name;
@@ -37,4 +38,28 @@ MethodMetadata setBackends(String[] backends) {
   public String[] getBackends() {
     return mBackends;
   }
+
+  /**
+   * @return Output tags
+   */
+  public int[] getOutputTags() {
+    return mOutputTags;
+  }
+
+  MethodMetadata setOutputTags(int[] outputTags) {
+    mOutputTags = outputTags;
+    return this;
+  }
+
+  /**
+   * @return Input tags
+   */
+  public int[] getInputTags() {
+    return mInputTags;
+  }
+
+  MethodMetadata setInputTags(int[] inputTags) {
+    mInputTags = inputTags;
+    return this;
+  }
 }

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

@@ -59,7 +59,13 @@ Map<String, MethodMetadata> populateMethodMeta() {
     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)));
+      metadata.put(
+          name,
+          new MethodMetadata()
+              .setName(name)
+              .setBackends(getUsedBackends(name))
+              .setInputTags(getInputTags(name))
+              .setOutputTags(getOutputTags(name)));
     }
 
     return metadata;
@@ -204,6 +210,12 @@ public String[] readLogBuffer() {
   @DoNotStrip
   private native String[] readLogBufferNative();
 
+  @DoNotStrip
+  private native int[] getInputTags(String method);
+
+  @DoNotStrip
+  private native int[] getOutputTags(String method);
+
   /**
    * Dump the ExecuTorch ETRecord file to /data/local/tmp/result.etdump.
    *

extension/android/jni/jni_layer.cpp

@@ -453,10 +453,11 @@ class ExecuTorchJni : public facebook::jni::HybridClass<ExecuTorchJni> {
 
   facebook::jni::local_ref<facebook::jni::JArrayClass<jstring>> getUsedBackends(
       facebook::jni::alias_ref<jstring> methodName) {
-    auto methodMeta = module_->method_meta(methodName->toStdString()).get();
+    auto method_meta =
+        module_->method_meta(methodName->toStdString()).get();
     std::unordered_set<std::string> backends;
-    for (auto i = 0; i < methodMeta.num_backends(); i++) {
-      backends.insert(methodMeta.get_backend_name(i).get());
+    for (auto i = 0; i < method_meta.num_backends(); i++) {
+      backends.insert(method_meta.get_backend_name(i).get());
     }
 
     facebook::jni::local_ref<facebook::jni::JArrayClass<jstring>> ret =
@@ -471,6 +472,36 @@ class ExecuTorchJni : public facebook::jni::HybridClass<ExecuTorchJni> {
     return ret;
   }
 
+  facebook::jni::local_ref<facebook::jni::JArrayInt> getInputTags(
+      facebook::jni::alias_ref<jstring> methodName) {
+    auto method_meta =
+        module_->method_meta(methodName->toStdString()).get();
+    auto num_inputs = method_meta.num_inputs();
+    facebook::jni::local_ref<facebook::jni::JArrayInt> ret =
+        facebook::jni::JArrayInt::newArray(num_inputs);
+
+    int i = 0;
+    for (int i = 0; i < num_inputs; i++) {
+      ret->pin()[i] = static_cast<uint32_t>(method_meta.input_tag(i).get());
+    }
+    return ret;
+  }
+
+  facebook::jni::local_ref<facebook::jni::JArrayInt> getOutputTags(
+      facebook::jni::alias_ref<jstring> methodName) {
+    auto method_meta =
+        module_->method_meta(methodName->toStdString()).get();
+    auto num_outputs = method_meta.num_outputs();
+    facebook::jni::local_ref<facebook::jni::JArrayInt> ret =
+        facebook::jni::JArrayInt::newArray(num_outputs);
+
+    int i = 0;
+    for (int i = 0; i < num_outputs; i++) {
+      ret->pin()[i] = static_cast<uint32_t>(method_meta.output_tag(i).get());
+    }
+    return ret;
+  }
+
   static void registerNatives() {
     registerHybrid({
         makeNativeMethod("initHybrid", ExecuTorchJni::initHybrid),
@@ -480,6 +511,8 @@ class ExecuTorchJni : public facebook::jni::HybridClass<ExecuTorchJni> {
         makeNativeMethod("etdump", ExecuTorchJni::etdump),
         makeNativeMethod("getMethods", ExecuTorchJni::getMethods),
         makeNativeMethod("getUsedBackends", ExecuTorchJni::getUsedBackends),
+        makeNativeMethod("getInputTags", ExecuTorchJni::getInputTags),
+        makeNativeMethod("getOutputTags", ExecuTorchJni::getOutputTags),
     });
   }
 };