From f42b9ece25b7d9e494565b1f53953aba6b5cdf8a Mon Sep 17 00:00:00 2001
From: Guillaume Fraux <guillaume.fraux@epfl.ch>
Date: Fri, 29 Nov 2024 14:13:37 +0100
Subject: [PATCH] Build wheels for aarch64-linux

---
 .github/workflows/build-wheels.yml               | 16 ++++++++++++----
 .../rustc-manylinux2014_aarch64/Dockerfile       | 11 +++++++++++
 .../rustc-manylinux2014_x86_64/Dockerfile        |  2 +-
 3 files changed, 24 insertions(+), 5 deletions(-)
 create mode 100644 python/scripts/rustc-manylinux2014_aarch64/Dockerfile

diff --git a/.github/workflows/build-wheels.yml b/.github/workflows/build-wheels.yml
index e4afe7f3e..06c5da867 100644
--- a/.github/workflows/build-wheels.yml
+++ b/.github/workflows/build-wheels.yml
@@ -25,6 +25,10 @@ jobs:
             os: ubuntu-22.04
             rust-target: x86_64-unknown-linux-gnu
             cibw-arch: x86_64
+          - name: arm64 Linux
+            os: ubuntu-22.04
+            rust-target: aarch64-unknown-linux-gnu
+            cibw-arch: aarch64
           - name: x86_64 macOS
             os: macos-13
             rust-target: x86_64-apple-darwin
@@ -35,7 +39,6 @@ jobs:
             cibw-arch: arm64
           - name: x86_64 Windows
             os: windows-2019
-            # TODO: add a 32-bit windows builder?
             rust-target: x86_64-pc-windows-msvc
             cibw-arch: AMD64
     steps:
@@ -57,18 +60,23 @@ jobs:
       - name: install dependencies
         run: python -m pip install cibuildwheel twine
 
+      - name: Set up QEMU for docker
+        if: matrix.os == 'ubuntu-22.04'
+        uses: docker/setup-qemu-action@v3
+
       - name: build manylinux with rust docker image
         if: matrix.os == 'ubuntu-22.04'
-        run: docker build -t rustc-manylinux2014_x86_64 python/scripts/rustc-manylinux2014_x86_64
+        run: docker build -t rustc-manylinux2014_${{ matrix.cibw-arch }} python/scripts/rustc-manylinux2014_${{ matrix.cibw-arch }}
 
       - name: build featomic wheel
         run: python -m cibuildwheel python/featomic
         env:
-          CIBW_BUILD: cp310-*
+          CIBW_BUILD: cp312-*
           CIBW_SKIP: "*musllinux*"
           CIBW_ARCHS: ${{ matrix.cibw-arch }}
           CIBW_BUILD_VERBOSITY: 2
           CIBW_MANYLINUX_X86_64_IMAGE: rustc-manylinux2014_x86_64
+          CIBW_MANYLINUX_AARCH64_IMAGE: rustc-manylinux2014_aarch64
           CIBW_REPAIR_WHEEL_COMMAND_MACOS: "delocate-wheel --ignore-missing-dependencies --require-archs {delocate_archs} -w {dest_dir} -v {wheel}"
           CIBW_REPAIR_WHEEL_COMMAND_LINUX: "auditwheel repair --exclude libmetatensor.so -w {dest_dir} {wheel}"
           CIBW_ENVIRONMENT: >
@@ -139,7 +147,7 @@ jobs:
           pattern: featomic-wheel-*
           merge-multiple: true
 
-      - name: Download other wheels and sdists
+      - name: Download sdists
         uses: actions/download-artifact@v4
         with:
           path: wheels
diff --git a/python/scripts/rustc-manylinux2014_aarch64/Dockerfile b/python/scripts/rustc-manylinux2014_aarch64/Dockerfile
new file mode 100644
index 000000000..4fa131934
--- /dev/null
+++ b/python/scripts/rustc-manylinux2014_aarch64/Dockerfile
@@ -0,0 +1,11 @@
+# Use manylinux docker image as a base
+FROM quay.io/pypa/manylinux2014_aarch64
+
+RUN yum install git -y
+RUN git config --global --add safe.directory /code
+
+# Download rustup-init and install
+RUN curl https://sh.rustup.rs -sSf | sh -s -- -y --profile minimal --default-toolchain 1.75
+
+ENV PATH="/root/.cargo/bin:${PATH}"
+ENV RUST_BUILD_TARGET="aarch64-unknown-linux-gnu"
diff --git a/python/scripts/rustc-manylinux2014_x86_64/Dockerfile b/python/scripts/rustc-manylinux2014_x86_64/Dockerfile
index 4745f73d2..d4fe44c2f 100644
--- a/python/scripts/rustc-manylinux2014_x86_64/Dockerfile
+++ b/python/scripts/rustc-manylinux2014_x86_64/Dockerfile
@@ -4,7 +4,7 @@ FROM quay.io/pypa/manylinux2014_x86_64
 RUN yum install git -y
 RUN git config --global --add safe.directory /code
 
-# Download rustup-init asn install
+# Download rustup-init and install
 RUN curl https://sh.rustup.rs -sSf | sh -s -- -y --profile minimal --default-toolchain 1.75
 
 ENV PATH="/root/.cargo/bin:${PATH}"