Update fbjni and enable PyTorch JNI build

dreiss · facebook-github-bot · commit d22f61432dc7 · 2019-11-15T13:59:44.000-08:00
Summary:
- Add a "BUILD_JNI" option that enables building PyTorch JNI bindings and
  fbjni.  This is off by default because it adds a dependency on jni.h.
- Update to the latest fbjni so we can inhibit building its tests,
  because they depend on gtest.
- Set JAVA_HOME and BUILD_JNI in Linux binary build configurations if we
  can find jni.h in Docker.

Test Plan:
- Built on dev server.
- Verified that libpytorch_jni links after libtorch when both are built
  in a parallel build.

Differential Revision: D18536828

fbshipit-source-id: 19cb3be8298d3619352d02bb9446ab802c27ec66
diff --git a/.circleci/scripts/binary_populate_env.sh b/.circleci/scripts/binary_populate_env.sh
@@ -62,6 +62,28 @@ else
 fi
 export PYTORCH_BUILD_NUMBER=1
 
+
+JAVA_HOME=
+BUILD_JNI=OFF
+if [[ "$PACKAGE_TYPE" == libtorch ]]; then
+  POSSIBLE_JAVA_HOMES=()
+  POSSIBLE_JAVA_HOMES+=(/usr/local)
+  POSSIBLE_JAVA_HOMES+=(/usr/lib/jvm/java-8-openjdk-amd64)
+  # TODO: Fix Mac Java build
+  #POSSIBLE_JAVA_HOMES+=(/Library/Java/JavaVirtualMachines/*.jdk/Contents/Home)
+  for JH in "${POSSIBLE_JAVA_HOMES[@]}" ; do
+    if [[ -e "$JH/include/jni.h" ]] ; then
+      echo "Found jni.h under $JH"
+      JAVA_HOME="$JH"
+      BUILD_JNI=ON
+      break
+    fi
+  done
+  if [ -z "$JAVA_HOME" ]; then
+    echo "Did not find jni.h"
+  fi
+fi
+
 cat >>"$envfile" <<EOL
 # =================== The following code will be executed inside Docker container ===================
 export TZ=UTC
@@ -85,6 +107,8 @@ export TORCH_PACKAGE_NAME='torch'
 export TORCH_CONDA_BUILD_FOLDER='pytorch-nightly'
 
 export USE_FBGEMM=1
+export JAVA_HOME=$JAVA_HOME
+export BUILD_JNI=$BUILD_JNI
 export PIP_UPLOAD_FOLDER="$PIP_UPLOAD_FOLDER"
 export DOCKER_IMAGE="$DOCKER_IMAGE"
 
diff --git a/.jenkins/pytorch/build.sh b/.jenkins/pytorch/build.sh
@@ -61,6 +61,25 @@ if ! which conda; then
   fi
 fi
 
+if [[ "$BUILD_ENVIRONMENT" == *libtorch* ]]; then
+  POSSIBLE_JAVA_HOMES=()
+  POSSIBLE_JAVA_HOMES+=(/usr/local)
+  POSSIBLE_JAVA_HOMES+=(/usr/lib/jvm/java-8-openjdk-amd64)
+  # TODO: Fix Mac Java build
+  #POSSIBLE_JAVA_HOMES+=(/Library/Java/JavaVirtualMachines/*.jdk/Contents/Home)
+  for JH in "${POSSIBLE_JAVA_HOMES[@]}" ; do
+    if [[ -e "$JH/include/jni.h" ]] ; then
+      echo "Found jni.h under $JH"
+      export JAVA_HOME="$JH"
+      export BUILD_JNI=ON
+      break
+    fi
+  done
+  if [ -z "$JAVA_HOME" ]; then
+    echo "Did not find jni.h"
+  fi
+fi
+
 # Use special scripts for Android builds
 if [[ "${BUILD_ENVIRONMENT}" == *-android* ]]; then
   export ANDROID_NDK=/opt/ndk
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -120,6 +120,7 @@ cmake_dependent_option(
     CAFFE2_USE_MSVC_STATIC_RUNTIME "Using MSVC static runtime libraries" ON
     "NOT BUILD_SHARED_LIBS" OFF)
 option(BUILD_TEST "Build C++ test binaries (need gtest and gbenchmark)" OFF)
+option(BUILD_JNI "Build JNI bindings" OFF)
 cmake_dependent_option(
     INSTALL_TEST "Install test binaries if BUILD_TEST is on" ON
     "BUILD_TEST" OFF)
@@ -646,5 +647,12 @@ if (BUILD_BINARY)
   add_subdirectory(binaries)
 endif()
 
+# ---[ JNI
+if (BUILD_JNI)
+  set(BUILD_LIBTORCH_WITH_JNI 1)
+  set(FBJNI_SKIP_TESTS 1)
+  add_subdirectory(android/pytorch_android)
+endif()
+
 include(cmake/Summary.cmake)
 caffe2_print_configuration_summary()
diff --git a/android/libs/fbjni b/android/libs/fbjni
@@ -1 +1 @@
-Subproject commit e109650ca03e3db54ab973c14d10f176b6436c48
+Subproject commit f908b58be482874137fa4c0e71333e4eca481706
diff --git a/android/pytorch_android/CMakeLists.txt b/android/pytorch_android/CMakeLists.txt
@@ -19,7 +19,10 @@ set(pytorch_android_DIR ${CMAKE_CURRENT_LIST_DIR}/src/main/cpp)
 if (ANDROID_ABI)
   set(libtorch_include_DIR ${pytorch_android_DIR}/libtorch_include/${ANDROID_ABI})
   set(BUILD_SUBDIR ${ANDROID_ABI})
+elseif(BUILD_LIBTORCH_WITH_JNI)
+  # Don't need LIBTORCH_HOME if we're building from within PyTorch.
 else()
+  # Building against a pre-built libtorch.
   if (NOT LIBTORCH_HOME)
     message(FATAL_ERROR
       "pytorch_android requires LIBTORCH_HOME to be defined for non-Android builds.")
diff --git a/cmake/Summary.cmake b/cmake/Summary.cmake
@@ -44,6 +44,7 @@ function (caffe2_print_configuration_summary)
   message(STATUS "  BUILD_CAFFE2_OPS      : ${BUILD_CAFFE2_OPS}")
   message(STATUS "  BUILD_SHARED_LIBS     : ${BUILD_SHARED_LIBS}")
   message(STATUS "  BUILD_TEST            : ${BUILD_TEST}")
+  message(STATUS "  BUILD_JNI             : ${BUILD_JNI}")
 
   message(STATUS "  INTERN_BUILD_MOBILE   : ${INTERN_BUILD_MOBILE}")
 
diff --git a/tools/setup_helpers/cmake.py b/tools/setup_helpers/cmake.py
@@ -230,6 +230,7 @@ def generate(self, version, cmake_python_library, build_python, build_test, my_e
              'CUDNN_INCLUDE_DIR',
              'EXPERIMENTAL_SINGLE_THREAD_POOL',
              'INSTALL_TEST',
+             'JAVA_HOME',
              'MKL_THREADING',
              'MKLDNN_THREADING',
              'MSVC_Z7_OVERRIDE',
