Merge pull request #76 from keichi/add-cuda13-wheel

Add CUDA 13 wheel build

Author: keichi

.github/workflows/main.yml

@@ -145,14 +145,41 @@ jobs:
             yum install -y cuda-compiler-12-9 libcublas-devel-12-9 cuda-cudart-devel-12-9 cuda-driver-devel-12-9
           CIBW_TEST_SKIP: "*"
           CIBW_ENVIRONMENT: >
-            CMAKE_ARGS='-DKEDM_ENABLE_GPU=ON -DKokkos_ARCH_VOLTA70=ON'
+            CMAKE_ARGS='-DKEDM_ENABLE_GPU=ON -DKokkos_ARCH_AMPERE80=ON'
             PATH=/usr/local/cuda/bin:$PATH
           CIBW_REPAIR_WHEEL_COMMAND_LINUX: auditwheel repair -w {dest_dir} --exclude libcublas.so.12 --exclude libcublasLt.so.12 --exclude libcuda.so.1 {wheel}
       - uses: actions/upload-artifact@v4
         with:
           name: wheels-cuda12x
           path: ./wheelhouse/*.whl
 
+  build-wheels-cuda13x:
+    runs-on: ubuntu-22.04
+    if: github.event_name == 'push' && startsWith(github.event.ref, 'refs/tags/v')
+    steps:
+      - uses: actions/checkout@v4
+      - name: Replace project name with kedm-cuda13x
+        run: sed -i 's/name = "kedm"/name = "kedm-cuda13x"/' pyproject.toml
+      - name: Build wheels
+        uses: pypa/cibuildwheel@v3.3.0
+        env:
+          CIBW_BUILD: cp3*-manylinux_x86_64
+          CIBW_SKIP: cp38-* cp39-*
+          CIBW_BEFORE_ALL_LINUX: >
+            yum install -y yum-utils &&
+            yum-config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo &&
+            yum install -y cuda-compiler-13-0 libcublas-devel-13-0 cuda-cudart-devel-13-0 cuda-driver-devel-13-0 devtoolset-11-gcc-c++
+          CIBW_TEST_SKIP: "*"
+          CIBW_ENVIRONMENT: >
+            CMAKE_ARGS='-DKEDM_ENABLE_GPU=ON -DKokkos_ARCH_AMPERE80=ON'
+            PATH=/usr/local/cuda/bin:$PATH
+            CXX=/opt/rh/devtoolset-11/root/bin/g++
+          CIBW_REPAIR_WHEEL_COMMAND_LINUX: auditwheel repair -w {dest_dir} --exclude libcublas.so.13 --exclude libcublasLt.so.13 --exclude libcuda.so.1 {wheel}
+      - uses: actions/upload-artifact@v4
+        with:
+          name: wheels-cuda13x
+          path: ./wheelhouse/*.whl
+
   build-wheels-linux:
     runs-on: ubuntu-22.04
     if: github.event_name == 'push' && startsWith(github.event.ref, 'refs/tags/v')
@@ -195,7 +222,7 @@ jobs:
           path: ./wheelhouse/*.whl
 
   publish-wheels:
-    needs: [build-wheels-linux, build-wheels-macos, build-wheels-cuda12x]
+    needs: [build-wheels-linux, build-wheels-macos, build-wheels-cuda12x, build-wheels-cuda13x]
     runs-on: ubuntu-latest
     permissions:
       id-token: write

README.md

@@ -23,12 +23,18 @@ CPU (Linux and macOS)
 pip3 install kedm
 ```
 
-NVIDIA GPU (CUDA 12.0 or later)
+NVIDIA GPU (CUDA 12)
 
 ```
 pip3 install kedm-cuda12x
 ```
 
+NVIDIA GPU (CUDA 13)
+
+```
+pip3 install kedm-cuda13x
+```
+
 ## Citing
 
 Please cite the following papers if you find kEDM useful:

docs/setup.rst

@@ -5,9 +5,8 @@ Prebuilt wheels
 ---------------
 
 The easiest way to install kEDM is to use the prebuilt wheels. Prebuilt wheels
-for Linux (x86_64) and macOS (x86_64 and arm64) are available on PyPI. Note that
-these wheels are not built for GPUs. If you want to use GPUs, you need to
-install :ref:`from source <pip-source>`.
+for Linux (x86_64) and macOS (x86_64 and arm64) are available on PyPI. GPU
+wheels are distributed as separate packages, one for each CUDA major version.
 
 Requirements
 ============
@@ -18,10 +17,24 @@ Requirements
 Installation
 ============
 
+CPU (Linux and macOS):
+
 .. code-block:: bash
 
     $ pip3 install kedm
 
+NVIDIA GPU (CUDA 12):
+
+.. code-block:: bash
+
+    $ pip3 install kedm-cuda12x
+
+NVIDIA GPU (CUDA 13):
+
+.. code-block:: bash
+
+    $ pip3 install kedm-cuda13x
+
 .. _pip-source:
 
 Source