Add linux aarch64 wheel build support

Signed-off-by: odidev <odidev@puresoftware.com>
odidev · Apr 12, 2022 · 9fc79c5 · 9fc79c5
1 parent 3e1b48c
commit 9fc79c5
Show file tree

Hide file tree

Showing 4 changed files with 63 additions and 10 deletions.
diff --git a/.travis.yml b/.travis.yml
@@ -9,6 +9,7 @@ env:
 branches:
   only:
     - master
+    - arm64_test
     - /^[0-9.]*$/
 stages:
   - test
@@ -107,6 +108,15 @@ jobs:
       language: python
       python: 3.6
       env: PYVER=3.6 PYNUM=3 PYTHON_INSTALL=pip BUILD_ARCH=x86_64
+    - os: linux
+      dist: xenial
+      arch: arm64-graviton2
+      virt: lxd
+      group: edge
+      compiler: gcc
+      language: python
+      python: 3.6
+      env: PYVER=3.6 PYNUM=3 PYTHON_INSTALL=pip BUILD_ARCH=aarch64
     - os: linux
       compiler: gcc
       language: cpp
@@ -145,6 +155,24 @@ jobs:
       language: python
       python: 3.8
       env: PYVER=3.8 PYNUM=3 PYTHON_INSTALL=pip BUILD_ARCH=x86_64
+    - os: linux
+      dist: xenial
+      arch: arm64-graviton2
+      virt: lxd
+      group: edge
+      compiler: gcc
+      language: python
+      python: 3.7
+      env: PYVER=3.7 PYNUM=3 PYTHON_INSTALL=pip BUILD_ARCH=aarch64
+    - os: linux
+      dist: xenial
+      arch: arm64-graviton2
+      virt: lxd
+      group: edge
+      compiler: gcc
+      language: python
+      python: 3.8
+      env: PYVER=3.8 PYNUM=3 PYTHON_INSTALL=pip BUILD_ARCH=aarch64
     - os: linux
       compiler: gcc
       language: cpp

diff --git a/.travis/build_linux_wheel.sh b/.travis/build_linux_wheel.sh
@@ -2,7 +2,11 @@
 set -xe
 
 # To be run inside docker container
-export CMAKE=cmake28 EIGEN3_INCLUDE_DIR="$TRAVIS_BUILD_DIR/eigen" LD_LIBRARY_PATH="$TRAVIS_BUILD_DIR/build/dynet:$LD_LIBRARY_PATH"
+if [[ "$BUILD_ARCH" == aarch64 ]]; then
+  export CMAKE=cmake EIGEN3_INCLUDE_DIR="$TRAVIS_BUILD_DIR/eigen" LD_LIBRARY_PATH="$TRAVIS_BUILD_DIR/build/dynet:$LD_LIBRARY_PATH"
+else
+  export CMAKE=cmake28 EIGEN3_INCLUDE_DIR="$TRAVIS_BUILD_DIR/eigen" LD_LIBRARY_PATH="$TRAVIS_BUILD_DIR/build/dynet:$LD_LIBRARY_PATH"
+fi
 cd "$TRAVIS_BUILD_DIR"
 
 if [[ "$BUILD_ARCH" == i686 ]]; then
@@ -11,7 +15,7 @@ else
   yum install -y gmp-devel
 fi
 # Compile wheels
-for PYBIN in /opt/python/*${PYVER/./}*/bin; do
+for PYBIN in /opt/python/cp*${PYVER/./}*/bin; do
   "$PYBIN/pip" install -U pip
   "$PYBIN/pip" install --prefer-binary cryptography
   "$PYBIN/pip" install -U numpy twine cython

diff --git a/docker/Dockerfile-aarch64 b/docker/Dockerfile-aarch64
@@ -0,0 +1,4 @@
+FROM quay.io/pypa/manylinux2014_aarch64
+MAINTAINER Daniel Hershcovich
+ADD .travis/ /root/.travis
+WORKDIR /root/.travis
diff --git a/dynet/mem.cc b/dynet/mem.cc
@@ -10,8 +10,10 @@
 
 #include <fcntl.h>
 #if !_WINDOWS
+#if !__ARM_ARCH_ISA_A64
 #include <mm_malloc.h>
 #endif
+#endif
 #include "dynet/except.h"
 #include "dynet/devices.h"
 #if HAVE_CUDA
@@ -27,17 +29,32 @@ namespace dynet {
 MemAllocator::~MemAllocator() {}
 
 void* CPUAllocator::malloc(size_t n) {
-  void* ptr = _mm_malloc(n, align);
-  if (!ptr) {
-    show_pool_mem_info();
-    cerr << "CPU memory allocation failed n=" << n << " align=" << align << endl;
-    throw dynet::out_of_memory("CPU memory allocation failed");
-  }
-  return ptr;
+  #if !__ARM_ARCH_ISA_A64
+    void* ptr = _mm_malloc(n, align);
+    if (!ptr) {
+      show_pool_mem_info();
+      cerr << "CPU memory allocation failed n=" << n << " align=" << align << endl;
+      throw dynet::out_of_memory("CPU memory allocation failed");
+    }
+    return ptr;
+  #else
+    void* ptr = aligned_alloc(align, n);
+    if (!ptr) {
+      show_pool_mem_info();
+      cerr << "CPU memory allocation failed n=" << n << " align=" << align << endl;
+      throw dynet::out_of_memory("CPU memory allocation failed");
+    }
+    return ptr;
+  #endif
+
 }
 
 void CPUAllocator::free(void* mem) {
-  _mm_free(mem);
+  #if !__ARM_ARCH_ISA_A64
+    _mm_free(mem);
+  #else
+    free(mem);
+  #endif
 }
 
 void CPUAllocator::zero(void* p, size_t n) {