[vulkan] jni build support USE_VULKAN (#39188)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/39188

Extracting Vulkan_LIBS and Vulkan_INCLUDES setup from `cmake/Dependencies.cmake` to `cmake/VulkanDependencies.cmake` and reuse it in android/pytorch_android/CMakeLists.txt

Adding control to build with Vulkan setting env variable `USE_VULKAN` for `scripts/build_android.sh` `scripts/build_pytorch_android.sh`

We do not use Vulkan backend in pytorch_android, but with this build option we can track android aar change with `USE_VULKAN` added.

Currently it is 88Kb.

Test Plan: Imported from OSS

Differential Revision: D21770892

Pulled By: IvanKobzarev

fbshipit-source-id: a39433505fdcf43d3b524e0fe08062d5ebe0d872

android/pytorch_android/CMakeLists.txt

@@ -13,6 +13,13 @@ endif()
 if(NOT TRACE_ENABLED)
   message(STATUS "TRACE_ENABLED OFF")
 endif()
+set(USE_VULKAN OFF)
+if(DEFINED ENV{USE_VULKAN})
+  if($ENV{USE_VULKAN} STREQUAL "1")
+    message(STATUS "USE_VULKAN ON")
+    set(USE_VULKAN ON)
+  endif()
+endif()
 
 set(pytorch_android_DIR ${CMAKE_CURRENT_LIST_DIR}/src/main/cpp)
 
@@ -69,8 +76,14 @@ set(fbjni_BUILD_DIR ${CMAKE_BINARY_DIR}/fbjni/${BUILD_SUBDIR})
 
 add_subdirectory(${fbjni_DIR} ${fbjni_BUILD_DIR})
 
-if(ANDROID_ABI)
+# ---[ Vulkan deps
+if(USE_VULKAN)
+  set(Vulkan_LIBS)
+  set(Vulkan_INCLUDES)
+  include(${CMAKE_CURRENT_LIST_DIR}/../../cmake/VulkanDependencies.cmake)
+endif()
 
+if(ANDROID_ABI)
   function(import_static_lib name)
     add_library(${name} STATIC IMPORTED)
     set_property(
@@ -89,8 +102,8 @@ if(ANDROID_ABI)
   import_static_lib(libcpuinfo)
   import_static_lib(libclog)
 
-    # Link most things statically on Android.
+  # Link most things statically on Android.
-  target_link_libraries(pytorch_jni
+  set(pytorch_jni_LIBS
       fbjni
       -Wl,--gc-sections
       -Wl,--whole-archive
@@ -105,11 +118,9 @@ if(ANDROID_ABI)
       libcpuinfo
       libclog
   )
-
 else()
-
   # Prefer dynamic linking on the host
-  target_link_libraries(pytorch_jni
+  set(pytorch_jni_LIBS
       fbjni
       torch
       torch_cpu
@@ -120,5 +131,10 @@ else()
       cpuinfo
       clog
   )
-
 endif()
+
+if(USE_VULKAN)
+  list(APPEND pytorch_jni_LIBS ${Vulkan_LIBS})
+endif()
+
+target_link_libraries(pytorch_jni ${pytorch_jni_LIBS})

cmake/Dependencies.cmake

@@ -476,161 +476,11 @@ endif()
 
 # ---[ Vulkan deps
 if(USE_VULKAN)
-  if(ANDROID)
+  set(Vulkan_LIBS)
-    if(NOT ANDROID_NDK)
+  set(Vulkan_INCLUDES)
-      message(FATAL_ERROR "USE_VULKAN requires ANDROID_NDK set")
+  include(${CMAKE_CURRENT_LIST_DIR}/VulkanDependencies.cmake)
-    endif()
+  list(APPEND Caffe2_DEPENDENCY_LIBS ${Vulkan_LIBS})
-
+  include_directories(SYSTEM ${Vulkan_INCLUDES})
-    # Vulkan from ANDROID_NDK
-    set(VULKAN_INCLUDE_DIR "${ANDROID_NDK}/sources/third-party/vulkan/src/include")
-    message(STATUS "VULKAN_INCLUDE_DIR:${VULKAN_INCLUDE_DIR}")
-
-    set(VULKAN_ANDROID_NDK_WRAPPER_DIR "${ANDROID_NDK}/sources/third_party/vulkan/src/common")
-    message(STATUS "Vulkan_ANDROID_NDK_WRAPPER_DIR:${VULKAN_ANDROID_NDK_WRAPPER_DIR}")
-    set(VULKAN_WRAPPER_DIR "${VULKAN_ANDROID_NDK_WRAPPER_DIR}")
-
-    add_library(
-      VulkanWrapper
-      STATIC
-      ${VULKAN_WRAPPER_DIR}/vulkan_wrapper.h
-      ${VULKAN_WRAPPER_DIR}/vulkan_wrapper.cpp)
-
-    target_include_directories(VulkanWrapper PUBLIC .)
-    target_include_directories(VulkanWrapper PUBLIC "${VULKAN_INCLUDE_DIR}")
-
-    target_link_libraries(VulkanWrapper ${CMAKE_DL_LIBS})
-
-    include_directories(SYSTEM ${VULKAN_WRAPPER_DIR})
-    list(APPEND Caffe2_DEPENDENCY_LIBS VulkanWrapper)
-
-    # Shaderc
-    if(USE_VULKAN_SHADERC_RUNTIME)
-      # Shaderc from ANDROID_NDK
-      set(Shaderc_ANDROID_NDK_INCLUDE_DIR "${ANDROID_NDK}/sources/third_party/shaderc/include")
-      message(STATUS "Shaderc_ANDROID_NDK_INCLUDE_DIR:${Shaderc_ANDROID_NDK_INCLUDE_DIR}")
-
-      find_path(
-        GOOGLE_SHADERC_INCLUDE_DIRS
-        NAMES shaderc/shaderc.hpp
-        PATHS "${Shaderc_ANDROID_NDK_INCLUDE_DIR}")
-
-      set(Shaderc_ANDROID_NDK_LIB_DIR "${ANDROID_NDK}/sources/third_party/shaderc/libs/${ANDROID_STL}/${ANDROID_ABI}")
-      message(STATUS "Shaderc_ANDROID_NDK_LIB_DIR:${Shaderc_ANDROID_NDK_LIB_DIR}")
-
-      find_library(
-        GOOGLE_SHADERC_LIBRARIES
-        NAMES shaderc
-        PATHS "${Shaderc_ANDROID_NDK_LIB_DIR}")
-
-      # Shaderc in NDK is not prebuilt
-      if(NOT GOOGLE_SHADERC_LIBRARIES)
-        set(NDK_SHADERC_DIR "${ANDROID_NDK}/sources/third_party/shaderc")
-        set(NDK_BUILD_CMD "${ANDROID_NDK}/ndk-build")
-
-        execute_process(
-          COMMAND ${NDK_BUILD_CMD}
-          NDK_PROJECT_PATH=${NDK_SHADERC_DIR}
-          APP_BUILD_SCRIPT=${NDK_SHADERC_DIR}/Android.mk
-          APP_PLATFORM=${ANDROID_PLATFORM}
-          APP_STL=${ANDROID_STL}
-          APP_ABI=${ANDROID_ABI}
-          libshaderc_combined -j8
-          WORKING_DIRECTORY "${NDK_SHADERC_DIR}"
-          RESULT_VARIABLE error_code)
-
-        if(error_code)
-          message(FATAL_ERROR "Failed to build ANDROID_NDK shaderc error_code:${error_code}")
-        else()
-          unset(GOOGLE_SHADERC_LIBRARIES CACHE)
-          find_library(
-            GOOGLE_SHADERC_LIBRARIES
-            NAMES shaderc
-            HINTS "${Shaderc_ANDROID_NDK_LIB_DIR}")
-        endif()
-      endif(NOT GOOGLE_SHADERC_LIBRARIES)
-
-      if(GOOGLE_SHADERC_INCLUDE_DIRS AND GOOGLE_SHADERC_LIBRARIES)
-        message(STATUS "shaderc FOUND include:${GOOGLE_SHADERC_INCLUDE_DIRS}")
-        message(STATUS "shaderc FOUND libs:${GOOGLE_SHADERC_LIBRARIES}")
-        set(Shaderc_FOUND TRUE)
-      endif()
-
-      include_directories(SYSTEM ${GOOGLE_SHADERC_INCLUDE_DIRS})
-      list(APPEND Caffe2_DEPENDENCY_LIBS ${GOOGLE_SHADERC_LIBRARIES})
-    endif(USE_VULKAN_SHADERC_RUNTIME)
-  else()
-    # USE_VULKAN AND NOT ANDROID
-    if(NOT DEFINED ENV{VULKAN_SDK})
-      message(FATAL_ERROR "USE_VULKAN requires environment var VULKAN_SDK set")
-    endif()
-    message(STATUS "VULKAN_SDK:$ENV{VULKAN_SDK}")
-
-    set(VULKAN_INCLUDE_DIR "$ENV{VULKAN_SDK}/source/Vulkan-Headers/include")
-    message(STATUS "VULKAN_INCLUDE_DIR:${VULKAN_INCLUDE_DIR}")
-
-    if(USE_VULKAN_WRAPPER)
-      # Vulkan wrapper from VULKAN_SDK
-      set(VULKAN_SDK_WRAPPER_DIR "$ENV{VULKAN_SDK}/source/Vulkan-Tools/common")
-      message(STATUS "Vulkan_SDK_WRAPPER_DIR:${VULKAN_SDK_WRAPPER_DIR}")
-      set(VULKAN_WRAPPER_DIR "${VULKAN_SDK_WRAPPER_DIR}")
-
-      add_library(
-        VulkanWrapper
-        STATIC
-        ${VULKAN_WRAPPER_DIR}/vulkan_wrapper.h
-        ${VULKAN_WRAPPER_DIR}/vulkan_wrapper.cpp)
-
-      target_include_directories(VulkanWrapper PUBLIC .)
-      target_include_directories(VulkanWrapper PUBLIC "${VULKAN_INCLUDE_DIR}")
-
-      target_link_libraries(VulkanWrapper ${CMAKE_DL_LIBS})
-
-      include_directories(SYSTEM ${VULKAN_WRAPPER_DIR})
-      list(APPEND Caffe2_DEPENDENCY_LIBS VulkanWrapper)
-    else(USE_VULKAN_WRAPPER)
-      find_library(VULKAN_LIBRARY
-        NAMES vulkan
-        PATHS
-        "$ENV{VULKAN_SDK}/${CMAKE_HOST_SYSTEM_PROCESSOR}/lib")
-
-      if(NOT VULKAN_LIBRARY)
-        message(FATAL_ERROR "USE_VULKAN: Vulkan library not found")
-      endif()
-
-      message(STATUS "VULKAN_LIBRARY:${VULKAN_LIBRARY}")
-      message(STATUS "VULKAN_INCLUDE_DIR:${VULKAN_INCLUDE_DIR}")
-
-      include_directories(SYSTEM ${VULKAN_INCLUDE_DIR})
-      list(APPEND Caffe2_DEPENDENCY_LIBS ${VULKAN_LIBRARY})
-    endif(USE_VULKAN_WRAPPER)
-
-    if(USE_VULKAN_SHADERC_RUNTIME)
-      # shaderc from VULKAN_SDK
-      find_path(
-          GOOGLE_SHADERC_INCLUDE_DIRS
-          NAMES shaderc/shaderc.hpp
-          PATHS $ENV{VULKAN_SDK}/${CMAKE_HOST_SYSTEM_PROCESSOR}/include)
-
-      find_library(
-          GOOGLE_SHADERC_LIBRARIES
-          NAMES shaderc_combined
-          PATHS $ENV{VULKAN_SDK}/${CMAKE_HOST_SYSTEM_PROCESSOR}/lib)
-
-      find_package_handle_standard_args(
-          Shaderc
-          DEFAULT_MSG
-          GOOGLE_SHADERC_INCLUDE_DIRS
-          GOOGLE_SHADERC_LIBRARIES)
-      if(NOT Shaderc_FOUND)
-        message(FATAL_ERROR "USE_VULKAN: Shaderc not found in VULKAN_SDK")
-      else()
-        message(STATUS "shaderc FOUND include:${GOOGLE_SHADERC_INCLUDE_DIRS}")
-        message(STATUS "shaderc FOUND libs:${GOOGLE_SHADERC_LIBRARIES}")
-      endif()
-      include_directories(SYSTEM ${GOOGLE_SHADERC_INCLUDE_DIRS})
-      list(APPEND Caffe2_DEPENDENCY_LIBS ${GOOGLE_SHADERC_LIBRARIES})
-    endif(USE_VULKAN_SHADERC_RUNTIME)
-  endif()
 endif()
 
 # ---[ gflags

cmake/VulkanDependencies.cmake

@@ -0,0 +1,160 @@
+if(NOT USE_VULKAN)
+  return()
+endif()
+
+if(ANDROID)
+  if(NOT ANDROID_NDK)
+    message(FATAL_ERROR "USE_VULKAN requires ANDROID_NDK set")
+  endif()
+
+  # Vulkan from ANDROID_NDK
+  set(VULKAN_INCLUDE_DIR "${ANDROID_NDK}/sources/third-party/vulkan/src/include")
+  message(STATUS "VULKAN_INCLUDE_DIR:${VULKAN_INCLUDE_DIR}")
+
+  set(VULKAN_ANDROID_NDK_WRAPPER_DIR "${ANDROID_NDK}/sources/third_party/vulkan/src/common")
+  message(STATUS "Vulkan_ANDROID_NDK_WRAPPER_DIR:${VULKAN_ANDROID_NDK_WRAPPER_DIR}")
+  set(VULKAN_WRAPPER_DIR "${VULKAN_ANDROID_NDK_WRAPPER_DIR}")
+
+  add_library(
+    VulkanWrapper
+    STATIC
+    ${VULKAN_WRAPPER_DIR}/vulkan_wrapper.h
+    ${VULKAN_WRAPPER_DIR}/vulkan_wrapper.cpp)
+
+  target_include_directories(VulkanWrapper PUBLIC .)
+  target_include_directories(VulkanWrapper PUBLIC "${VULKAN_INCLUDE_DIR}")
+
+  target_link_libraries(VulkanWrapper ${CMAKE_DL_LIBS})
+
+  list(APPEND Vulkan_INCLUDES ${VULKAN_WRAPPER_DIR})
+  list(APPEND Vulkan_LIBS VulkanWrapper)
+
+  # Shaderc
+  if(USE_VULKAN_SHADERC_RUNTIME)
+    # Shaderc from ANDROID_NDK
+    set(Shaderc_ANDROID_NDK_INCLUDE_DIR "${ANDROID_NDK}/sources/third_party/shaderc/include")
+    message(STATUS "Shaderc_ANDROID_NDK_INCLUDE_DIR:${Shaderc_ANDROID_NDK_INCLUDE_DIR}")
+
+    find_path(
+      GOOGLE_SHADERC_INCLUDE_DIRS
+      NAMES shaderc/shaderc.hpp
+      PATHS "${Shaderc_ANDROID_NDK_INCLUDE_DIR}")
+
+    set(Shaderc_ANDROID_NDK_LIB_DIR "${ANDROID_NDK}/sources/third_party/shaderc/libs/${ANDROID_STL}/${ANDROID_ABI}")
+    message(STATUS "Shaderc_ANDROID_NDK_LIB_DIR:${Shaderc_ANDROID_NDK_LIB_DIR}")
+
+    find_library(
+      GOOGLE_SHADERC_LIBRARIES
+      NAMES shaderc
+      PATHS "${Shaderc_ANDROID_NDK_LIB_DIR}")
+
+    # Shaderc in NDK is not prebuilt
+    if(NOT GOOGLE_SHADERC_LIBRARIES)
+      set(NDK_SHADERC_DIR "${ANDROID_NDK}/sources/third_party/shaderc")
+      set(NDK_BUILD_CMD "${ANDROID_NDK}/ndk-build")
+
+      execute_process(
+        COMMAND ${NDK_BUILD_CMD}
+        NDK_PROJECT_PATH=${NDK_SHADERC_DIR}
+        APP_BUILD_SCRIPT=${NDK_SHADERC_DIR}/Android.mk
+        APP_PLATFORM=${ANDROID_PLATFORM}
+        APP_STL=${ANDROID_STL}
+        APP_ABI=${ANDROID_ABI}
+        libshaderc_combined -j8
+        WORKING_DIRECTORY "${NDK_SHADERC_DIR}"
+        RESULT_VARIABLE error_code)
+
+      if(error_code)
+        message(FATAL_ERROR "Failed to build ANDROID_NDK shaderc error_code:${error_code}")
+      else()
+        unset(GOOGLE_SHADERC_LIBRARIES CACHE)
+        find_library(
+          GOOGLE_SHADERC_LIBRARIES
+          NAMES shaderc
+          HINTS "${Shaderc_ANDROID_NDK_LIB_DIR}")
+      endif()
+    endif(NOT GOOGLE_SHADERC_LIBRARIES)
+
+    if(GOOGLE_SHADERC_INCLUDE_DIRS AND GOOGLE_SHADERC_LIBRARIES)
+      message(STATUS "shaderc FOUND include:${GOOGLE_SHADERC_INCLUDE_DIRS}")
+      message(STATUS "shaderc FOUND libs:${GOOGLE_SHADERC_LIBRARIES}")
+      set(Shaderc_FOUND TRUE)
+    endif()
+
+    list(APPEND Vulkan_INCLUDES ${GOOGLE_SHADERC_INCLUDE_DIRS})
+    list(APPEND Vulkan_LIBS ${GOOGLE_SHADERC_LIBRARIES})
+  endif(USE_VULKAN_SHADERC_RUNTIME)
+else()
+  # USE_VULKAN AND NOT ANDROID
+  if(NOT DEFINED ENV{VULKAN_SDK})
+    message(FATAL_ERROR "USE_VULKAN requires environment var VULKAN_SDK set")
+  endif()
+  message(STATUS "VULKAN_SDK:$ENV{VULKAN_SDK}")
+
+  set(VULKAN_INCLUDE_DIR "$ENV{VULKAN_SDK}/source/Vulkan-Headers/include")
+  message(STATUS "VULKAN_INCLUDE_DIR:${VULKAN_INCLUDE_DIR}")
+
+  if(USE_VULKAN_WRAPPER)
+    # Vulkan wrapper from VULKAN_SDK
+    set(VULKAN_SDK_WRAPPER_DIR "$ENV{VULKAN_SDK}/source/Vulkan-Tools/common")
+    message(STATUS "Vulkan_SDK_WRAPPER_DIR:${VULKAN_SDK_WRAPPER_DIR}")
+    set(VULKAN_WRAPPER_DIR "${VULKAN_SDK_WRAPPER_DIR}")
+
+    add_library(
+      VulkanWrapper
+      STATIC
+      ${VULKAN_WRAPPER_DIR}/vulkan_wrapper.h
+      ${VULKAN_WRAPPER_DIR}/vulkan_wrapper.cpp)
+
+    target_include_directories(VulkanWrapper PUBLIC .)
+    target_include_directories(VulkanWrapper PUBLIC "${VULKAN_INCLUDE_DIR}")
+
+    target_link_libraries(VulkanWrapper ${CMAKE_DL_LIBS})
+
+    list(APPEND Vulkan_INCLUDES ${VULKAN_WRAPPER_DIR})
+    list(APPEND Vulkan_LIBS VulkanWrapper)
+  else(USE_VULKAN_WRAPPER)
+    find_library(VULKAN_LIBRARY
+      NAMES vulkan
+      PATHS
+      "$ENV{VULKAN_SDK}/${CMAKE_HOST_SYSTEM_PROCESSOR}/lib")
+
+    if(NOT VULKAN_LIBRARY)
+      message(FATAL_ERROR "USE_VULKAN: Vulkan library not found")
+    endif()
+
+    message(STATUS "VULKAN_LIBRARY:${VULKAN_LIBRARY}")
+    message(STATUS "VULKAN_INCLUDE_DIR:${VULKAN_INCLUDE_DIR}")
+
+    list(APPEND Vulkan_INCLUDES ${VULKAN_INCLUDE_DIR})
+    list(APPEND Vulkan_LIBS ${VULKAN_LIBRARY})
+  endif(USE_VULKAN_WRAPPER)
+
+  if(USE_VULKAN_SHADERC_RUNTIME)
+    # shaderc from VULKAN_SDK
+    find_path(
+        GOOGLE_SHADERC_INCLUDE_DIRS
+        NAMES shaderc/shaderc.hpp
+        PATHS $ENV{VULKAN_SDK}/${CMAKE_HOST_SYSTEM_PROCESSOR}/include)
+
+    find_library(
+        GOOGLE_SHADERC_LIBRARIES
+        NAMES shaderc_combined
+        PATHS $ENV{VULKAN_SDK}/${CMAKE_HOST_SYSTEM_PROCESSOR}/lib)
+
+    find_package_handle_standard_args(
+        Shaderc
+        DEFAULT_MSG
+        GOOGLE_SHADERC_INCLUDE_DIRS
+        GOOGLE_SHADERC_LIBRARIES)
+    if(NOT Shaderc_FOUND)
+      message(FATAL_ERROR "USE_VULKAN: Shaderc not found in VULKAN_SDK")
+    else()
+      message(STATUS "shaderc FOUND include:${GOOGLE_SHADERC_INCLUDE_DIRS}")
+      message(STATUS "shaderc FOUND libs:${GOOGLE_SHADERC_LIBRARIES}")
+    endif()
+    list(APPEND Vulkan_INCLUDES ${GOOGLE_SHADERC_INCLUDE_DIRS})
+    list(APPEND Vulkan_LIBS ${GOOGLE_SHADERC_LIBRARIES})
+  endif(USE_VULKAN_SHADERC_RUNTIME)
+endif()
+

scripts/build_android.sh

@@ -117,6 +117,10 @@ if [ "${ANDROID_DEBUG_SYMBOLS:-}" == '1' ]; then
   CMAKE_ARGS+=("-DANDROID_DEBUG_SYMBOLS=1")
 fi
 
+if [ -n "${USE_VULKAN}" ]; then
+  CMAKE_ARGS+=("-DUSE_VULKAN=ON")
+fi
+
 # Use-specified CMake arguments go last to allow overridding defaults
 CMAKE_ARGS+=($@)