Building DJL For aarch64

.github/workflows/native_jni_s3_pytorch.yml

@@ -47,7 +47,6 @@ jobs:
           PYTORCH_VERSION=${PYTORCH_VERSION:-$(cat gradle.properties | awk -F '=' '/pytorch_version/ {print $2}')}
           aws s3 sync engines/pytorch/pytorch-native/jnilib s3://djl-ai/publish/pytorch/${PYTORCH_VERSION}/jnilib
           aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/pytorch/${PYTORCH_VERSION}/jnilib*"
-
   build-pytorch-jni-linux:
     if: github.repository == 'deepjavalibrary/djl'
     runs-on: ubuntu-latest
@@ -67,11 +66,11 @@ jobs:
       - name: Install Environment
         run: |
           apt-get update
-          DEBIAN_FRONTEND=noninteractive apt-get install -y locales cmake curl unzip software-properties-common
+          DEBIAN_FRONTEND=noninteractive apt-get install -y locales cmake curl unzip software-properties-common gcc-aarch64-linux-gnu g++-aarch64-linux-gnu
           add-apt-repository -y ppa:deadsnakes/ppa
           apt-get update
           apt-get install -y python3 python3-distutils
-          curl -O https://bootstrap.pypa.io/get-pip.py
+          curl -O https://bootstrap.pypa.io/pip/3.6/get-pip.py
           python3 get-pip.py
           pip3 install awscli --upgrade
       - name: Release JNI prep
@@ -85,6 +84,8 @@ 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:

engines/pytorch/pytorch-native/build.gradle

