[pytorch] Upgrade to pytorch 1.11.0 (#1583)

* [pytorch] Upgrade PyTorch to 1.11.0 Change-Id: I9ee7d61a1c3fa68df50da4cf9dc8f54f125b6b43 * Building Pytorch for precxx11 aarch64 using v1.11.0 * Fix build.sh typo * Adding support for precxx11/cxx11 aarch64 builds * Adding support for precxx11/cxx11 aarch64 build in github action * [pytorch] Upgrade PyTorch to 1.11.0 Change-Id: I9ee7d61a1c3fa68df50da4cf9dc8f54f125b6b43 Co-authored-by: Harris Dizdarevic <dizdah@amazon.com>
deepjavalibrary · Apr 19, 2022 · 06997ce · 06997ce
1 parent e65ba99
commit 06997ce
Show file tree

Hide file tree

Showing 13 changed files with 609 additions and 42 deletions.
diff --git a/.github/workflows/native_jni_s3_pytorch.yml b/.github/workflows/native_jni_s3_pytorch.yml
@@ -66,13 +66,14 @@ jobs:
       - name: Install Environment
         run: |
           apt-get update
-          DEBIAN_FRONTEND=noninteractive apt-get install -y locales cmake curl unzip software-properties-common gcc-aarch64-linux-gnu g++-aarch64-linux-gnu
+          DEBIAN_FRONTEND=noninteractive apt-get install -y locales curl unzip software-properties-common
           add-apt-repository -y ppa:deadsnakes/ppa
           apt-get update
           apt-get install -y python3 python3-distutils
           curl -O https://bootstrap.pypa.io/pip/3.6/get-pip.py
           python3 get-pip.py
-          pip3 install awscli --upgrade
+          pip3 install awscli cmake
+          ln -s /usr/local/bin/cmake /usr/bin/cmake
       - name: Release JNI prep
         run: |
           PYTORCH_VERSION=${{ github.event.inputs.pt_version }}
@@ -84,8 +85,6 @@ jobs:
           ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcu10 -Ppt_version=$PYTORCH_VERSION
           ./gradlew :engines:pytorch:pytorch-native:cleanJNI
           ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcu11 -Ppt_version=$PYTORCH_VERSION
-          ./gradlew :engines:pytorch:pytorch-native:cleanJNI
-          CXX=aarch64-linux-gnu-gcc ./gradlew :engines:pytorch:pytorch-native:compileJNI -Paarch64 -Ppt_version=$PYTORCH_VERSION
       - name: Configure AWS Credentials
         uses: aws-actions/configure-aws-credentials@v1
         with:
@@ -161,33 +160,27 @@ jobs:
           key: ${{ runner.os }}-gradle-${{ hashFiles('**/*.gradle*') }}
           restore-keys: |
             ${{ runner.os }}-gradle-
-      - name: Install CUDA 10.2
+      - name: Install CUDA 11.3
         shell: cmd
         run: |
-          curl.exe -L http://developer.download.nvidia.com/compute/cuda/10.2/Prod/local_installers/cuda_10.2.89_441.22_windows.exe -o cuda102.exe
-          curl.exe -L https://developer.download.nvidia.com/compute/redist/cudnn/v7.6.4/cudnn-10.1-windows7-x64-v7.6.4.38.zip -o cudnn.zip
-          cuda102.exe -s
+          call "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvarsall.bat" amd64
+          curl.exe -L https://developer.download.nvidia.com/compute/cuda/11.3.1/local_installers/cuda_11.3.1_465.89_win10.exe -o cuda.exe
+          curl.exe -L https://developer.download.nvidia.com/compute/redist/cudnn/v8.2.1/cudnn-11.3-windows-x64-v8.2.1.32.zip -o cudnn.zip
+          cuda.exe -s
           mkdir cuda
           unzip.exe cudnn.zip
-          cp.exe -a cuda/include cuda/lib cuda/bin "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v10.2/"
-          rm.exe -Rf cuda102.exe cuda.exe cudnn.zip cuda
+          cp.exe -a cuda/include cuda/lib cuda/bin "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v11.3/"
+          rm.exe -Rf cuda.exe cuda.exe cudnn.zip cuda
       - name: Release CPU JNI
         shell: cmd
         run: |
-          call "C:\Program Files (x86)\Microsoft Visual Studio\2019\Enterprise\VC\Auxiliary\Build\vcvarsall.bat" amd64
+          call "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvarsall.bat" amd64
           gradlew :engines:pytorch:pytorch-native:compileJNI -Ppt_version=${{ github.event.inputs.pt_version }}
-      - name: Release cuda10 JNI
-        shell: cmd
-        run: |
-          call "C:\Program Files (x86)\Microsoft Visual Studio\2019\Enterprise\VC\Auxiliary\Build\vcvarsall.bat" amd64
-          set "CUDA_PATH=%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v10.2"
-          set "PATH=%CUDA_PATH%\bin;%CUDA_PATH%\libnvvp;%PATH%"
-          gradlew :engines:pytorch:pytorch-native:cleanJNI :engines:pytorch:pytorch-native:compileJNI -Pcu10 -Ppt_version=${{ github.event.inputs.pt_version }}
       - name: Release cuda11 JNI
         shell: cmd
         run: |
-          call "C:\Program Files (x86)\Microsoft Visual Studio\2019\Enterprise\VC\Auxiliary\Build\vcvarsall.bat" amd64
-          set "CUDA_PATH=%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v10.2"
+          call "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvarsall.bat" amd64
+          set "CUDA_PATH=%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v11.3"
           set "PATH=%CUDA_PATH%\bin;%CUDA_PATH%\libnvvp;%PATH%"
           gradlew :engines:pytorch:pytorch-native:cleanJNI :engines:pytorch:pytorch-native:compileJNI -Pcu11 -Ppt_version=${{ github.event.inputs.pt_version }}
       - name: Configure AWS Credentials

diff --git a/.github/workflows/nightly_publish.yml b/.github/workflows/nightly_publish.yml
@@ -108,8 +108,8 @@ jobs:
       - name: Publish to snapshot repository
         if: ${{ github.event.inputs.mode == '' || github.event.inputs.mode == 'snapshot' }}
         run: |
-          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.8.1 -Psnapshot
           ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.9.1 -Psnapshot
+          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.10.0 -Psnapshot
           ./gradlew clean publish -Psnapshot
           cd bom
           ./gradlew publish -Psnapshot
@@ -121,8 +121,8 @@ jobs:
       - name: Publish to staging repository
         if: ${{ github.event.inputs.mode == 'staging' }}
         run: |
-          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.8.1 -P${{ github.event.inputs.mode }}
           ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.9.1 -P${{ github.event.inputs.mode }}
+          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.10.0 -P${{ github.event.inputs.mode }}
           ./gradlew clean publish -P${{ github.event.inputs.mode }}
           cd bom
           ./gradlew publish -P${{ github.event.inputs.mode }}

diff --git a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/MkldnnTest.java b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/MkldnnTest.java
@@ -17,13 +17,18 @@
 import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.Shape;
 import java.util.Arrays;
+import org.testng.SkipException;
 import org.testng.annotations.Test;
 
 /** The file is for testing PyTorch MKLDNN functionalities. */
 public class MkldnnTest {
 
     @Test
     public void testMkldnn() {
+        if (!"amd64".equals(System.getProperty("os.arch"))) {
+            throw new SkipException("MKLDNN Test requires x86_64 arch.");
+        }
+
         System.setProperty("ai.djl.pytorch.use_mkldnn", "true");
         try (NDManager manager = NDManager.newBaseManager()) {
             NDArray[] arrays = {

diff --git a/engines/pytorch/pytorch-jni/build.gradle b/engines/pytorch/pytorch-jni/build.gradle
@@ -23,17 +23,23 @@ processResources {
                 "linux-x86_64/cu102/libdjl_torch.so",
                 "linux-x86_64/cpu-precxx11/libdjl_torch.so",
                 "osx-x86_64/cpu/libdjl_torch.dylib",
-                "win-x86_64/cpu/djl_torch.dll",
-                "win-x86_64/cu102/djl_torch.dll"
+                "win-x86_64/cpu/djl_torch.dll"
         ]
-        if (ptVersion.startsWith("1.10.")) {
+        if (ptVersion.startsWith("1.11.")) {
+            files.add("linux-aarch64/cpu/libdjl_torch.so")
             files.add("linux-x86_64/cu113/libdjl_torch.so")
             files.add("linux-x86_64/cu113-precxx11/libdjl_torch.so")
             files.add("win-x86_64/cu113/djl_torch.dll")
+        } else if (ptVersion.startsWith("1.10.")) {
+            files.add("linux-x86_64/cu113/libdjl_torch.so")
+            files.add("linux-x86_64/cu113-precxx11/libdjl_torch.so")
+            files.add("win-x86_64/cu113/djl_torch.dll")
+            files.add("win-x86_64/cu102/djl_torch.dll")
         } else {
             files.add("linux-x86_64/cu111/libdjl_torch.so")
             files.add("linux-x86_64/cu111-precxx11/libdjl_torch.so")
             files.add("win-x86_64/cu111/djl_torch.dll")
+            files.add("win-x86_64/cu102/djl_torch.dll")
         }
         String jnilibDir = "${project.projectDir}/jnilib/${djl_version}"
         files.each { entry ->

diff --git a/engines/pytorch/pytorch-native/CMakeLists.txt b/engines/pytorch/pytorch-native/CMakeLists.txt
@@ -48,6 +48,10 @@ set(SOURCE_FILES
     "src/main/native/ai_djl_pytorch_jni_cache.h"
     "src/main/native/ai_djl_pytorch_jni_cache.cc")
 
+if(PT_OLD_VERSION)
+    add_compile_definitions(V1_10_X)
+endif()
+
 add_library(djl_torch SHARED ${SOURCE_FILES})
 # build host
 if(NOT BUILD_ANDROID)

diff --git a/engines/pytorch/pytorch-native/build.cmd b/engines/pytorch/pytorch-native/build.cmd
@@ -5,7 +5,7 @@
 @rem choco install jdk8 -y
 
 set FILEPATH="libtorch"
-set VERSION="%1"
+set VERSION=%1
 if "%2" == "cpu" (
     set DOWNLOAD_URL="https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-%VERSION%%%2Bcpu.zip"
 ) else if "%2" == "cu102" (
@@ -28,9 +28,19 @@ if exist %FILEPATH% (
     echo Finished downloading libtorch
 )
 
+if "%VERSION%" == "1.11.0" (
+    copy /y src\main\patch\cuda.cmake libtorch\share\cmake\Caffe2\public\
+)
+if "%VERSION%" == "1.10.0" (
+    set PT_OLD_VERSION=1
+)
+if "%VERSION%" == "1.9.1" (
+    set PT_OLD_VERSION=1
+)
+
 if exist build rd /q /s build
 md build\classes
 cd build
 javac -sourcepath ..\..\pytorch-engine\src\main\java\ ..\..\pytorch-engine\src\main\java\ai\djl\pytorch\jni\PyTorchLibrary.java -h include -d classes
-cmake -DCMAKE_PREFIX_PATH=libtorch ..
+cmake -DCMAKE_PREFIX_PATH=libtorch -DPT_OLD_VERSION=%PT_OLD_VERSION% ..
 cmake --build . --config Release
diff --git a/engines/pytorch/pytorch-native/build.gradle b/engines/pytorch/pytorch-native/build.gradle
@@ -12,7 +12,7 @@ if (project.hasProperty("pt_version") && project.property("pt_version") != "") {
 }
 boolean isRelease = project.hasProperty("release") || project.hasProperty("staging")
 boolean isPrecxx11 = project.hasProperty("precxx11")
-boolean isAarch64 = project.hasProperty("aarch64")
+boolean isAarch64 = project.hasProperty("aarch64") || System.properties["os.arch"] == "aarch64"
 
 String FLAVOR = "cpu"
 if (project.hasProperty("cu10")) {
@@ -91,14 +91,13 @@ def prepareNativeLib(String binaryRoot, String ver) {
             "cpu/libtorch-win-shared-with-deps-${ver}%2Bcpu.zip"              : "cpu/win-x86_64",
             "cu102/libtorch-cxx11-abi-shared-with-deps-${ver}%2Bcu102.zip"    : "cu102/linux-x86_64",
             "${cu11}/libtorch-cxx11-abi-shared-with-deps-${ver}%2B${cu11}.zip": "${cu11}/linux-x86_64",
-            "cu102/libtorch-win-shared-with-deps-${ver}%2Bcu102.zip"          : "cu102/win-x86_64",
             "${cu11}/libtorch-win-shared-with-deps-${ver}%2B${cu11}.zip"      : "${cu11}/win-x86_64",
             "cpu/libtorch-shared-with-deps-${ver}%2Bcpu.zip"                  : "cpu-precxx11/linux-x86_64",
             "${cu11}/libtorch-shared-with-deps-${ver}%2B${cu11}.zip"          : "${cu11}-precxx11/linux-x86_64"
     ]
 
     def aarch64Files = [
-            "${ver}/libtorch-cxx11-shared-with-deps-${ver}-aarch64.zip"       : "cpu/linux-aarch64"
+            "${ver}/libtorch-cxx11-shared-with-deps-${ver}-aarch64.zip": "cpu/linux-aarch64"
     ]
 
     copyNativeLibToOutputDir(files, binaryRoot, officialPytorchUrl)
@@ -130,7 +129,7 @@ def copyNativeLibToOutputDir(Map<String, String> fileStoreMap, String binaryRoot
             }
             into("${outputDir}/native/lib")
         }
-        if (entry.value.contains("-precxx11")) {
+        if (entry.value.contains("-precxx11") && !entry.value.contains("aarch64")) {
             def libstd = new File("${outputDir}/native/lib/libstdc++.so.6")
             new URL("https://publish.djl.ai/extra/libstdc%2B%2B.so.6").withInputStream {
                 i -> libstd.withOutputStream { it << i }
@@ -253,11 +252,10 @@ task uploadS3 {
                 "${BINARY_ROOT}/cpu/win-x86_64/native/lib/",
                 "${BINARY_ROOT}/cpu-precxx11/linux-x86_64/native/lib/",
                 "${BINARY_ROOT}/cu102/linux/native/lib/",
-                "${BINARY_ROOT}/cu102/win-x86_64/native/lib/",
                 "${BINARY_ROOT}/cu113/linux-x86_64/native/lib/",
                 "${BINARY_ROOT}/cu113/win-x86_64/native/lib/",
                 "${BINARY_ROOT}/cu113-precxx11/linux-x86_64/native/lib/",
-                "${BINARY_ROOT}/cpu/linux-aarch64/native/lib"
+                "${BINARY_ROOT}/cpu/linux-aarch64/native/lib/"
         ]
         uploadDirs.each { item ->
             fileTree(item).files.name.each {
@@ -315,11 +313,11 @@ flavorNames.each { flavor ->
                     libstd.text = new URL("https://publish.djl.ai/extra/THIRD-PARTY-LICENSES_qHnMKgbdWa.txt").text
                 }
             }
-            from ("${BINARY_ROOT}/${flavor}/${osName}/native/lib") {
-                into ("pytorch/${flavor}/${osName}")
+            from("${BINARY_ROOT}/${flavor}/${osName}/native/lib") {
+                into("pytorch/${flavor}/${osName}")
             }
-            from ("${BINARY_ROOT}/pytorch.properties") {
-                into ("native/lib")
+            from("${BINARY_ROOT}/pytorch.properties") {
+                into("native/lib")
             }
             from "src/main/resources"
             archiveClassifier = "${osName}"

diff --git a/engines/pytorch/pytorch-native/build.sh b/engines/pytorch/pytorch-native/build.sh
@@ -13,9 +13,11 @@ fi
 PLATFORM=$(uname | tr '[:upper:]' '[:lower:]')
 VERSION=$1
 FLAVOR=$2
+AARCH64_CXX11ABI="-cxx11"
 CXX11ABI="-cxx11-abi"
 if [[ $3 == "precxx11" ]]; then
   CXX11ABI=""
+  AARCH64_CXX11ABI=""
 fi
 ARCH=$4
 
@@ -27,7 +29,7 @@ if [[ ! -d "libtorch" ]]; then
     fi
 
     if [[ $ARCH == 'aarch64' ]]; then
-      curl -s https://djl-ai.s3.amazonaws.com/publish/pytorch/${VERSION}/libtorch-cxx11-shared-with-deps-${VERSION}-aarch64.zip | jar xv
+      curl -s https://djl-ai.s3.amazonaws.com/publish/pytorch/${VERSION}/libtorch${AARCH64_CXX11ABI}-shared-with-deps-${VERSION}-aarch64.zip | jar xv
     else
       curl -s https://download.pytorch.org/libtorch/${FLAVOR}/libtorch${CXX11ABI}-shared-with-deps-${VERSION}%2B${FLAVOR}.zip | jar xv
     fi
@@ -40,13 +42,16 @@ if [[ ! -d "libtorch" ]]; then
   fi
 fi
 
+if [[ "$VERSION" =~ ^1\.10\..*|^1\.9\..* ]]; then
+  PT_OLD_VERSION=1
+fi
 pushd .
 
 rm -rf build
 mkdir build && cd build
 mkdir classes
 javac -sourcepath ../../pytorch-engine/src/main/java/ ../../pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java -h include -d classes
-cmake -DCMAKE_PREFIX_PATH=libtorch ..
+cmake -DCMAKE_PREFIX_PATH=libtorch -DPT_OLD_VERSION=${PT_OLD_VERSION} ..
 cmake --build . --config Release -- -j "${NUM_PROC}"
 
 if [[ $PLATFORM == 'darwin' ]]; then

diff --git a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_system.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_system.cc
@@ -12,7 +12,12 @@
  */
 #include <torch/torch.h>
 // clang-format off
-#include <torch/csrc/jit/frontend/code_template.h>
+#ifdef V1_10_X
+    #include <torch/csrc/jit/frontend/code_template.h>
+#else
+    #include <ATen/code_template.h>
+#endif
+#include <ATen/core/jit_type.h>
 // clang-format on
 
 #include <sstream>
@@ -162,7 +167,8 @@ inline std::string FormatMemory(int64_t bytes) {
   return oss.str();
 }
 
-// the code snippet is copied from torch/csrc/autograd/profiler.cpp
+// the code snippet is copied from torch/csrc/autograd/profiler_legacy.cpp
+#ifdef V1_10_X
 static torch::jit::CodeTemplate event_template(R"(
 {
   "name": "${name}",
@@ -175,6 +181,20 @@ static torch::jit::CodeTemplate event_template(R"(
   "cpu mem": "${cpu_mem}",
   "args": {}
 })");
+#else
+static const at::jit::CodeTemplate event_template(R"(
+{
+  "name": "${name}",
+  "ph": "X",
+  "ts": ${ts},
+  "dur": ${dur},
+  "tid": ${tid},
+  "pid": "CPU Functions",
+  "shape": ${shape},
+  "cpu mem": "${cpu_mem}",
+  "args": {}
+})");
+#endif
 
 // The function doesn't support GPU yet
 // You can refer to
@@ -227,7 +247,11 @@ void WriteProfilerEventsToStream(std::ostream& out, const std::vector<std::vecto
         LegacyEvent* start = it->second;
         int64_t memory_usage = mem_it->second;
 
+#ifdef V1_10_X
         torch::jit::TemplateEnv env;
+#else
+        at::jit::TemplateEnv env;
+#endif
         env.s("name", start->name());
         env.d("ts", profiler_start->cpuElapsedUs(*start));
         env.d("dur", start->cpuElapsedUs(*evt));

diff --git a/engines/pytorch/pytorch-native/src/main/native/djl_pytorch_utils.h b/engines/pytorch/pytorch-native/src/main/native/djl_pytorch_utils.h
@@ -30,9 +30,15 @@ namespace utils {
 
 #if !defined(__ANDROID__)
 // for image interpolation
+#ifdef V1_10_X
 typedef torch::variant<torch::enumtype::kNearest, torch::enumtype::kLinear, torch::enumtype::kBilinear,
     torch::enumtype::kBicubic, torch::enumtype::kTrilinear, torch::enumtype::kArea>
     mode_t;
+#else
+typedef torch::variant<torch::enumtype::kNearest, torch::enumtype::kLinear, torch::enumtype::kBilinear,
+    torch::enumtype::kBicubic, torch::enumtype::kTrilinear, torch::enumtype::kArea, torch::enumtype::kNearestExact>
+    mode_t;
+#endif
 #endif
 
 inline jint GetDTypeFromScalarType(const torch::ScalarType& type) {
@@ -108,6 +114,10 @@ inline mode_t GetInterpolationMode(jint jmode) {
       return torch::kTrilinear;
     case 5:
       return torch::kArea;
+    case 6:
+#ifndef V1_10_X
+      return torch::kNearestExact;
+#endif
     default:
       throw;
   }