Pre-Build Wheels

.github/workflows/build_wheels.yml

@@ -0,0 +1,207 @@
+name: Build and Deploy Manylinux Wheels
+
+# Trigger on version tags and manual dispatch
+on:
+  push:
+    tags:
+      - 'v*'
+  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: '{"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", "12.9": "8.0;8.6;8.9;9.0;10.0;12.0;+PTX", "13.0": "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:
+          # One build per CUDA version, using the oldest torch that supports it.
+          # This gives the widest torch runtime compatibility for each CUDA variant.
+          - torch-version: '2.6.0'
+            torch-cuda: 'cu126'
+            cuda-version: '12.6'
+            python-builds: 'cp39-* cp310-* cp311-* cp312-*'
+          - torch-version: '2.7.0'
+            torch-cuda: 'cu128'
+            cuda-version: '12.8'
+            python-builds: 'cp39-* cp310-* cp311-* cp312-*'
+          - torch-version: '2.8.0'
+            torch-cuda: 'cu129'
+            cuda-version: '12.9'
+            python-builds: 'cp310-* cp311-* cp312-*'
+          - torch-version: '2.9.1'
+            torch-cuda: 'cu130'
+            cuda-version: '13.0'
+            python-builds: 'cp310-* cp311-* cp312-*'
+          - torch-version: '2.11.0'
+            torch-cuda: 'cu130'
+            cuda-version: '13.0'
+            python-builds: 'cp310-* cp311-* cp312-*'
+          # CPU-only build based on latest PyTorch
+          - torch-version: '2.11.0'
+            torch-cuda: 'cpu'
+            cuda-version: 'none'
+            python-builds: 'cp310-* cp311-* cp312-*'
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - 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: Extract base version from pyproject.toml
+      run: |
+        BASE_VERSION=$(python -c "import tomllib; print(tomllib.load(open('pyproject.toml', 'rb'))['tool']['setuptools_scm']['fallback_version'])" 2>/dev/null || \
+                       python -c "import tomli; print(tomli.load(open('pyproject.toml', 'rb'))['tool']['setuptools_scm']['fallback_version'])")
+        echo "BASE_VERSION=$BASE_VERSION" >> $GITHUB_ENV
+        echo "Base version: $BASE_VERSION"
+
+    - name: Build manylinux wheels
+      env:
+        CIBW_BEFORE_BUILD: |
+          set -e
+          pip install --upgrade pip setuptools wheel setuptools-scm
+          pip install --no-cache-dir torch==${{ matrix.torch-version }} --index-url https://download.pytorch.org/whl/${{ matrix.torch-cuda }}
+          pip install numpy
+        CIBW_BEFORE_TEST: |
+          set -e
+          pip install torch==${{ matrix.torch-version }} --index-url https://download.pytorch.org/whl/${{ matrix.torch-cuda }}
+          pip install numpy
+        CIBW_BUILD: ${{ matrix.python-builds }}
+        TORCH_CUDA_VERSION: ${{ matrix.cuda-version }}
+        TORCH_CUDA_ARCH_LIST: ${{ env.TORCH_CUDA_ARCH_LIST }}
+        CIBW_ENVIRONMENT: "SETUPTOOLS_SCM_PRETEND_VERSION=${{ env.BASE_VERSION }}+torch${{ matrix.torch-version }}.${{ matrix.torch-cuda }} TORCH_CUDA_VERSION=${{ matrix.cuda-version }} TORCH_HARMONICS_ENABLE_OPENMP=1 TORCHINDUCTOR_COMPILE_THREADS=1"
+        CIBW_REPAIR_WHEEL_COMMAND_LINUX: |
+          set -e
+          TORCH_LIB=$(python -c "import torch, os; print(os.path.join(os.path.dirname(torch.__file__), 'lib'))")
+          EXCLUDES=$(find "$TORCH_LIB" -maxdepth 1 -name '*.so*' | xargs -I{} basename {} | sort -u | sed 's/^/--exclude /' | tr '\n' ' ')
+          LD_LIBRARY_PATH="$TORCH_LIB:${LD_LIBRARY_PATH:-}" auditwheel repair $EXCLUDES -w {dest_dir} {wheel}
+          for whl in {dest_dir}/*.whl; do
+            python -m wheel unpack "$whl" --dest /tmp/strip_tmp
+            rm -f "$whl"
+            find /tmp/strip_tmp -name '*.so' -exec strip --strip-debug {} \;
+            python -m wheel pack /tmp/strip_tmp/* --dest {dest_dir}
+            rm -rf /tmp/strip_tmp
+          done
+
+      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 whl in wheelhouse/*.whl; do
+          echo "==== $whl ===="
+          tmpdir=$(mktemp -d)
+          python -m zipfile -e "$whl" "$tmpdir"
+          for so in $(find "$tmpdir" -name '*.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"
+          done
+          rm -rf "$tmpdir"
+        done
+
+    - name: Upload wheel artifacts
+      uses: actions/upload-artifact@v4
+      with:
+        name: wheel-torch-${{ matrix.torch-version }}-${{ matrix.torch-cuda }}
+        path: wheelhouse/*.whl
+
+    - name: Repack wheels for PyPI (strip local version identifier)
+      if: matrix.torch-version == '2.11.0' && matrix.torch-cuda == 'cu130'
+      run: |
+        pip install wheel
+        mkdir -p pypi-wheelhouse
+        for whl in wheelhouse/*.whl; do
+          tmpdir=$(mktemp -d)
+          python -m wheel unpack "$whl" --dest "$tmpdir"
+          pkg_dir="$tmpdir/$(ls "$tmpdir")"
+          dist_info=$(ls "$pkg_dir" | grep '\.dist-info$')
+          base_ver=$(grep '^Version:' "$pkg_dir/$dist_info/METADATA" | sed 's/Version: //' | sed 's/+.*//')
+          pkg_name=$(echo "$dist_info" | sed 's/-[^-]*\.dist-info//')
+          sed -i "s/^Version:.*/Version: $base_ver/" "$pkg_dir/$dist_info/METADATA"
+          new_dist_info="${pkg_name}-${base_ver}.dist-info"
+          [ "$dist_info" != "$new_dist_info" ] && mv "$pkg_dir/$dist_info" "$pkg_dir/$new_dist_info"
+          new_pkg_dir="$tmpdir/${pkg_name}-${base_ver}"
+          [ "$pkg_dir" != "$new_pkg_dir" ] && mv "$pkg_dir" "$new_pkg_dir"
+          python -m wheel pack "$new_pkg_dir" --dest pypi-wheelhouse/
+          rm -rf "$tmpdir"
+        done
+
+    - name: Inspect repacked wheels and .so files
+      if: always() && (matrix.torch-version == '2.11.0' && matrix.torch-cuda == 'cu130')
+      run: |
+        pip install twine auditwheel
+
+        echo "=== twine check ==="
+        python -m twine check pypi-wheelhouse/*.whl
+
+        echo "=== auditwheel show ==="
+        for whl in wheelhouse/*.whl; do
+          echo "-- $whl --"
+          auditwheel show "$whl"
+        done
+
+        echo "=== Listing pypi-wheelhouse contents ==="
+        ls -lh pypi-wheelhouse/
+
+        echo "=== Checking .so symbols and ABI ==="
+        for whl in pypi-wheelhouse/*.whl; do
+          echo "==== $whl ===="
+          tmpdir=$(mktemp -d)
+          python -m zipfile -e "$whl" "$tmpdir"
+          for so in $(find "$tmpdir" -name '*.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"
+          done
+          rm -rf "$tmpdir"
+        done
+
+        echo "=== Install + import test ==="
+        python -m venv /tmp/test_venv
+        /tmp/test_venv/bin/pip install --quiet torch==${{ matrix.torch-version }} --index-url https://download.pytorch.org/whl/cpu
+        /tmp/test_venv/bin/pip install --quiet pypi-wheelhouse/*cp311*.whl
+        cd /tmp && /tmp/test_venv/bin/python -c "import torch_harmonics; print('import OK, version:', torch_harmonics.__version__)"
+
+    - name: Check wheel sizes (must be <= 100 MB)
+      run: |
+        WHL_DIR=$([ -d pypi-wheelhouse ] && echo pypi-wheelhouse || echo wheelhouse)
+        MAX_SIZE_BYTES=100000000
+        FAILED=0
+        for whl in "$WHL_DIR"/*.whl; do
+          SIZE_BYTES=$(stat -c%s "$whl")
+          SIZE_MB=$(awk "BEGIN {printf \"%.2f\", $SIZE_BYTES / 1000000}")
+          if [ "$SIZE_BYTES" -gt "$MAX_SIZE_BYTES" ]; then
+            echo "FAIL: $whl is ${SIZE_MB} MB (> 100 MB)"
+            FAILED=1
+          else
+            echo "OK:   $whl is ${SIZE_MB} MB"
+          fi
+        done
+        [ "$FAILED" -eq 0 ] || exit 1

.github/workflows/release_checklist.yml

@@ -0,0 +1,34 @@
+name: Release Checklist
+on:
+  release:
+    types: [published]
+
+permissions:
+  contents: read
+  issues: write
+
+jobs:
+  create-checklist:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Create release checklist issue
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          TAG: ${{ github.event.release.tag_name }}
+        run: |
+          gh issue create \
+            --title "Release Checklist: ${TAG}" \
+            --body "$(cat <<EOF
+          ## Release ${TAG} Checklist
+
+          - [ ] Changelog updated
+          - [ ] All tests passing
+          - [ ] Version number bumped in \`__init__.py\` and \`pyproject.toml\`
+          - [ ] Wheels build successful
+          - [ ] Release announcement drafted
+
+          ## Post-release
+
+          - Push wheels to PyPI and NVIDIA PyPI
+          EOF
+          )"

.github/workflows/tests.yml

@@ -29,12 +29,25 @@ jobs:
     - name: Test with pytest
       run: |
         python3 -m pip install pytest pytest-cov parameterized
-        python3 -m pytest --cov-report term --cov-config=.coveragerc --cov=torch_harmonics \
-          ./tests/test_sht.py \
-          ./tests/test_convolution.py \
-          ./tests/test_attention.py \
-          ./tests/test_cache.py \
-          ./tests/test_quadrature.py \
-          ./tests/test_resample.py \
-          ./tests/test_random_fields.py \
-          ./tests/test_spectral_convolution.py
+        python3 -m pytest --cov-report term --cov-report json --cov-config=.coveragerc --cov=torch_harmonics \
+          --ignore-glob="**/test_distributed_*" ./tests/
+
+    - name: Update coverage badge
+      if: github.ref == 'refs/heads/main'
+      env:
+        GIST_TOKEN: ${{ secrets.COVERAGE_GIST_TOKEN }}
+      run: |
+        COV=$(python3 -c "import json; print(json.load(open('coverage.json'))['totals']['percent_covered_display'])")
+        COLOR=$(python3 -c "
+        c = float('$COV')
+        print('brightgreen' if c >= 90 else 'green' if c >= 75 else 'yellow' if c >= 60 else 'red')
+        ")
+        PAYLOAD=$(python3 -c "
+        import json
+        print(json.dumps({'files': {'torch_harmonics_coverage.json': {'content': json.dumps({'schemaVersion': 1, 'label': 'coverage', 'message': '${COV}%', 'color': '${COLOR}'})}}}))
+        ")
+        curl -s -X PATCH \
+          -H "Authorization: token ${GIST_TOKEN}" \
+          -H "Accept: application/vnd.github+json" \
+          -d "$PAYLOAD" \
+          "https://api.github.com/gists/5d08af0f7f08ac865934a1929f198ffd"

.gitignore

@@ -1,4 +1,8 @@
 *.DS_Store
 __pycache__
 *.so
-checkpoints
+checkpoints
+
+# Kitmaker deployment secrets (never commit)
+scripts/kitmaker.env
+*.env

Changelog.md

@@ -2,6 +2,31 @@
 
 ## Versioning
 
+### v0.9.0
+
+* New CPU backend (OpenMP-accelerated) for both DISCO convolution and attention layers
+* Pre-compiled manylinux wheels for multiple PyTorch and CUDA versions, available on PyPI and pypi.nvidia.com
+* Revised truncation logic for the SHT: centralized in new `truncation.py` module, enforcing triangular truncation (`lmax = min(lmax, mmax)`) across all SHT classes. Note: truncation for equiangular/equidistant grids changed from `nlat` to `(nlat+1)//2`
+* SHT performance improvements: contraction dimensions are now transposed to be stride-1 before einsum, and real/imaginary parts are split into separate contiguous tensors
+* New `fft.py` wrapper module with proper Hermitian symmetry enforcement in `irfft` and explicit mode truncation in `rfft`
+* Full PyTorch 2 custom operator compatibility for DISCO and attention layers using `torch.library` registration, enabling `torch.compile` and `torch.export`
+* Restructured DISCO convolution and attention code into proper subpackages (`torch_harmonics/disco/`, `torch_harmonics/attention/`)
+* Added double precision support for DISCO convolution
+* Fixed Schmidt normalization for derivatives of associated Legendre polynomials
+* Fixed up/downsampling in attention layers when input and output shapes differ
+* Fixed `GaussianRandomFieldS2` to use `isht.lmax`/`isht.mmax` for compatibility with revised truncation logic
+* Distributed module: added shape verification for transpose and gather operations, controllable via `TORCH_HARMONICS_DISTRIBUTED_DEBUG`
+* Distributed module: fixed `finalize()` bug where process group was not properly destroyed
+* Query functions `torch_harmonics.disco.optimized_kernels_is_available` and `torch_harmonics.attention.optimized_kernels_is_available` for checking optimized layer availability
+* Quadrature helper functions `precompute_latitudes` and `precompute_longitudes` are now public API
+* added new tests:
+    * Comprehensive SHT test suite now covering vector SHT, Schmidt normalization, batch dimensions, and multiple grid types
+    * New test suites for `SpectralConvS2`, `QuadratureS2`, `GaussianRandomFieldS2`, and `ResampleS2`
+     Enhanced DISCO convolution tests covering different input/output channel counts and double precision
+    * Enhanced attention tests with up/downsampling and `opcheck` integration
+    * New distributed tests for primitives, quadrature, and spectral convolution
+    * Shared test utilities module (`testutils.py`)
+
 ### v0.8.2
 * Adding Driscoll-Healy (spectral) convolutions
 * Adding QuadratureS2 method which allows to integrate a spherical field over one of the supported grids