Skip to content

Android input and output tag JNI #11057

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 8 commits into
base: main
Choose a base branch
from
Draft

Android input and output tag JNI #11057

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
0d38a99
input and output tag API
kirklandsign May 22, 2025
d956392
Fix
kirklandsign May 22, 2025
13916d5
Merge remote-tracking branch 'origin/main' into android-metadata-2
kirklandsign May 22, 2025
9e414ec
Merge remote-tracking branch 'origin/main' into android-metadata-2
kirklandsign May 22, 2025
0dc2ba5
Linter
kirklandsign May 22, 2025
6660276
fmt
kirklandsign May 22, 2025
a3a02ef
format
kirklandsign May 22, 2025
06f8e7b
Update
kirklandsign May 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions ...ecutorch_android/src/androidTest/java/org/pytorch/executorch/ModuleInstrumentationTest.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
27 changes: 26 additions & 1 deletion ...nsion/android/executorch_android/src/main/java/org/pytorch/executorch/MethodMetadata.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
}
}
14 changes: 13 additions & 1 deletion extension/android/executorch_android/src/main/java/org/pytorch/executorch/Module.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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.
*
Expand Down
34 changes: 31 additions & 3 deletions extension/android/jni/jni_layer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -453,10 +453,10 @@ 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 =
Expand All @@ -471,6 +471,32 @@ 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);

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);

for (int i = 0; i < num_outputs; i++) {
Copy link
Preview

Copilot AI May 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The variable 'i' is declared before the for-loop and then re-declared in the loop, which shadows the outer variable. Removing the redundant declaration will prevent potential confusion.

Suggested change
for (int i = 0; i < num_outputs; i++) {
for (i = 0; i < num_outputs; i++) {

Copilot uses AI. Check for mistakes.

ret->pin()[i] = static_cast<uint32_t>(method_meta.output_tag(i).get());
}
return ret;
}

static void registerNatives() {
registerHybrid({
makeNativeMethod("initHybrid", ExecuTorchJni::initHybrid),
Expand All @@ -480,6 +506,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),
});
}
};
Expand Down
Loading