Skip to content

Commit d4a9a4b

Browse files
frostedoysterfrostedoyster
authored
Add ARM wheels (#195)
1 parent 3704357 commit d4a9a4b

File tree

5 files changed

+75
-21
lines changed

5 files changed

+75
-21
lines changed

‎.github/workflows/build-jax-wheels.yml‎

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -52,12 +52,17 @@ jobs:
5252
strategy:
5353
matrix:
5454
python-version: ["39", "310", "311", "312", "313"]
55-
arch: ["arm64", "x86_64"]
56-
os: ["ubuntu-22.04", "macos-14"]
57-
exclude:
58-
# remove mismatched arch/os pairs
59-
- {os: macos-14, arch: x86_64}
60-
- {os: ubuntu-22.04, arch: arm64}
55+
os: [ubuntu-22.04, macos-14, ubuntu-22.04-arm]
56+
include:
57+
- os: ubuntu-22.04
58+
arch: x86_64
59+
name: x86_64 Linux
60+
- os: macos-14
61+
arch: arm64
62+
name: arm64 macOS
63+
- os: ubuntu-22.04-arm
64+
arch: aarch64
65+
name: aarch64 Linux
6166

6267
name: "Py${{matrix.python-version}} ${{matrix.os}} ${{matrix.arch}}"
6368

@@ -66,7 +71,7 @@ jobs:
6671
uses: actions/checkout@v4
6772

6873
- name: Build custom manylinux Docker Image with CUDA
69-
if: matrix.os == 'ubuntu-22.04'
74+
if: startsWith(matrix.os, 'ubuntu-')
7075
run: |
7176
docker build --no-cache \
7277
-t sphericart_manylinux_2_28_"${{ matrix.arch }}" \
@@ -88,6 +93,7 @@ jobs:
8893
CIBW_SKIP: "*-musllinux*"
8994
CIBW_ARCHS: "${{ matrix.arch }}"
9095
CIBW_MANYLINUX_X86_64_IMAGE: sphericart_manylinux_2_28_x86_64
96+
CIBW_MANYLINUX_AARCH64_IMAGE: sphericart_manylinux_2_28_aarch64
9197
CIBW_ENVIRONMENT: >
9298
SPHERICART_ARCH_NATIVE=OFF
9399
SPHERICART_NO_LOCAL_DEPS=1

‎.github/workflows/build-torch-wheels.yml‎

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -52,12 +52,17 @@ jobs:
5252
strategy:
5353
matrix:
5454
pytorch-version: ["2.4", "2.5", "2.6", "2.7", "2.8"]
55-
arch: ["arm64", "x86_64"]
56-
os: ["ubuntu-22.04", "macos-14"]
57-
exclude:
58-
# remove mismatched arch/os pairs
59-
- {os: macos-14, arch: x86_64}
60-
- {os: ubuntu-22.04, arch: arm64}
55+
os: [ubuntu-22.04, macos-14, ubuntu-22.04-arm]
56+
include:
57+
- os: ubuntu-22.04
58+
arch: x86_64
59+
name: x86_64 Linux
60+
- os: macos-14
61+
arch: arm64
62+
name: arm64 macOS
63+
- os: ubuntu-22.04-arm
64+
arch: aarch64
65+
name: aarch64 Linux
6166

6267
name: "Torch ${{matrix.pytorch-version}} ${{matrix.os}} ${{matrix.arch}}"
6368

@@ -66,7 +71,7 @@ jobs:
6671
uses: actions/checkout@v4
6772

6873
- name: Free disk space (Ubuntu)
69-
if: matrix.os == 'ubuntu-22.04'
74+
if: startsWith(matrix.os, 'ubuntu-')
7075
run: |
7176
sudo rm -rf /usr/local/.ghcup
7277
sudo rm -rf /opt/hostedtoolcache/CodeQL
@@ -76,7 +81,7 @@ jobs:
7681
sudo rm -rf /usr/local/share/boost
7782
7883
- name: Build custom manylinux Docker Image with CUDA
79-
if: matrix.os == 'ubuntu-22.04'
84+
if: startsWith(matrix.os, 'ubuntu-')
8085
run: |
8186
docker build --no-cache \
8287
-t sphericart_manylinux_2_28_"${{ matrix.arch }}" \
@@ -98,6 +103,7 @@ jobs:
98103
CIBW_SKIP: "*-musllinux*"
99104
CIBW_ARCHS: "${{ matrix.arch }}"
100105
CIBW_MANYLINUX_X86_64_IMAGE: sphericart_manylinux_2_28_x86_64
106+
CIBW_MANYLINUX_AARCH64_IMAGE: sphericart_manylinux_2_28_aarch64
101107
CIBW_REPAIR_WHEEL_COMMAND_LINUX: |
102108
auditwheel repair \
103109
--exclude libc10.so \
@@ -134,6 +140,9 @@ jobs:
134140
- name: arm64 macOS
135141
os: macos-14
136142
arch: arm64
143+
- name: aarch64 Linux
144+
os: ubuntu-22.04-arm
145+
arch: aarch64
137146
steps:
138147
- uses: actions/checkout@v4
139148

‎.github/workflows/build-wheels.yml‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,13 +62,16 @@ jobs:
6262
- name: arm64 macOS
6363
os: macos-14
6464
arch: arm64
65+
- name: aarch64 Linux
66+
os: ubuntu-22.04-arm
67+
arch: aarch64
6568

6669
steps:
6770
- name: Checkout code
6871
uses: actions/checkout@v4
6972

7073
- name: Build custom manylinux Docker Image with CUDA
71-
if: matrix.os == 'ubuntu-22.04'
74+
if: startsWith(matrix.os, 'ubuntu-')
7275
run: |
7376
docker build --no-cache \
7477
-t sphericart_manylinux_2_28_"${{ matrix.arch }}" \
@@ -90,6 +93,7 @@ jobs:
9093
CIBW_SKIP: "*-musllinux*"
9194
CIBW_ARCHS: "${{ matrix.arch }}"
9295
CIBW_MANYLINUX_X86_64_IMAGE: sphericart_manylinux_2_28_x86_64
96+
CIBW_MANYLINUX_AARCH64_IMAGE: sphericart_manylinux_2_28_aarch64
9397
CIBW_ENVIRONMENT: >
9498
SPHERICART_ARCH_NATIVE=OFF
9599
SPHERICART_NO_LOCAL_DEPS=1

‎scripts/manylinux_2_28_aarch64/Dockerfile‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
# Inspired by https://github.com/ameli/manylinux-cuda, adding cuda-nvtx
2+
3+
FROM quay.io/pypa/manylinux_2_28_aarch64
4+
5+
# ------------
6+
# Install cuda
7+
# ------------
8+
9+
ARG CUDA_VER="12-8"
10+
11+
RUN yum install -y yum-utils
12+
RUN yum-config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/sbsa/cuda-rhel8.repo
13+
RUN yum -y install cuda-compiler-${CUDA_VER} \
14+
cuda-libraries-${CUDA_VER} \
15+
cuda-libraries-devel-${CUDA_VER} \
16+
cuda-nvtx-${CUDA_VER}
17+
RUN yum clean all
18+
RUN rm -rf /var/cache/yum/*
19+
RUN echo "/usr/local/cuda/lib64" >> /etc/ld.so.conf.d/999_nvidia_cuda.conf
20+
21+
# -------------------------
22+
# Set environment variables
23+
# -------------------------
24+
25+
ENV PATH="/usr/local/cuda/bin:${PATH}"
26+
ENV LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}"
27+
ENV CUDA_HOME=/usr/local/cuda
28+
ENV CUDA_ROOT=/usr/local/cuda
29+
ENV CUDA_PATH=/usr/local/cuda
30+
ENV CUDADIR=/usr/local/cuda
31+
32+
# --------
33+
# Commands
34+
# --------
35+
36+
CMD ["/bin/bash"]

‎scripts/manylinux_2_28_x86_64/Dockerfile‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,14 +7,13 @@ FROM quay.io/pypa/manylinux_2_28_x86_64
77
# ------------
88

99
ARG CUDA_VER="12-8"
10-
ARG ARCH="x86_64"
1110

1211
RUN yum install -y yum-utils
1312
RUN yum-config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo
14-
RUN yum -y install cuda-compiler-${CUDA_VER}.${ARCH} \
15-
cuda-libraries-${CUDA_VER}.${ARCH} \
16-
cuda-libraries-devel-${CUDA_VER}.${ARCH} \
17-
cuda-nvtx-${CUDA_VER}.${ARCH}
13+
RUN yum -y install cuda-compiler-${CUDA_VER} \
14+
cuda-libraries-${CUDA_VER} \
15+
cuda-libraries-devel-${CUDA_VER} \
16+
cuda-nvtx-${CUDA_VER}
1817
RUN yum clean all
1918
RUN rm -rf /var/cache/yum/*
2019
RUN echo "/usr/local/cuda/lib64" >> /etc/ld.so.conf.d/999_nvidia_cuda.conf

0 commit comments

Comments
 (0)