[android] test_app example linking to pytorch_android aar content (pytorch#39587)

IvanKobzarev · facebook-github-bot · commit 52a2adb3f482 · 2020-06-18T16:26:25.000-07:00
Summary: Pull Request resolved: pytorch#39587 Example of using direct linking to pytorch_jni library from aar and updating android/README.md with the tutorial how to do it. Adding `nativeBuild` dimension to `test_app`, using direct aar dependencies, as headers packaging is not landed yet, excluding `nativeBuild` from building by default for CI. Additional change to `scripts/build_pytorch_android.sh`: Skipping clean task here as android gradle plugin 3.3.2 exteralNativeBuild has problems with it when abiFilters are specified. Will be returned back in the following diffs with upgrading of gradle and android gradle plugin versions. Test Plan: Imported from OSS Differential Revision: D22118945 Pulled By: IvanKobzarev fbshipit-source-id: 31c54b49b1f262cbe5f540461d3406f74851db6c
diff --git a/android/README.md b/android/README.md
@@ -106,6 +106,116 @@ As `pytorch_android` [depends](https://github.com/pytorch/pytorch/blob/master/an
 
 You can check out [test app example](https://github.com/pytorch/pytorch/blob/master/android/test_app/app/build.gradle) that uses aars directly.
 
-## More Details
+## Linking to prebuilt libtorch library from gradle dependency
+
+In some cases, you may want to use libtorch from your android native build.
+You can do it without building libtorch android, using native libraries from PyTorch android gradle dependency.
+For that, you will need to add the next lines to your gradle build.
+```
+android {
+...
+    configurations {
+       extractForNativeBuild
+    }
+...
+    compileOptions {
+        externalNativeBuild {
+            cmake {
+                arguments "-DANDROID_STL=c++_shared"
+            }
+        }
+    }
+...
+    externalNativeBuild {
+        cmake {
+            path "CMakeLists.txt"
+        }
+    }
+}
+
+dependencies {
+    extractForNativeBuild('org.pytorch:pytorch_android:1.6.0')
+}
+
+task extractAARForNativeBuild {
+    doLast {
+        configurations.extractForNativeBuild.files.each {
+            def file = it.absoluteFile
+            copy {
+                from zipTree(file)
+                into "$buildDir/$file.name"
+                include "headers/**"
+                include "jni/**"
+            }
+        }
+    }
+}
+
+tasks.whenTaskAdded { task ->
+  if (task.name.contains('externalNativeBuild')) {
+    task.dependsOn(extractAARForNativeBuild)
+  }
+}
+```
+
+pytorch_android aar contains headers to link in `headers` folder and native libraries in `jni/$ANDROID_ABI/`.
+As PyTorch native libraries use `ANDROID_STL` - we should use `ANDROID_STL=c++_shared` to have only one loaded binary of STL.
+
+The added task will unpack them to gradle build directory.
+
+In your native build you can link to them adding these lines to your CMakeLists.txt:
+
+
+```
+# Relative path of gradle build directory to CMakeLists.txt
+set(build_DIR ${CMAKE_SOURCE_DIR}/build)
+
+file(GLOB PYTORCH_INCLUDE_DIRS "${build_DIR}/pytorch_android*.aar/headers")
+file(GLOB PYTORCH_LINK_DIRS "${build_DIR}/pytorch_android*.aar/jni/${ANDROID_ABI}")
+
+set(BUILD_SUBDIR ${ANDROID_ABI})
+target_include_directories(${PROJECT_NAME} PRIVATE
+  ${PYTORCH_INCLUDE_DIRS}
+)
+
+find_library(PYTORCH_LIBRARY pytorch_jni
+  PATHS ${PYTORCH_LINK_DIRS}
+  NO_CMAKE_FIND_ROOT_PATH)
+
+target_link_libraries(${PROJECT_NAME}
+  ${PYTORCH_LIBRARY})
+
+```
+If your CMakeLists.txt file is located in the same directory as your build.gradle, `set(build_DIR ${CMAKE_SOURCE_DIR}/build)` should work for you. But if you have another location of it, you may need to change it.
+
+After that, you can use libtorch C++ API from your native code.
+```
+#include <string>
+#include <ATen/NativeFunctions.h>
+#include <torch/script.h>
+namespace pytorch_testapp_jni {
+namespace {
+    struct JITCallGuard {
+      torch::autograd::AutoGradMode no_autograd_guard{false};
+      torch::AutoNonVariableTypeMode non_var_guard{true};
+      torch::jit::GraphOptimizerEnabledGuard no_optimizer_guard{false};
+    };
+}
+
+void loadAndForwardModel(const std::string& modelPath) {
+  JITCallGuard guard;
+  torch::jit::Module module = torch::jit::load(modelPath);
+  module.eval();
+  torch::Tensor t = torch::randn({1, 3, 224, 224});
+  c10::IValue t_out = module.forward({t});
+}
+}
+```
+
+To load torchscript model for mobile we need some special setup which is placed in `struct JITCallGuard` in this example. It may change in future, you can track the latest changes keeping an eye in our [pytorch android jni code]([https://github.com/pytorch/pytorch/blob/master/android/pytorch_android/src/main/cpp/pytorch_jni_jit.cpp#L28)
+
+[Example of linking to libtorch from aar](https://github.com/pytorch/pytorch/tree/master/android/test_app)
+
+## PyTorch Android API Javadoc
 
 You can find more details about the PyTorch Android API in the [Javadoc](https://pytorch.org/docs/stable/packages.html).
diff --git a/android/test_app/app/CMakeLists.txt b/android/test_app/app/CMakeLists.txt
@@ -0,0 +1,38 @@
+cmake_minimum_required(VERSION 3.4.1)
+set(PROJECT_NAME pytorch_testapp_jni)
+project(${PROJECT_NAME} CXX)
+set(CMAKE_CXX_STANDARD 14)
+set(CMAKE_VERBOSE_MAKEFILE ON)
+
+set(build_DIR ${CMAKE_SOURCE_DIR}/build)
+
+set(pytorch_testapp_cpp_DIR ${CMAKE_CURRENT_LIST_DIR}/src/main/cpp)
+message(STATUS "ANDROID_STL:${ANDROID_STL}")
+file(GLOB pytorch_testapp_SOURCES
+  ${pytorch_testapp_cpp_DIR}/pytorch_testapp_jni.cpp
+)
+
+add_library(${PROJECT_NAME} SHARED
+    ${pytorch_testapp_SOURCES}
+)
+
+file(GLOB PYTORCH_INCLUDE_DIRS "${build_DIR}/pytorch_android*.aar/headers")
+file(GLOB PYTORCH_LINK_DIRS "${build_DIR}/pytorch_android*.aar/jni/${ANDROID_ABI}")
+
+target_compile_options(${PROJECT_NAME} PRIVATE
+  -fexceptions
+)
+
+set(BUILD_SUBDIR ${ANDROID_ABI})
+
+target_include_directories(${PROJECT_NAME} PRIVATE
+  ${PYTORCH_INCLUDE_DIRS}
+)
+
+find_library(PYTORCH_LIBRARY pytorch_jni
+  PATHS ${PYTORCH_LINK_DIRS}
+  NO_CMAKE_FIND_ROOT_PATH)
+
+target_link_libraries(${PROJECT_NAME}
+  ${PYTORCH_LIBRARY}
+  log)
diff --git a/android/test_app/app/build.gradle b/android/test_app/app/build.gradle
@@ -11,6 +11,9 @@ repositories {
 }
 
 android {
+    configurations {
+        extractForNativeBuild
+    }
     compileOptions {
         sourceCompatibility 1.8
         targetCompatibility 1.8
@@ -28,17 +31,28 @@ android {
         }
         externalNativeBuild {
             cmake {
+                abiFilters ABI_FILTERS.split(",")
                 arguments "-DANDROID_STL=c++_shared"
             }
         }
         buildConfigField("String", "MODULE_ASSET_NAME", "\"mobilenet2q.pt\"")
         buildConfigField("String", "LOGCAT_TAG", "@string/app_name")
         buildConfigField("long[]", "INPUT_TENSOR_SHAPE", "new long[]{1, 3, 224, 224}")
+        buildConfigField("boolean", "NATIVE_BUILD", 'false')
         addManifestPlaceholders([APP_NAME: "@string/app_name", MAIN_ACTIVITY: "org.pytorch.testapp.MainActivity"])
     }
     buildTypes {
         debug {
             minifyEnabled false
+            debuggable true
+        }
+        release {
+            minifyEnabled false
+        }
+    }
+    externalNativeBuild {
+        cmake {
+            path "CMakeLists.txt"
         }
     }
     flavorDimensions "model", "build", "activity"
@@ -66,6 +80,10 @@ android {
         aar {
             dimension "build"
         }
+        nativeBuild {
+            dimension "build"
+            buildConfigField("boolean", "NATIVE_BUILD", "true")
+        }
         camera {
             dimension "activity"
             addManifestPlaceholders([MAIN_ACTIVITY: "org.pytorch.testapp.CameraActivity"])
@@ -79,7 +97,6 @@ android {
                     }
                 }
             }
-
         }
     }
     packagingOptions {
@@ -92,7 +109,8 @@ android {
             def names = variant.flavors*.name
             if (names.contains("nightly")
                 || names.contains("camera")
-                || names.contains("aar")) {
+                || names.contains("aar")
+                || names.contains("nativeBuild")) {
                 setIgnore(true)
             }
         }
@@ -101,9 +119,16 @@ android {
 
 dependencies {
     implementation 'com.android.support:appcompat-v7:28.0.0'
+    implementation 'com.facebook.soloader:nativeloader:0.8.0'
 
     localImplementation project(':pytorch_android')
     localImplementation project(':pytorch_android_torchvision')
+
+    nativeBuildImplementation(name: 'pytorch_android-release', ext: 'aar')
+    nativeBuildImplementation(name: 'pytorch_android_torchvision-release', ext: 'aar')
+    // Commented due to dependency on local copy of pytorch_android aar to aars folder
+    //extractForNativeBuild(name: 'pytorch_android-release', ext: 'aar')
+
     nightlyImplementation 'org.pytorch:pytorch_android:1.6.0-SNAPSHOT'
     nightlyImplementation 'org.pytorch:pytorch_android_torchvision:1.6.0-SNAPSHOT'
 
@@ -117,3 +142,23 @@ dependencies {
     cameraImplementation "androidx.camera:camera-camera2:$camerax_version"
     cameraImplementation 'com.google.android.material:material:1.0.0-beta01'
 }
+
+task extractAARForNativeBuild {
+    doLast {
+        configurations.extractForNativeBuild.files.each {
+            def file = it.absoluteFile
+            copy {
+                from zipTree(file)
+                into "$buildDir/$file.name"
+                include "headers/**"
+                include "jni/**"
+            }
+        }
+    }
+}
+
+tasks.whenTaskAdded { task ->
+  if (task.name.contains('externalNativeBuild')) {
+    task.dependsOn(extractAARForNativeBuild)
+  }
+}
diff --git a/android/test_app/app/src/main/cpp/pytorch_testapp_jni.cpp b/android/test_app/app/src/main/cpp/pytorch_testapp_jni.cpp
@@ -0,0 +1,78 @@
+#include <android/log.h>
+#include <pthread.h>
+#include <unistd.h>
+#include <cassert>
+#include <cmath>
+#include <vector>
+#define ALOGI(...) \
+  __android_log_print(ANDROID_LOG_INFO, "PyTorchTestAppJni", __VA_ARGS__)
+#define ALOGE(...) \
+  __android_log_print(ANDROID_LOG_ERROR, "PyTorchTestAppJni", __VA_ARGS__)
+
+#include "jni.h"
+
+#include <torch/script.h>
+
+namespace pytorch_testapp_jni {
+namespace {
+
+template <typename T>
+void log(const char* m, T t) {
+  std::ostringstream os;
+  os << t << std::endl;
+  ALOGI("%s %s", m, os.str().c_str());
+}
+
+struct JITCallGuard {
+  torch::autograd::AutoGradMode no_autograd_guard{false};
+  torch::AutoNonVariableTypeMode non_var_guard{true};
+  torch::jit::GraphOptimizerEnabledGuard no_optimizer_guard{false};
+};
+} // namespace
+
+static void loadAndForwardModel(JNIEnv* env, jclass, jstring jModelPath) {
+  const char* modelPath = env->GetStringUTFChars(jModelPath, 0);
+  assert(modelPath);
+
+  // To load torchscript model for mobile we need set these guards,
+  // because mobile build doesn't support features like autograd for smaller
+  // build size which is placed in `struct JITCallGuard` in this example. It may
+  // change in future, you can track the latest changes keeping an eye in
+  // android/pytorch_android/src/main/cpp/pytorch_jni_jit.cpp
+  JITCallGuard guard;
+  torch::jit::Module module = torch::jit::load(modelPath);
+  module.eval();
+  torch::Tensor t = torch::randn({1, 3, 224, 224});
+  log("input tensor:", t);
+  c10::IValue t_out = module.forward({t});
+  log("output tensor:", t_out);
+  env->ReleaseStringUTFChars(jModelPath, modelPath);
+}
+} // namespace pytorch_testapp_jni
+
+JNIEXPORT jint JNI_OnLoad(JavaVM* vm, void*) {
+  JNIEnv* env;
+  if (vm->GetEnv(reinterpret_cast<void**>(&env), JNI_VERSION_1_6) != JNI_OK) {
+    return JNI_ERR;
+  }
+
+  jclass c =
+      env->FindClass("org/pytorch/testapp/LibtorchNativeClient$NativePeer");
+  if (c == nullptr) {
+    return JNI_ERR;
+  }
+
+  static const JNINativeMethod methods[] = {
+      {"loadAndForwardModel",
+       "(Ljava/lang/String;)V",
+       (void*)pytorch_testapp_jni::loadAndForwardModel},
+  };
+  int rc = env->RegisterNatives(
+      c, methods, sizeof(methods) / sizeof(JNINativeMethod));
+
+  if (rc != JNI_OK) {
+    return rc;
+  }
+
+  return JNI_VERSION_1_6;
+}
diff --git a/android/test_app/app/src/main/java/org/pytorch/testapp/LibtorchNativeClient.java b/android/test_app/app/src/main/java/org/pytorch/testapp/LibtorchNativeClient.java
@@ -0,0 +1,21 @@
+package org.pytorch.testapp;
+import com.facebook.soloader.nativeloader.NativeLoader;
+import com.facebook.soloader.nativeloader.SystemDelegate;
+
+public final class LibtorchNativeClient {
+
+  public static void loadAndForwardModel(final String modelPath) {
+    NativePeer.loadAndForwardModel(modelPath); 
+  }
+
+  private static class NativePeer {
+    static {
+      if (!NativeLoader.isInitialized()) {
+        NativeLoader.init(new SystemDelegate());
+      }
+      NativeLoader.loadLibrary("pytorch_testapp_jni");
+    }
+
+    private static native void loadAndForwardModel(final String modelPath);
+  }
+}
diff --git a/android/test_app/app/src/main/java/org/pytorch/testapp/MainActivity.java b/android/test_app/app/src/main/java/org/pytorch/testapp/MainActivity.java
diff --git a/scripts/build_pytorch_android.sh b/scripts/build_pytorch_android.sh
