gpu-build.yml

name: GPU Build

on:
  push:
    branches:
      - main
  pull_request:

jobs:
  build:
    runs-on: ubuntu-22.04
    container: nvidia/cuda:12.2.2-devel-ubuntu22.04

    steps:
      - name: Install git and FFTW
        run: apt-get update && apt-get install -y git libfftw3-dev
        env:
          DEBIAN_FRONTEND: "noninteractive"

      - name: Get the version info for NVCC
        run: nvcc --version

      - uses: actions/checkout@v4
        with:
          submodules: true
          fetch-depth: 0

      - uses: actions/setup-python@v5
        with:
          python-version: 3.12

      - name: Install Python dependencies
        run: |
          python -m pip install -U pip
          python -m pip install -U jax[cpu]

      - name: Compile the extension
        run: python -m pip install -v .
        env:
          CMAKE_ARGS: "-DJAX_FINUFFT_USE_CUDA=ON"

      - name: Check that the GPU extension was built
        run: python -c "import jax_finufft.jax_finufft_gpu"