@@ -12,6 +12,8 @@ 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")
+
 String FLAVOR = "cpu"
 if (project.hasProperty("cu10")) {
     FLAVOR = "cu102"
@@ -26,15 +28,16 @@ String BINARY_ROOT = "${project.buildDir}/download"
 
 version = VERSION + (isRelease ? "" : "-SNAPSHOT")
 
-def downloadBuild(String ver, String os, String flavor, Boolean isPrecxx11 = false) {
+def downloadBuild(String ver, String os, String flavor, Boolean isPrecxx11 = false, Boolean isAarch64 = false) {
+    def arch = isAarch64 ? "aarch64" : "x86_64"
     exec {
         if (os == "win") {
             commandLine "${project.projectDir}/build.cmd", "${ver}", "${flavor}"
         } else {
             if (isPrecxx11) {
-                commandLine 'bash', 'build.sh', "${ver}", "${flavor}", "precxx11"
+                commandLine 'bash', 'build.sh', "${ver}", "${flavor}", "precxx11", "${arch}"
             } else {
-                commandLine 'bash', 'build.sh', "${ver}", "${flavor}"
+                commandLine 'bash', 'build.sh', "${ver}", "${flavor}", "cxx11", "${arch}"
             }
         }
     }
@@ -43,7 +46,7 @@ def downloadBuild(String ver, String os, String flavor, Boolean isPrecxx11 = fal
     // the reason why we duplicate the folder here is to insert djl_version into the path
     // so different versions of JNI wouldn't override each other. We don't also want publishDir
     // to have djl_version as engine would require to know that during the System.load()
-    String classifier = "${os}-x86_64"
+    String classifier = "${os}-${arch}"
     String ciDir = "${project.projectDir}/jnilib/${djl_version}/${classifier}/${flavor}"
     if (isPrecxx11) {
         ciDir = "${ciDir}-precxx11"
@@ -73,7 +76,8 @@ def downloadBuildAndroid(String ver) {
 }
 
 def prepareNativeLib(String binaryRoot, String ver) {
-    def url = "https://download.pytorch.org/libtorch"
+    def officialPytorchUrl = "https://download.pytorch.org/libtorch"
+    def aarch64PytorchUrl = "https://djl-ai.s3.amazonaws.com/publish/pytorch"
     String cu11
     if (ver.startsWith("1.8.") || ver.startsWith("1.9.")) {
         cu11 = "cu111"
@@ -92,7 +96,17 @@ def prepareNativeLib(String binaryRoot, String ver) {
             "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"
     ]
-    files.each { entry ->
+
+    def aarch64Files = [
+            "${ver}/libtorch-cxx11-shared-with-deps-${ver}-aarch64.zip"       : "cpu/linux-aarch64"
+    ]
+
+    copyNativeLibToOutputDir(files, binaryRoot, officialPytorchUrl)
+    copyNativeLibToOutputDir(aarch64Files, binaryRoot, aarch64PytorchUrl)
+}
+
+def copyNativeLibToOutputDir(Map<String, String> fileStoreMap, String binaryRoot, String url) {
+    fileStoreMap.each { entry ->
         project.logger.lifecycle("Downloading ${url}/${entry.key}")
         def outputDir = new File("${binaryRoot}/${entry.value}")
         def file = new File("${outputDir}/libtorch.zip")
@@ -148,9 +162,9 @@ task compileJNI {
         if (System.properties['os.name'].toLowerCase(Locale.ROOT).contains("windows")) {
             downloadBuild("${VERSION}", "win", "${FLAVOR}")
         } else if (System.properties['os.name'].toLowerCase(Locale.ROOT).contains("mac")) {
-            downloadBuild("${VERSION}", "osx", "${FLAVOR}")
+            downloadBuild("${VERSION}", "osx", "${FLAVOR}", false, isAarch64)
         } else if (System.properties['os.name'].toLowerCase(Locale.ROOT).contains("linux")) {
-            downloadBuild("${VERSION}", "linux", "${FLAVOR}", isPrecxx11)
+            downloadBuild("${VERSION}", "linux", "${FLAVOR}", isPrecxx11, isAarch64)
         } else {
             throw new IllegalStateException("Unknown Architecture " + System.properties['os.name'] + "-${FLAVOR}")
         }
@@ -227,7 +241,7 @@ task downloadPyTorchNativeLib() {
 task uploadS3 {
     doLast {
         delete "${BINARY_ROOT}"
-        prepareNativeLib("${BINARY_ROOT}", "${version}")
+        prepareNativeLib("${BINARY_ROOT}", "${VERSION}")
 
         exec {
             commandLine "sh", "-c", "find ${BINARY_ROOT} -type f | xargs gzip"
@@ -242,7 +256,8 @@ task uploadS3 {
                 "${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}/cu113-precxx11/linux-x86_64/native/lib/",
+                "${BINARY_ROOT}/cpu/linux-aarch64/native/lib"
         ]
         uploadDirs.each { item ->
             fileTree(item).files.name.each {

engines/pytorch/pytorch-native/build.sh

@@ -17,14 +17,21 @@ CXX11ABI="-cxx11-abi"
 if [[ $3 == "precxx11" ]]; then
   CXX11ABI=""
 fi
+ARCH=$4
 
 if [[ ! -d "libtorch" ]]; then
   if [[ $PLATFORM == 'linux' ]]; then
     if [[ ! "$FLAVOR" =~ ^(cpu|cu102|cu111|cu113)$ ]]; then
       echo "$FLAVOR is not supported."
       exit 1
     fi
-    curl -s https://download.pytorch.org/libtorch/${FLAVOR}/libtorch${CXX11ABI}-shared-with-deps-${VERSION}%2B${FLAVOR}.zip | jar xv
+
+    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
+    else
+      curl -s https://download.pytorch.org/libtorch/${FLAVOR}/libtorch${CXX11ABI}-shared-with-deps-${VERSION}%2B${FLAVOR}.zip | jar xv
+    fi
+
   elif [[ $PLATFORM == 'darwin' ]]; then
     curl -s https://download.pytorch.org/libtorch/cpu/libtorch-macos-${VERSION}.zip | jar xv
   else

gradle.properties

@@ -9,7 +9,7 @@ systemProp.org.gradle.internal.publish.checksums.insecure=true
 
 djl_version=0.16.0
 mxnet_version=1.9.0
-pytorch_version=1.10.0
+pytorch_version=1.10.2
 tensorflow_version=2.7.0
 tflite_version=2.6.2
 dlr_version=1.6.0