Add Java runtime API and JNI layer

Author: keyprocedure

extension/android/CMakeLists.txt

@@ -64,7 +64,7 @@ set(executorch_DIR ${CMAKE_CURRENT_BINARY_DIR}/../../lib/cmake/ExecuTorch)
 find_package(executorch CONFIG REQUIRED)
 target_link_options_shared_lib(executorch)
 
-add_library(executorch_jni SHARED jni/jni_layer.cpp jni/log.cpp)
+add_library(executorch_jni SHARED jni/jni_layer.cpp jni/log.cpp jni/jni_layer_runtime.cpp)
 
 set(link_libraries)
 list(

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

@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+package org.pytorch.executorch;
+
+import org.pytorch.executorch.annotations.Experimental;
+import com.facebook.jni.annotations.DoNotStrip;
+import com.facebook.soloader.nativeloader.NativeLoader;
+import com.facebook.soloader.nativeloader.SystemDelegate;
+
+/**
+ * Java API to access native runtime functionality on Android.
+ *
+ * <p>Warning: These APIs are experimental and subject to change without notice.
+ */
+@Experimental
+@DoNotStrip
+public class Runtime {
+    private static boolean initialized = false;
+
+    static {
+        if (!NativeLoader.isInitialized()) {
+            NativeLoader.init(new SystemDelegate());
+        }
+        // Loads libexecutorch.so from jniLibs
+        NativeLoader.loadLibrary("executorch");
+        initialized = true;
+    }
+
+    public static boolean isInitialized() {
+        return initialized;
+    }
+
+    @DoNotStrip
+    public static native String[] getRegisteredOps();
+
+    @DoNotStrip
+    public static native String[] getRegisteredBackends();
+}

extension/android/jni/jni_layer.cpp

@@ -491,9 +491,11 @@ extern void register_natives_for_llm();
 // No op if we don't build LLM
 void register_natives_for_llm() {}
 #endif
+extern void register_natives_for_runtime();
 JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM* vm, void*) {
   return facebook::jni::initialize(vm, [] {
     executorch::extension::ExecuTorchJni::registerNatives();
     register_natives_for_llm();
+    register_natives_for_runtime();
   });
 }

extension/android/jni/jni_layer_runtime.cpp

@@ -0,0 +1,69 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <jni.h>
+#include <fbjni/fbjni.h>
+
+#include <executorch/runtime/kernel/operator_registry.h>
+#include <executorch/runtime/backend/interface.h>
+
+namespace executorch_jni {
+namespace runtime = ::executorch::ET_RUNTIME_NAMESPACE;
+
+class AndroidRuntimeJni : public facebook::jni::JavaClass<AndroidRuntimeJni> {
+ public:
+  constexpr static const char* kJavaDescriptor = "Lorg/pytorch/executorch/Runtime;";
+
+  static void registerNatives() {
+    javaClassStatic()->registerNatives({
+        makeNativeMethod("getRegisteredOps", AndroidRuntimeJni::getRegisteredOps),
+        makeNativeMethod("getRegisteredBackends", AndroidRuntimeJni::getRegisteredBackends),
+    });
+  }
+
+  // Returns a string array of all registered ops
+  static facebook::jni::local_ref<facebook::jni::JArrayClass<jstring>>
+  getRegisteredOps(facebook::jni::alias_ref<jclass>) {
+    auto kernels = runtime::get_registered_kernels();
+    auto result = facebook::jni::JArrayClass<jstring>::newArray(kernels.size());
+
+    for (size_t i = 0; i < kernels.size(); ++i) {
+        auto op = facebook::jni::make_jstring(kernels[i].name_);
+        result->setElement(i, op.get());
+    }
+
+    return result;
+  }
+
+  // Returns a string array of all registered backends
+  static facebook::jni::local_ref<facebook::jni::JArrayClass<jstring>>
+  getRegisteredBackends(facebook::jni::alias_ref<jclass>) {
+    int num_backends = runtime::get_num_registered_backends();
+    auto result = facebook::jni::JArrayClass<jstring>::newArray(num_backends);
+
+    for (int i = 0; i < num_backends; ++i) {
+        auto name_result = runtime::get_backend_name(i);
+        const char* name = "";
+
+        if (name_result.ok()) {
+            name = *name_result;
+        }
+
+        auto backend_str = facebook::jni::make_jstring(name);
+        result->setElement(i, backend_str.get());
+    }
+
+    return result;
+  }
+};
+
+} // namespace executorch_jni
+
+void register_natives_for_runtime() {
+  executorch_jni::AndroidRuntimeJni::registerNatives();
+}