Try to build CUDA wheels

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

@@ -0,0 +1,98 @@
+# inspired by https://github.com/AutoGPTQ/AutoGPTQ/blob/main/.github/workflows/build_wheels_cuda.yml
+
+name: Build Python wheels with CUDA
+
+on:
+  push:
+    branches: [main]
+    tags: ["*"]
+  pull_request:
+    # Check all PR
+
+env:
+  SPHERICART_NO_LOCAL_DEPS: "1"
+
+jobs:
+  build_wheels:
+    name: ${{ matrix.os }}, torch ${{ matrix.torch }}, cuda ${{ matrix.cuda }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-20.04]
+        torch: ["2.0.0", "2.1.0", "2.2.0"]
+        cuda: ["11.8", "12.1"]
+        exclude:
+          # torch 1.13.0 does not support cuda 12.1
+          - os: ubuntu-20.04
+            torch: "1.13.1"
+            cuda: "12.1"
+
+    steps:
+      - name: Prepare cuda_no_point variable
+        id: prepare_cuda_no_point
+        run: |
+          export cuda_no_point=$(echo ${{ matrix.cuda }} | sed 's/\.//')
+          echo ::set-output name=cuda_no_point::$cuda_no_point
+
+      - name: free disk space
+        run: sudo rm -rf /usr/share/dotnet /usr/local/lib/android || true
+
+      - uses: actions/checkout@v3
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.10"
+
+      - name: Install cibuildwheel
+        run: python -m pip install cibuildwheel build
+
+      - name: Build sphericart wheels
+        run: |
+          # ensure we build the wheel from the sdist, not the checkout
+          python -m build --sdist . --outdir dist
+          python -m cibuildwheel dist/*.tar.gz --output-dir dist
+        env:
+          CIBW_BUILD_VERBOSITY: 3
+          # build wheels on CPython 3.10
+          CIBW_BUILD: cp39-* cp310-* cp311-* cp312-*
+          # skip musl and 32-bit builds
+          CIBW_SKIP: "*-musllinux* *-win32 *-manylinux_i686"
+          # on macOS, build both Intel & Apple Silicon versions
+          CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014
+          # do not build wheels with -march=native
+          CIBW_ENVIRONMENT: SPHERICART_ARCH_NATIVE=OFF
+
+      - name: Build sphericart-torch wheels
+        run: |
+          # ensure we build the wheel from the sdist, not the checkout
+          python -m build --sdist sphericart-torch --outdir sphericart-torch/dist
+          python -m cibuildwheel sphericart-torch/dist/*.tar.gz --output-dir sphericart-torch/dist
+        env:
+          CIBW_BEFORE_ALL: bash /host/home/runner/work/sphericart/sphericart/scripts/cibw-cuda-setup.sh ${{ matrix.cuda }}
+          CIBW_BUILD_VERBOSITY: 3
+          CIBW_BUILD: >
+            bash -c 'if [[ "${{ matrix.torch }}" == "2.2.0" ]]; then
+              echo "cp39-* cp310-* cp311-* cp312-*"
+            else
+              echo "cp39-* cp310-* cp311-*"
+            fi'
+          CIBW_SKIP: "*-musllinux* *-win32 *-manylinux_i686"
+          CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014
+          # set environment variables for sphericart-torch build
+          CIBW_ENVIRONMENT: SPHERICART_ARCH_NATIVE=OFF CUDACXX=/usr/local/cuda/bin/nvcc TORCH_CUDA_ARCH_LIST=All CUDAARCHS=all SPHERICART_TORCH_TORCH_VERSION=${{ matrix.torch }} PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/cu${{ steps.prepare_cuda_no_point.outputs.cuda_no_point }}
+          # do not complain for missing libtorch.so in sphericart-torch wheel
+          CIBW_REPAIR_WHEEL_COMMAND_LINUX: |
+            auditwheel repair --exclude libtorch.so --exclude libtorch_cpu.so --exclude libtorch_cuda.so --exclude libc10.so --exclude libc10_cuda.so -w {dest_dir} {wheel}
+
+      # - name: Build sphericart-jax wheels
+      #   run: |
+      #     # ensure we build the wheel from the sdist, not the checkout
+      #     python -m build --sdist sphericart-jax --outdir sphericart-jax/dist
+      #     python -m cibuildwheel sphericart-jax/dist/*.tar.gz --output-dir sphericart-jax/dist
+      #   env:
+      #     CIBW_BUILD_VERBOSITY: 3
+      #     CIBW_BUILD: cp310-*
+      #     CIBW_SKIP: "*-musllinux* *-win32 *-manylinux_i686"
+      #     CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014
+      #     CIBW_ENVIRONMENT: SPHERICART_ARCH_NATIVE=OFF

scripts/cibw-cuda-setup.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+
+# Set CUDA version and architecture
+CU_VER=${1//./-}
+ARCH="x86_64"
+
+# Install CUDA compiler and libraries
+yum install -y yum-utils
+yum-config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel7/${ARCH}/cuda-rhel7.repo
+yum -y install cuda-toolkit-${CU_VER}.${ARCH} \
+              nvidia-driver-latest-dkms
+
+# Clean up YUM caches
+yum clean all
+rm -rf /var/cache/yum/*
+
+# Configure dynamic linker run-time bindings
+echo "/usr/local/cuda/lib64" >> /etc/ld.so.conf.d/999_nvidia_cuda.conf
+
+# Set environment variables
+export PATH="/usr/local/cuda/bin:${PATH}"
+export LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}"
+export CUDA_HOME=/usr/local/cuda
+export CUDA_ROOT=/usr/local/cuda
+export CUDA_PATH=/usr/local/cuda
+export CUDADIR=/usr/local/cuda

sphericart-torch/MANIFEST.in

@@ -7,5 +7,7 @@ recursive-include sphericart *
 recursive-include src *
 recursive-include include *
 
+recursive-include build-backend *.py
+
 include pyproject.toml
 include README.md

sphericart-torch/build-backend/backend.py

@@ -0,0 +1,26 @@
+# this is a custom Python build backend wrapping setuptool's to set a
+# specific torch version as a build dependency, based on an environment
+# variable
+import os
+
+from setuptools import build_meta
+
+TORCH_VERSION = os.environ.get("SPHERICART_TORCH_TORCH_VERSION")
+
+if TORCH_VERSION is not None:
+    TORCH_DEP = f"torch =={TORCH_VERSION}"
+else:
+    TORCH_DEP = "torch >=1.13"
+
+
+prepare_metadata_for_build_wheel = build_meta.prepare_metadata_for_build_wheel
+build_wheel = build_meta.build_wheel
+build_sdist = build_meta.build_sdist
+
+
+def get_requires_for_build_wheel(config_settings=None):
+    defaults = build_meta.get_requires_for_build_wheel(config_settings)
+    return defaults + [TORCH_DEP]
+
+
+get_requires_for_build_sdist = build_meta.get_requires_for_build_sdist

sphericart-torch/pyproject.toml

@@ -42,9 +42,12 @@ requires = [
     "setuptools >=44",
     "wheel >=0.36",
     "cmake",
-    "torch >= 1.13",
 ]
-build-backend = "setuptools.build_meta"
+
+# use a custom build backend to add a dependency on a specific
+# version of torch+cuda
+build-backend = "backend"
+backend-path = ["build-backend"]
 
 
 [tool.setuptools]