build-windows-wheel-test.yml

name: Build Cross-Platform Wheels

on:
  push:
    tags:
      - 'v*'
  workflow_dispatch:
    inputs:
      python_versions:
        description: 'Python versions to build (comma-separated)'
        required: false
        default: '3.11'
      cuda_versions:
        description: 'CUDA versions to build (comma-separated)'
        required: false
        default: '12.6'
      pytorch_version:
        description: 'PyTorch version to install (e.g., "2.1.0", "2.4.0", or "latest" for newest)'
        required: false
        default: '2.6.0'
      platforms:
        description: 'Platforms to build for (comma-separated: windows,linux)'
        required: false
        default: 'windows'
      build_latest_only:
        description: 'Build only latest PyTorch version for each CUDA version'
        required: false
        default: 'false'

jobs:
  generate-matrix:
    runs-on: ubuntu-latest
    outputs:
      matrix: ${{ steps.set-matrix.outputs.matrix }}
    steps:
    - name: Generate build matrix
      id: set-matrix
      run: |
        # Get input parameters or use defaults
        python_versions="${{ github.event.inputs.python_versions }}"
        cuda_versions="${{ github.event.inputs.cuda_versions }}"
        platforms="${{ github.event.inputs.platforms }}"

        # Convert comma-separated strings to arrays
        IFS=',' read -ra PYTHON_ARRAY <<< "$python_versions"
        IFS=',' read -ra CUDA_ARRAY <<< "$cuda_versions"
        IFS=',' read -ra PLATFORM_ARRAY <<< "$platforms"

        # Map platforms to OS
        os_array=()
        for platform in "${PLATFORM_ARRAY[@]}"; do
          case "$platform" in
            "windows") os_array+=("windows-2022") ;;
            "linux") os_array+=("ubuntu-22.04") ;;
          esac
        done

        # Build matrix JSON
        matrix_json="{\"include\":["
        first=true

        for os in "${os_array[@]}"; do
          for python in "${PYTHON_ARRAY[@]}"; do
            for cuda in "${CUDA_ARRAY[@]}"; do
              # Skip Python 3.12 with CUDA 11.8 for compatibility
              if [[ "$python" == "3.12" && "$cuda" == "11.8" ]]; then
                continue
              fi

              if [ "$first" = true ]; then
                first=false
              else
                matrix_json+=","
              fi

              matrix_json+="{\"os\":\"$os\",\"python-version\":\"$python\",\"cuda-version\":\"$cuda\"}"
            done
          done
        done

        matrix_json+="]}"

        echo "Generated matrix: $matrix_json"
        echo "matrix=$matrix_json" >> $GITHUB_OUTPUT

  build-wheels:
    needs: generate-matrix
    strategy:
      matrix: ${{ fromJson(needs.generate-matrix.outputs.matrix) }}
      fail-fast: false

    runs-on: ${{ matrix.os }}

    steps:
    - name: Install Linux dependencies (for system sparsehash)
      if: runner.os == 'Linux'
      run: |
        sudo apt-get update
        sudo apt-get install -y build-essential libsparsehash-dev

    # REMOVED: "Configure sparsehash for MSVC" step
    # REMOVED: "Set MSVC compiler options" step (flags now in setup.py)

    - name: Checkout repository
      uses: actions/checkout@v4
      with:
        submodules: recursive # Good practice

    - name: Set up Python ${{ matrix.python-version }}
      uses: actions/setup-python@v5
      with:
        python-version: ${{ matrix.python-version }}
        cache: 'pip'
        cache-dependency-path: |
          requirements.txt
          setup.py

    - name: Cache CUDA Toolkit (using NVIDIA action path)
      if: runner.os == 'Windows'
      id: cache-cuda-toolkit
      uses: actions/cache@v4
      with:
        # The exact path might vary slightly based on the action, check its documentation
        # or observe where it installs on the first run.
        path: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v${{ matrix.cuda-version }}
        key: ${{ runner.os }}-nvidia-cuda-${{ matrix.cuda-version }}

    - name: Install CUDA Toolkit ${{ matrix.cuda-version }}
      if: runner.os == 'Windows' && steps.cache-cuda-toolkit.outputs.cache-hit != 'true'
      id: cuda-toolkit
      shell: pwsh
      run: scripts/setup_cuda.ps1
      env:
        INPUT_CUDA_VERSION: ${{ matrix.cuda-version }}.0

    # - name: Install CUDA Toolkit ${{ matrix.cuda-version }}
    #   if: runner.os == 'Windows' && steps.cache-cuda-toolkit.outputs.cache-hit != 'true'
    #   uses: Jimver/cuda-toolkit@v0.2.23
    #   with:
    #     cuda: ${{ matrix.cuda-version }}.0
    #     use-github-cache: false
    #     method: 'network'
    #     sub-packages: '["nvcc", "visual_studio_integration", "cudart", "cublas", "cublas_dev", "cusparse", "cusparse_dev", "cusolver", "cusolver_dev", "cupti", "cufft", "cufft_dev", "curand", "curand_dev", "nvrtc_dev", "nvrtc"]'

    - name: Setup Visual Studio environment (Windows)
      if: runner.os == 'Windows'
      uses: ilammy/msvc-dev-cmd@v1
      with:
        arch: amd64

    # REMOVED: "Install sparsehash (Windows)" step
    # REMOVED: "Set sparsehash environment (Windows)" step

    - name: Install Python dependencies (build tools)
      run: |
        python -m pip install --upgrade pip
        pip install wheel setuptools ninja # Ninja for faster C++ extension builds

    - name: Install PyTorch (Windows)
      if: runner.os == 'Windows'
      shell: pwsh
      run: |
        $cuda_version_input = "${{ matrix.cuda-version }}"
        $pytorch_version_input = "${{ github.event.inputs.pytorch_version || 'latest' }}"
        $cuda_short_tag = switch ($cuda_version_input) {
          "11.8" { "cu118" }
          "12.1" { "cu121" }
          "12.4" { "cu124" }
          "12.6" { "cu126" }
          default { Write-Error "Unsupported CUDA version for PyTorch wheels: $cuda_version_input"; exit 1 }
        }
        $pytorch_index_url = "https://download.pytorch.org/whl/$cuda_short_tag"
        Write-Host "PyTorch Index URL: $pytorch_index_url"

        if ($pytorch_version_input -eq "latest") {
          Write-Host "Installing latest PyTorch for CUDA $cuda_version_input ($cuda_short_tag)"
          pip install torch torchvision torchaudio --index-url $pytorch_index_url
        } else {
          Write-Host "Installing PyTorch $pytorch_version_input for CUDA $cuda_version_input ($cuda_short_tag)"
          $torch_package_version = "$pytorch_version_input+$cuda_short_tag"
          if ($pytorch_version_input -eq "2.6.0") {
            pip install "torch==$torch_package_version" "torchvision==0.21.0+$cuda_short_tag" "torchaudio==2.6.0+$cuda_short_tag" --index-url $pytorch_index_url
          } else {
            pip install "torch==$torch_package_version" torchvision torchaudio --index-url $pytorch_index_url
          }
        }
        python -c "import torch; print(f'PyTorch version: {torch.__version__}'); print(f'Torch CUDA version: {torch.version.cuda}')"

    - name: Install PyTorch (Linux)
      if: runner.os == 'Linux'
      run: |
        cuda_version_input="${{ matrix.cuda-version }}"
        pytorch_version_input="${{ github.event.inputs.pytorch_version || 'latest' }}"
        cuda_short_tag=""
        case "$cuda_version_input" in
          "11.8") cuda_short_tag="cu118" ;;
          "12.1") cuda_short_tag="cu121" ;;
          "12.4") cuda_short_tag="cu124" ;;
          "12.6") cuda_short_tag="cu126" ;;
          *) echo "Error: Unsupported CUDA version for PyTorch wheels: $cuda_version_input"; exit 1 ;;
        esac
        pytorch_index_url="https://download.pytorch.org/whl/$cuda_short_tag"
        echo "PyTorch Index URL: $pytorch_index_url"

        if [ "$pytorch_version_input" = "latest" ]; then
          echo "Installing latest PyTorch for CUDA $cuda_version_input ($cuda_short_tag)"
          pip install torch torchvision torchaudio --index-url "$pytorch_index_url"
        else
          echo "Installing PyTorch $pytorch_version_input for CUDA $cuda_version_input ($cuda_short_tag)"
          torch_package_version="$pytorch_version_input+$cuda_short_tag"
          if [ "$pytorch_version_input" = "2.6.0" ]; then
            pip install "torch==$torch_package_version" "torchvision==0.21.0+$cuda_short_tag" "torchaudio==2.6.0+$cuda_short_tag" --index-url "$pytorch_index_url"
          else
            pip install "torch==$torch_package_version" torchvision torchaudio --index-url "$pytorch_index_url"
          fi
        fi
        python -c "import torch; print(f'PyTorch version: {torch.__version__}'); print(f'Torch CUDA version: {torch.version.cuda}')"
       
    - name: Set up CUDA environment (Windows)
      if: runner.os == 'Windows'
      shell: pwsh
      run: |
        $cuda_version_matrix = "${{ matrix.cuda-version }}"
        $cuda_root_dir = "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v$cuda_version_matrix"
        
        Write-Host "Initial environment (before CUDA setup):"
        Get-ChildItem Env: | Where-Object {$_.Name -in @("INCLUDE", "LIB", "PATH")} | Format-Table -AutoSize
        
        $cuda_include_paths = @(
            "$cuda_root_dir\include",
            "$cuda_root_dir\extras\CUPTI\include",
            "$cuda_root_dir\nvvm\include"
        ) -join ";"
        $cuda_lib_paths = @(
            "$cuda_root_dir\lib\x64",
            "$cuda_root_dir\extras\CUPTI\lib64"
        ) -join ";"

        echo "CUDA_PATH=$cuda_root_dir" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append
        echo "CUDA_HOME=$cuda_root_dir" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append
        echo "CUDA_TOOLKIT_ROOT_DIR=$cuda_root_dir" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append
        
        echo "INCLUDE=$cuda_include_paths;$($env:INCLUDE)" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append
        echo "LIB=$cuda_lib_paths;$($env:LIB)" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append
        echo "PATH=$cuda_root_dir\bin;$($env:PATH)" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append
        
        Write-Host "Verifying CUDA paths..."
        if (-not (Test-Path "$cuda_root_dir\include")) { Write-Error "❌ Missing: $cuda_root_dir\include"; exit 1 } else { Write-Host "✅ Found: $cuda_root_dir\include" }
        if (-not (Test-Path "$cuda_root_dir\lib\x64")) { Write-Error "❌ Missing: $cuda_root_dir\lib\x64"; exit 1 } else { Write-Host "✅ Found: $cuda_root_dir\lib\x64" }
        Write-Host "NVCC version:"
        nvcc --version

    - name: Set build environment (Windows)
      if: runner.os == 'Windows'
      shell: pwsh
      run: |
        echo "DISTUTILS_USE_SDK=1" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append
        echo "MSSdk=1" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append
        echo "FORCE_CUDA=1" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append
        echo "TORCH_CUDA_ARCH_LIST=7.5;8.0;8.6;8.9;9.0" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append
        echo "MAX_JOBS=$(Get-CimInstance Win32_Processor | Measure-Object -Property NumberOfLogicalProcessors -Sum).Sum" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append
        echo "_CRT_SECURE_NO_WARNINGS=1" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append
        echo "_SILENCE_TR1_NAMESPACE_DEPRECATION_WARNING=1" | Out-File -FilePath $env:GITHUB_ENV -Encoding utf8 -Append

    - name: Set build environment (Linux)
      if: runner.os == 'Linux'
      run: |
        echo "FORCE_CUDA=1" >> $GITHUB_ENV
        echo "TORCH_CUDA_ARCH_LIST=7.5;8.0;8.6;8.9;9.0" >> $GITHUB_ENV
        echo "MAX_JOBS=$(nproc)" >> $GITHUB_ENV
        # CXXFLAGS/CFLAGS from setup.py will be used for the extension

    - name: Build wheel
      shell: pwsh # Using pwsh for consistency, bash for Linux specific commands if any
      run: |
        Write-Host "--- Build Environment Variables ---"
        Get-ChildItem Env: | Sort-Object Name | Format-Table Name, Value -AutoSize
        Write-Host "-----------------------------------"
        Write-Host "Python version for build: $(python --version)"
        Write-Host "PyTorch version for build: $(python -c "import torch; print(torch.__version__)")"
        Write-Host "CUDA_PATH: $env:CUDA_PATH"
        Write-Host "CUDA_HOME: $env:CUDA_HOME"
        
        if ($env:RUNNER_OS -eq 'Windows') {
            if (Test-Path "$env:CUDA_PATH\lib\x64\cusparse.lib") {
                Write-Host "✅ Windows: CUDA cusparse.lib found"
            } else {
                Write-Error "❌ Windows: CUDA cusparse.lib not found at $env:CUDA_PATH\lib\x64\cusparse.lib. Full include path: $env:INCLUDE. Full lib path: $env:LIB"
                exit 1
            }
        }
        
        python setup.py bdist_wheel --verbose

    - name: Test wheel installation (Windows)
      if: runner.os == 'Windows'
      shell: pwsh
      run: |
        $wheel = Get-ChildItem -Path "dist" -Filter "*.whl" | Select-Object -First 1
        if (-not $wheel) { Write-Error "No wheel found in dist directory."; exit 1 }
        Write-Host "Installing wheel: $($wheel.FullName)"
        pip install $wheel.FullName
        python -c "import torchsparse; print(f'TorchSparse version: {torchsparse.__version__}')"
        python -c "import torch; import torchsparse; print('Basic import test passed for torchsparse on Windows')"

    - name: Test wheel installation (Linux)
      if: runner.os == 'Linux'
      run: |
        wheel_path=$(find dist -name "*.whl" | head -n 1)
        if [ -z "$wheel_path" ]; then echo "No wheel found in dist directory."; exit 1; fi
        echo "Installing wheel: $wheel_path"
        pip install "$wheel_path"
        python -c "import torchsparse; print(f'TorchSparse version: {torchsparse.__version__}')"
        python -c "import torch; import torchsparse; print('Basic import test passed for torchsparse on Linux')"

    - name: Upload wheel artifacts
      uses: actions/upload-artifact@v4
      with:
        name: wheels-${{ runner.os }}-py${{ matrix.python-version }}-cuda${{ matrix.cuda-version }}
        path: dist/*.whl

  create-release:
    needs: build-wheels
    runs-on: ubuntu-latest
    if: startsWith(github.ref, 'refs/tags/')
    steps:
    - name: Checkout repository
      uses: actions/checkout@v4
    - name: Download all wheel artifacts
      uses: actions/download-artifact@v4
      with:
        path: wheels
    - name: Organize wheels
      run: |
        mkdir -p release_wheels
        find wheels -name "*.whl" -exec cp -v {} release_wheels/ \;
        echo "Collected wheels in release_wheels/:"
        ls -lR release_wheels/
    - name: Create release notes
      id: create_release_notes
      run: |
        cat > release_notes.md <<EOF
        # TorchSparse v${{ github.ref_name }}

        Official release for TorchSparse version ${{ github.ref_name }}.

        ## Wheels
        Attached wheels provide support for various Python and CUDA versions on Linux and Windows.
        Please check the file names for platform, Python version, and CUDA compatibility.
        EOF
    - name: Create GitHub Release
      uses: softprops/action-gh-release@v1
      with:
        files: release_wheels/*.whl
        body_path: release_notes.md
        tag_name: ${{ github.ref_name }}
        name: TorchSparse ${{ github.ref_name }}
        draft: false
        prerelease: false
      env:
        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

  test-wheels:
    needs: [build-wheels] # generate-matrix not needed if matrix is static for test
    strategy:
      matrix:
        os: [windows-2022, ubuntu-22.04]
        python-version: ['3.11']
        cuda-version: ['12.6']
    runs-on: ${{ matrix.os }}
    steps:
    - name: Set up Python ${{ matrix.python-version }}
      uses: actions/setup-python@v5
      with:
        python-version: ${{ matrix.python-version }}
        cache: 'pip'

    - name: Install CUDA Toolkit (Test) - Optional
      if: runner.os == 'Windows' # Example: only install full toolkit on Windows for test if needed
      uses: Jimver/cuda-toolkit@v0.2.23
      with:
        cuda: ${{ matrix.cuda-version }}.0
        method: 'network'
        # sub-packages: '["cudart"]' # Minimal for runtime

    - name: Download wheel artifact for testing
      uses: actions/download-artifact@v4
      with:
        name: wheels-${{ runner.os }}-py${{ matrix.python-version }}-cuda${{ matrix.cuda-version }}
        path: downloaded_wheel

    - name: Install PyTorch for testing
      shell: bash
      run: |
        PYTORCH_CUDA_TAG=""
        TARGET_PYTORCH_VERSION=""
        TARGET_TORCHVISION_VERSION=""
        TARGET_TORCHAUDIO_VERSION=""

        if [ "${{ matrix.cuda-version }}" = "12.6" ]; then
          PYTORCH_CUDA_TAG="cu126"
          TARGET_PYTORCH_VERSION="2.6.0" # Align with build
          TARGET_TORCHVISION_VERSION="0.21.0"
          TARGET_TORCHAUDIO_VERSION="2.6.0"
        elif [ "${{ matrix.cuda-version }}" = "12.1" ]; then
          PYTORCH_CUDA_TAG="cu121"
          TARGET_PYTORCH_VERSION="2.1.0" # Example
          TARGET_TORCHVISION_VERSION="0.16.0"
          TARGET_TORCHAUDIO_VERSION="2.1.0"
        else
          echo "Unsupported CUDA version for PyTorch test install: ${{ matrix.cuda-version }}"
          exit 1
        fi
        
        echo "Installing PyTorch $TARGET_PYTORCH_VERSION+$PYTORCH_CUDA_TAG, torchvision $TARGET_TORCHVISION_VERSION+$PYTORCH_CUDA_TAG, torchaudio $TARGET_TORCHAUDIO_VERSION+$PYTORCH_CUDA_TAG"
        pip install "torch==$TARGET_PYTORCH_VERSION+$PYTORCH_CUDA_TAG" \
                    "torchvision==$TARGET_TORCHVISION_VERSION+$PYTORCH_CUDA_TAG" \
                    "torchaudio==$TARGET_TORCHAUDIO_VERSION+$PYTORCH_CUDA_TAG" \
                    --index-url https://download.pytorch.org/whl/$PYTORCH_CUDA_TAG

    - name: Test wheel installation and functionality
      shell: bash
      run: |
        WHEEL_FILE=$(find downloaded_wheel -name "*.whl" | head -n 1)
        if [ -z "$WHEEL_FILE" ]; then echo "No wheel found in downloaded_wheel directory."; exit 1; fi
        echo "Installing downloaded wheel: $WHEEL_FILE"
        pip install "$WHEEL_FILE" --force-reinstall # Ensure it installs over any other version

        echo "Running Python test script..."
        python -c "
        import torch
        import torchsparse
        import numpy as np # Often used, though not strictly in this basic test
        import os # For platform info

        print(f'--- Test Environment ---')
        print(f'TorchSparse version: {torchsparse.__version__}')
        print(f'PyTorch version: {torch.__version__}')
        print(f'PyTorch CUDA available: {torch.cuda.is_available()}')
        if torch.cuda.is_available():
            print(f'PyTorch CUDA version: {torch.version.cuda}')
            print(f'CUDA Device Name: {torch.cuda.get_device_name(0)}')
        print(f'Platform: {os.name} ({os.uname() if hasattr(os, \"uname\") else \"Windows\"})')
        print(f'------------------------')

        # Basic functionality test
        try:
            print('Creating sparse tensor data...')
            # Example: 3 points in a 4D space (e.g., batch, x, y, z)
            coords_data = [[0, 0, 0, 0], [0, 0, 1, 1], [0, 1, 2, 3]]
            # Features for each point (2 features per point)
            feats_data = [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]
            
            coords = torch.tensor(coords_data, dtype=torch.int)
            feats = torch.tensor(feats_data, dtype=torch.float)

            if torch.cuda.is_available():
                print('Moving tensors to CUDA device 0...')
                coords = coords.to('cuda:0')
                feats = feats.to('cuda:0')
            else:
                print('CUDA not available, using CPU for test.')

            print(f'Coords shape: {coords.shape}, Feats shape: {feats.shape}, Device: {coords.device}')
            
            # Create SparseTensor
            # Stride can be a single int or a list/tuple matching spatial dimensions
            # For 4D coords (B, Z, Y, X), spatial_dims is usually Z, Y, X.
            # If your data is (B, X, Y, Z), adjust accordingly.
            # Let's assume default stride and let SparseTensor infer shape.
            sparse_tensor = torchsparse.SparseTensor(coords=coords, feats=feats)
            
            print(f'Sparse tensor created successfully.')
            print(f'  Shape: {sparse_tensor.shape}')
            print(f'  Device: {sparse_tensor.device}')
            print(f'  Stride: {sparse_tensor.stride}')
            print(f'  Spatial Strides: {sparse_tensor.spatial_strides}') # Added this
            print(f'  Number of points: {len(sparse_tensor)}')

            # Test a simple operation if available, e.g., to_dense (can be memory intensive)
            # dense_tensor = sparse_tensor.to_dense()
            # print(f'  Dense tensor shape: {dense_tensor.shape}')

            print('✅ Basic SparseTensor creation and properties test passed!')
        except Exception as e:
            print(f'❌ Test failed: {e}')
            import traceback
            traceback.print_exc()
            exit(1)
        "