update

Author: akihironitta

.github/workflows/_build_docker.yml

@@ -10,15 +10,15 @@ defaults:
 
 jobs:
   docker:
-    runs-on: ubuntu-latest
+    runs-on: ${{ matrix.runner }}
     permissions:
       contents: read
       packages: write
     strategy:
       fail-fast: false
       matrix:
         cuda-version: [cu126, cu128, cu130, cu132]
-        arch: [x86_64]
+        arch: [x86_64, aarch64]
         include:
           - cuda-version: cu126
             cuda-number: "12.6"
@@ -28,6 +28,10 @@ jobs:
             cuda-number: "13.0"
           - cuda-version: cu132
             cuda-number: "13.2"
+          - arch: x86_64
+            runner: ubuntu-latest
+          - arch: aarch64
+            runner: ubuntu-24.04-arm
 
     steps:
       - name: Checkout repository

.github/workflows/docker/install_cuda.sh

@@ -3,30 +3,36 @@ set -ex
 
 CUDA_VERSION="${1:?Specify cuda version, e.g. 12.9}"
 
+ARCH=$(uname -m)
+if [ "$ARCH" != "x86_64" ] && [ "$ARCH" != "aarch64" ]; then
+    echo "Unsupported arch: $ARCH"
+    exit 1
+fi
+
 # See https://developer.nvidia.com/cuda-toolkit-archive
 if [ "$CUDA_VERSION" == "13.2" ]; then
-    URL=https://developer.download.nvidia.com/compute/cuda/13.2.0/local_installers/cuda-repo-rhel8-13-2-local-13.2.0_595.45.04-1.x86_64.rpm
+    URL=https://developer.download.nvidia.com/compute/cuda/13.2.0/local_installers/cuda-repo-rhel8-13-2-local-13.2.0_595.45.04-1.${ARCH}.rpm
     PACKAGE_NAME=cuda-toolkit-13-2
 elif [ "$CUDA_VERSION" == "13.0" ]; then
-    URL=https://developer.download.nvidia.com/compute/cuda/13.0.0/local_installers/cuda-repo-rhel8-13-0-local-13.0.0_580.65.06-1.x86_64.rpm
+    URL=https://developer.download.nvidia.com/compute/cuda/13.0.0/local_installers/cuda-repo-rhel8-13-0-local-13.0.0_580.65.06-1.${ARCH}.rpm
     PACKAGE_NAME=cuda-toolkit-13-0
 elif [ "$CUDA_VERSION" == "12.9" ]; then
-    URL=https://developer.download.nvidia.com/compute/cuda/12.9.1/local_installers/cuda-repo-rhel8-12-9-local-12.9.1_575.57.08-1.x86_64.rpm
+    URL=https://developer.download.nvidia.com/compute/cuda/12.9.1/local_installers/cuda-repo-rhel8-12-9-local-12.9.1_575.57.08-1.${ARCH}.rpm
     PACKAGE_NAME=cuda-toolkit-12-9
 elif [ "$CUDA_VERSION" == "12.8" ]; then
-    URL=https://developer.download.nvidia.com/compute/cuda/12.8.1/local_installers/cuda-repo-rhel8-12-8-local-12.8.1_570.124.06-1.x86_64.rpm
+    URL=https://developer.download.nvidia.com/compute/cuda/12.8.1/local_installers/cuda-repo-rhel8-12-8-local-12.8.1_570.124.06-1.${ARCH}.rpm
     PACKAGE_NAME=cuda-toolkit-12-8
 elif [ "$CUDA_VERSION" == "12.6" ]; then
-    URL=https://developer.download.nvidia.com/compute/cuda/12.6.3/local_installers/cuda-repo-rhel8-12-6-local-12.6.3_560.35.05-1.x86_64.rpm
+    URL=https://developer.download.nvidia.com/compute/cuda/12.6.3/local_installers/cuda-repo-rhel8-12-6-local-12.6.3_560.35.05-1.${ARCH}.rpm
     PACKAGE_NAME=cuda-toolkit-12-6
 elif [ "$CUDA_VERSION" == "12.4" ]; then
-    URL=https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda-repo-rhel8-12-4-local-12.4.1_550.54.15-1.x86_64.rpm
+    URL=https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda-repo-rhel8-12-4-local-12.4.1_550.54.15-1.${ARCH}.rpm
     PACKAGE_NAME=cuda-toolkit-12-4
 elif [ "$CUDA_VERSION" == "12.1" ]; then
-    URL=https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installers/cuda-repo-rhel8-12-1-local-12.1.1_530.30.02-1.x86_64.rpm
+    URL=https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installers/cuda-repo-rhel8-12-1-local-12.1.1_530.30.02-1.${ARCH}.rpm
     PACKAGE_NAME=cuda
 elif [ "$CUDA_VERSION" == "11.8" ]; then
-    URL=https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-rhel8-11-8-local-11.8.0_520.61.05-1.x86_64.rpm
+    URL=https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-rhel8-11-8-local-11.8.0_520.61.05-1.${ARCH}.rpm
     PACKAGE_NAME=cuda
 else
     echo "Invalid CUDA version: $CUDA_VERSION"