Bbonev/wheels

.github/workflows/build_and_deploy.yml

@@ -0,0 +1,193 @@
+name: Build and Deploy Manylinux Wheels
+
+# Trigger workflow on PRs, tags, and manual dispatch
+on:
+  push:
+    tags:
+      - 'v*'
+  pull_request:
+    branches:
+      - main
+  workflow_dispatch:
+    inputs:
+      torch_version:
+        description: 'PyTorch version to build (e.g., 2.4.0)'
+        required: false
+        default: ''
+      cuda_version:
+        description: 'CUDA version to build (e.g., cu121, cpu)'
+        required: false
+        default: ''
+
+env:
+  CUDA_ARCH_MAP: '{"11.8": "8.0;8.6;8.9;+PTX", "12.1": "8.0;8.6;8.9;9.0;+PTX", "12.4": "8.0;8.6;8.9;9.0;+PTX", "12.6": "8.0;8.6;8.9;9.0;+PTX", "12.8": "8.0;8.6;8.9;9.0;10.0;12.0;+PTX", "none": ""}'
+
+jobs:
+  build-manylinux-wheels:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        include:
+          # # PyTorch 2.4.x with multiple CUDA versions
+          # - torch-version: '2.4.0'
+          #   torch-cuda: 'cpu'
+          #   cuda-version: 'none'
+          # - torch-version: '2.4.0'
+          #   torch-cuda: 'cu118'
+          #   cuda-version: '11.8'
+          # - torch-version: '2.4.0'
+          #   torch-cuda: 'cu121'
+          #   cuda-version: '12.1'
+          # - torch-version: '2.4.0'
+          #   torch-cuda: 'cu124'
+          #   cuda-version: '12.4'
+
+          # # PyTorch 2.5.x with multiple CUDA versions
+          # - torch-version: '2.5.0'
+          #   torch-cuda: 'cpu'
+          #   cuda-version: 'none'
+          # - torch-version: '2.5.0'
+          #   torch-cuda: 'cu118'
+          #   cuda-version: '11.8'
+          # - torch-version: '2.5.0'
+          #   torch-cuda: 'cu121'
+          #   cuda-version: '12.1'
+          # - torch-version: '2.5.0'
+          #   torch-cuda: 'cu124'
+          #   cuda-version: '12.4'
+
+          # PyTorch 2.6.x with multiple CUDA versions
+          # - torch-version: '2.6.0'
+          #   torch-cuda: 'cpu'
+          #   cuda-version: 'none'
+          # - torch-version: '2.6.0'
+          #   torch-cuda: 'cu118'
+          #   cuda-version: '11.8'
+          # - torch-version: '2.6.0'
+          #   torch-cuda: 'cu124'
+          #   cuda-version: '12.4'
+          # - torch-version: '2.6.0'
+          #   torch-cuda: 'cu126'
+          #   cuda-version: '12.6'
+
+          # PyTorch 2.7.x with multiple CUDA versions
+          - torch-version: '2.7.0'
+            torch-cuda: 'cpu'
+            cuda-version: 'none'
+          # - torch-version: '2.7.0'
+          #   torch-cuda: 'cu118'
+          #   cuda-version: '11.8'
+          # - torch-version: '2.7.0'
+          #   torch-cuda: 'cu126'
+          #   cuda-version: '12.6'
+          # - torch-version: '2.7.0'
+          #   torch-cuda: 'cu128'
+          #   cuda-version: '12.8'
+
+    steps:
+    - uses: actions/checkout@v4
+
+    # - name: Free Disk Space
+    #   run: |
+    #     sudo rm -rf /usr/local/.ghcup
+    #     sudo rm -rf /opt/hostedtoolcache/CodeQL
+    #     sudo rm -rf /usr/local/lib/android/sdk/ndk
+    #     sudo rm -rf /usr/share/dotnet
+    #     sudo rm -rf /opt/ghc
+    #     sudo rm -rf /usr/local/share/boost
+    #     sudo apt-get clean
+    #     echo "Disk space after cleanup:"
+    #     df -h
+
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: '3.11'
+
+    - name: Install cibuildwheel
+      run: |
+        python -m pip install cibuildwheel
+
+    - name: Set CUDA architecture list
+      run: |
+        ARCH_LIST=$(echo "$CUDA_ARCH_MAP" | jq -r '.["${{ matrix.cuda-version }}"]')
+        echo "TORCH_CUDA_ARCH_LIST=$ARCH_LIST" >> $GITHUB_ENV
+
+    - name: Determine package version suffix
+      run: |
+        BASE_VERSION=$(git describe --tags --dirty 2>/dev/null | sed 's/^v//' || echo "0.8.1")
+        SUFFIX="torch${{ matrix.torch-version }}.${{ matrix.torch-cuda }}"
+        echo "PACKAGE_VERSION=${BASE_VERSION}+${SUFFIX}" >> $GITHUB_ENV
+
+    - name: Build manylinux wheels
+      env:
+        CIBW_BEFORE_BUILD: |
+          pip install --upgrade pip setuptools wheel setuptools-scm
+          if [ "${{ matrix.torch-cuda }}" = "cpu" ]; then
+            pip install torch==${{ matrix.torch-version }}+cpu --index-url https://download.pytorch.org/whl/cpu
+          else
+            pip install torch==${{ matrix.torch-version }}+${{ matrix.torch-cuda }} --index-url https://download.pytorch.org/whl/${{ matrix.torch-cuda }}
+          fi
+          pip install numpy
+          export TORCH_HARMONICS_ENABLE_OPENMP=1
+          if [ "${{ matrix.cuda-version }}" != "none" ]; then
+            export CUDA_HOME=/usr/local/cuda-${{ matrix.cuda-version }}
+            export PATH=$CUDA_HOME/bin:$PATH
+            export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
+          fi
+        CIBW_BEFORE_TEST: |
+          echo "Installing PyTorch for testing..."
+          if [ "${{ matrix.torch-cuda }}" = "cpu" ]; then
+            pip install torch==${{ matrix.torch-version }}+cpu --index-url https://download.pytorch.org/whl/cpu
+          else
+            pip install torch==${{ matrix.torch-version }}+${{ matrix.torch-cuda }} --index-url https://download.pytorch.org/whl/${{ matrix.torch-cuda }}
+          fi
+          pip install numpy
+          if [ "${{ matrix.cuda-version }}" != "none" ]; then
+            export CUDA_HOME=/usr/local/cuda-${{ matrix.cuda-version }}
+            export PATH=$CUDA_HOME/bin:$PATH
+            export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
+          fi
+        CIBW_TEST_INSIDE: true
+        TORCH_CUDA_VERSION: ${{ matrix.cuda-version }}
+        TORCH_CUDA_ARCH_LIST: ${{ env.TORCH_CUDA_ARCH_LIST }}
+        CIBW_ENVIRONMENT: "CXXFLAGS='-D_GLIBCXX_USE_CXX11_ABI=1' SETUPTOOLS_SCM_PRETEND_VERSION=${{ env.PACKAGE_VERSION }} TORCH_CUDA_VERSION=${{ matrix.cuda-version }}"
+
+      run: |
+        cibuildwheel --platform linux
+
+    - name: Inspect built wheels and .so files
+      if: always()
+      run: |
+        echo "Listing wheelhouse contents:"
+        ls -lh wheelhouse/
+        echo "Checking .so symbols and ABI:"
+        for so in wheelhouse/*.so; do
+          echo "==== $so ===="
+          nm -D "$so" | grep exception_ptr || echo "No exception_ptr symbols found"
+          strings "$so" | grep ABI || echo "No ABI macro found in $so"
+        done
+
+    - name: Upload wheel artifacts
+      uses: actions/upload-artifact@v4
+      with:
+        name: wheel-torch-${{ matrix.torch-version }}-${{ matrix.torch-cuda }}
+        path: wheelhouse/*.whl
+
+    - name: Upload to TestPyPI (for testing)
+      if: github.event_name == 'workflow_dispatch' || github.event_name == 'pull_request'
+      env:
+        TWINE_USERNAME: __token__
+        TWINE_PASSWORD: ${{ secrets.TEST_PYPI_PASSWORD }}
+      run: |
+        python -m pip install twine
+        python -m twine upload --verbose --repository testpypi wheelhouse/*.whl
+
+    - name: Upload to PyPI (production)
+      if: startsWith(github.ref, 'refs/tags/')
+      env:
+        TWINE_USERNAME: __token__
+        TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
+      run: |
+        python -m pip install twine
+        python -m twine upload wheelhouse/*.whl

.gitignore

@@ -1,4 +1,36 @@
 *.DS_Store
 __pycache__
 *.so
-checkpoints
+checkpoints
+
+# Build artifacts
+build/
+dist/
+wheelhouse/
+*.egg-info/
+
+# Compiled extensions (keep source, ignore compiled)
+*.so
+*.pyd
+
+# CUDA compilation artifacts
+*.cubin
+*.fatbin
+*.ptx
+
+# Python cache
+*.pyc
+*.pyo
+*.pyd
+__pycache__/
+*.egg-info/
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+
+# OS
+.DS_Store
+Thumbs.db

README.md

@@ -75,35 +75,44 @@ torch-harmonics has been used to implement a variety of differentiable PDE solve
 
 
 ## Installation
-A simple installation can be directly done from PyPI:
 
+In general we recommend building the package yourself with the exact PyTorch build you rely on (see the instructions below) to avoid binary incompatibilities. Many features in torch-harmonics make use of accelerated CUDA kernels which need to match the CUDA packaged with your PyTorch installation. To do so, you may build torch-harmonics using
 ```bash
-pip install torch-harmonics
+export FORCE_CUDA_EXTENSION=1
+export TORCH_CUDA_ARCH_LIST="8.0 8.6 8.7 9.0 10.0 12.0+PTX"
+pip install --no-binary=torch-harmonics --no-build-isolation torch-harmonics
 ```
-If you are planning to use spherical convolutions, we recommend building the corresponding custom CUDA kernels. To enforce this, you can set the `FORCE_CUDA_EXTENSION` flag. You may also want to set appropriate architectures with the `TORCH_CUDA_ARCH_LIST` flag. Finally, make sure to disable build isolation via the `--no-build-isolation` flag to ensure that the custom kernels are built with the existing torch installation.
+The `--no-build-isolation` flag is necessary to ensure that torch-harmonics extensions are built with the correct PyTorch libraries. For container builds, we recommend setting the flags `FORCE_CUDA_EXTENSION` and `TORCH_CUDA_ARCH_LIST`, as the containers may run on systems with different GPUs than those available on the host system building the container.
+
+:warning: Please note that the custom CUDA extensions only support CUDA architectures >= 7.0.
+
+A simple installation using one of the prebuilt wheels on PyPI is also possible. However, it is important that it matches your local PyTorch and CUDA build. To do so, use the helper command below to print the exact `pip install` command for your environment:
+
 ```bash
-export FORCE_CUDA_EXTENSION=1
-export TORCH_CUDA_ARCH_LIST="7.0 7.2 7.5 8.0 8.6 8.7 9.0+PTX"
-pip install --no-build-isolation torch-harmonics
+python -c "import torch, os; parts=torch.__version__.split('+', 1); torch_tag=parts[0]; cuda_tag=parts[1] if len(parts) > 1 else 'cpu'; print(f'pip install torch-harmonics==0.8.1+torch{torch_tag}.{cuda_tag}')"
 ```
-:warning: Please note that the custom CUDA extensions currently only support CUDA architectures >= 7.0.
+This will generate a command of the form
+```bash
+pip install torch-harmonics==0.8.1+torch{torch_tag}.{cuda_tag}
+```
+specifying the exact wheel, matching both PyTorch installation and its CUDA version.
 
 If you want to actively develop torch-harmonics, we recommend building it in your environment from github:
 
 ```bash
 git clone [email protected]:NVIDIA/torch-harmonics.git
 cd torch-harmonics
-pip install -e .
+pip install --no-build-isolation -e .
 ```
-
+<!--
 Alternatively, use the Dockerfile to build your custom container after cloning:
 
 ```bash
 git clone [email protected]:NVIDIA/torch-harmonics.git
 cd torch-harmonics
 docker build . -t torch_harmonics
 docker run --gpus all -it --rm --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 torch_harmonics
-```
+``` -->
 
 ## More about torch-harmonics