python-package.yml

name: Python package

on:
  push:
    branches: [ main ]
  pull_request:
    paths:
      - '.github/workflows/python-package.yml'
      - 'bitsandbytes/**'
      - 'csrc/**'
      - 'include/**'
      - 'tests/**'
      - 'CMakeLists.txt'
      - 'requirements*.txt'
      - 'setup.py'
      - 'pyproject.toml'
      - 'pytest.ini'
      - '**/*.md'
  release:
    branches: [ main ]
    types: [ published ]

concurrency:
  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
  cancel-in-progress: true

jobs:

  ##
  # This job matrix builds the non-CUDA versions of the libraries for all supported platforms.
  ##
  build-shared-libs:
    strategy:
      # Set fail-fast to false to ensure that feedback is delivered for all matrix combinations. Consider changing this to true when your workflow is stable.
      fail-fast: false

      matrix:
        os: [ubuntu-latest, macos-latest, windows-latest]
        arch: [x86_64, aarch64]
        build_type: [Release]
        exclude:
          - os: windows-latest # This probably requires arm64 Windows agents
            arch: aarch64
    runs-on: ${{ matrix.os }} # One day, we could run them on native agents. Azure supports this now but it's planned only for Q3 2023 for hosted agents
    steps:
      # Check out code
    - uses: actions/checkout@v4
    - name: Allow cross-compile on aarch64
      if: ${{ startsWith(matrix.os, 'ubuntu') && matrix.arch == 'aarch64' }}
      run: |
        # Allow cross-compile on aarch64
        sudo apt-get install -y g++-aarch64-linux-gnu binutils-aarch64-linux-gnu

    - name: Setup cmake
      uses: jwlawson/actions-setup-cmake@v1.14
      with:
        cmake-version: '3.26.x'
    - name: Setup MSVC
      if: startsWith(matrix.os, 'windows')
      #uses: microsoft/setup-msbuild@v1.1 # to use msbuild
      uses: ilammy/msvc-dev-cmd@v1.13.0 # to use cl
      # Check out dependencies code

    - name: Prep Compilers
      shell: bash -el {0}
      run: |
        python3 -m pip install cmake==3.27.9 ninja
        if [[ "${{ matrix.os }}" = windows-* ]]; then
          echo CXX_COMPILER=-DCMAKE_CXX_COMPILER=cl >> "$GITHUB_ENV"
        elif [[ "${{ matrix.os }}" = ubuntu-* ]] && [[ "${{ matrix.arch }}" = "aarch64" ]]; then
          echo CXX_COMPILER=-DCMAKE_CXX_COMPILER=aarch64-linux-gnu-g++ >> "$GITHUB_ENV"
        else
          echo CXX_COMPILER=-DCMAKE_CXX_COMPILER=g++ >> "$GITHUB_ENV"
        fi

        if [[ "${{ matrix.os }}" = macos-* ]] && [[ "${{ matrix.arch }}" = "aarch64" ]]; then
          echo DCMAKE_OSX_ARCHITECTURES=-DCMAKE_OSX_ARCHITECTURES=arm64 >> "$GITHUB_ENV"
        fi

    - name: Build CPU
      shell: bash -el {0}
      run: |
        cmake -B build \
        -G Ninja ${{ env.DCMAKE_OSX_ARCHITECTURES }} \
        ${{ env.CXX_COMPILER }} \
        -DCMAKE_BUILD_TYPE=${{ matrix.build_type }} \
        -DCOMPUTE_BACKEND=cpu \
        -S .

        cmake --build build --config ${{ matrix.build_type }}

        mkdir -p output/${{ matrix.os }}/${{ matrix.arch }}
        ls -l bitsandbytes
        echo " check lib files..."
        file bitsandbytes/libbitsandbytes*.*
        ( shopt -s nullglob && cp bitsandbytes/*.{so,dylib,dll} output/${{ matrix.os }}/${{ matrix.arch }}/ )
    - name: Upload build artifact
      uses: actions/upload-artifact@v4
      with:
        name: shared_library_${{ matrix.os }}_${{ matrix.arch }}
        path: output/*
  ##
  # This job matrix builds the CUDA versions of the libraries for platforms that support CUDA (Linux x64/aarch64 + Windows x64)
  ##
  build-shared-libs-cuda:
    strategy:
      # Set fail-fast to false to ensure that feedback is delivered for all matrix combinations. Consider changing this to true when your workflow is stable.
      fail-fast: false

      matrix:
        os: [ubuntu-latest, windows-latest]
        arch: ${{ startsWith(github.ref, 'refs/tags') && fromJson('["x86_64","aarch64"]') || fromJson('["x86_64"]') }}
        cuda-version: ${{ startsWith(github.ref, 'refs/tags') && fromJson('["11.8.0","12.1.1"]') || fromJson('["12.1.1"]') }}
        build_type: [Release]
        exclude:
          - os: windows-latest # This probably requires arm64 Windows agents
            arch: aarch64
    runs-on: ${{ matrix.os }} # One day, we could run them on native agents. Azure supports this now but it's planned only for Q3 2023 for hosted agents
    steps:
      # Check out code
    - uses: actions/checkout@v4
      # Linux: We use Docker to build cross platform Cuda (aarch64 is built in emulation)
    - name: Set up Docker multiarch
      if: startsWith(matrix.os, 'ubuntu')
      uses: docker/setup-qemu-action@v2
      # On Linux we use CMake within Docker
    - name: Setup cmake
      if: ${{ !startsWith(matrix.os, 'ubuntu') }}
      uses: jwlawson/actions-setup-cmake@v1.14
      with:
        cmake-version: '3.26.x'
      # Windows: We install Cuda on the agent (slow)
      # it is too slow and sometime break
    - uses: Jimver/cuda-toolkit@v0.2.14
      if: startsWith(matrix.os, 'x') # disabled
      id: cuda-toolkit
      with:
        cuda: ${{ matrix.cuda_version }}
        method: 'local'
        # sub-packages: '["nvcc","cudart","nvrtc_dev","cublas_dev","cusparse_dev","visual_studio_integration"]'
    - name: Set up Python 3.10
      uses: actions/setup-python@v5
      with:
        python-version: "3.10"

    - name: Setup MSVC
      if: startsWith(matrix.os, 'windows')
      #uses: microsoft/setup-msbuild@v1.1 # to use msbuild
      uses: ilammy/msvc-dev-cmd@v1.13.0 # to use cl
      # Check out dependencies code
    - name: Setup Mambaforge
      if: startsWith(matrix.os, 'windows')
      uses: conda-incubator/setup-miniconda@v3.0.1
      with:
        miniforge-variant: Mambaforge
        miniforge-version: latest
        activate-environment: bnb-env
        use-mamba: true

    - uses: conda-incubator/setup-miniconda@v3.0.1
      if: startsWith(matrix.os, 'windows')
      with:
        auto-update-conda: true
        activate-environment: bnb-env
        environment-file: environment-bnb.yml
        use-only-tar-bz2: false
        auto-activate-base: true
        python-version: "3.10"
        mamba-version: "*"

    - name: Setup CUDA
      if: startsWith(matrix.os, 'windows')
      shell: bash -el {0}
      run: |
        addon=""
        cuda_version=${{ matrix.cuda-version }}
        [[ "$cuda_version" = 12.1.* ]] && [[ "${{ matrix.os }}" = ubuntu-* ]] && addon="cuda-cudart-static cuda-nvrtc"
        [[ "$cuda_version" = 11.8.* ]] && [[ "${{ matrix.os }}" = windows-* ]] && addon="cuda-nvrtc"
        [[ "$cuda_version" = 11.8.* ]] && cuda__version="11.8"
        [[ "$cuda_version" = 12.1.* ]] && cuda__version="12.1"

        conda install pytorch-cuda=$cuda__version -c pytorch # it's dependency not correctly resolved sometime
        conda install cuda-python=$cuda__version cuda-libraries-dev cuda-nvcc cuda-nvtx cuda-cupti cuda-cudart cuda-cudart-dev cuda-runtime cuda-libraries $addon -c "nvidia/label/cuda-$cuda_version"

        CUDA_HOME="${{ env.CONDA }}/envs/bnb-env"
        echo CUDA_HOME=$CUDA_HOME >> "$GITHUB_ENV"
        echo CUDA_PATH=$CUDA_HOME >> "$GITHUB_ENV"

        if [[ "${{ matrix.os }}" = windows-* ]]; then
          echo CXX_COMPILER=-DCMAKE_CXX_COMPILER=cl >> "$GITHUB_ENV"
          # without -DCMAKE_CUDA_COMPILER=nvcc, cmake config always fail for cuda-11.8
          echo DCMAKE_CUDA_COMPILER=-DCMAKE_CUDA_COMPILER=nvcc >> "$GITHUB_ENV"
        else
          echo CXX_COMPILER=-DCMAKE_CXX_COMPILER=g++ >> "$GITHUB_ENV"
        fi

        if [[ "${{ matrix.os }}" = macos-* ]] && [[ "${{ matrix.arch }}" = "aarch64" ]]; then
          echo DCMAKE_OSX_ARCHITECTURES=-DCMAKE_OSX_ARCHITECTURES=arm64 >> "$GITHUB_ENV"
        fi

        nvcc --version

    - name: Select COMPUTE_CAPABILITY
      shell: bash -el {0}
      run: |
        # set COMPUTE_CAPABILITY
        COMPUTE_CAPABILITY="61;75;86;89"
        [[ "${{ env.GITHUB_REF }}" = refs/tags/* ]] && COMPUTE_CAPABILITY="50;52;60;61;62;70;72;75;80;86;87;89;90"
        echo "COMPUTE_CAPABILITY=$COMPUTE_CAPABILITY" >> "$GITHUB_ENV"

    - name: Update environment
      if: startsWith(matrix.os, 'windows')
      run: mamba env update -n bnb-env -f environment-bnb.yml

    - name: Prep build
      if: startsWith(matrix.os, 'windows')
      run: python -m pip install cmake==3.27.9 ninja

    - name: Build CUDA
      if: startsWith(matrix.os, 'windows')
      shell: bash -el {0}
      run: |
        cmake -B build \
        -G Ninja ${{ env.DCMAKE_CUDA_COMPILER }} \
        ${{ env.CXX_COMPILER }} \
        -DCMAKE_BUILD_TYPE=${{ matrix.build_type }} \
        -DCOMPUTE_CAPABILITY="${{ env.COMPUTE_CAPABILITY }}" \
        -DCOMPUTE_BACKEND=cuda \
        -S .

        cmake --build build --config ${{ matrix.build_type }}

    - name: Build CUDA (docker)
      if: startsWith(matrix.os, 'ubuntu')
      uses: addnab/docker-run-action@v3
      with:
        image: ${{ format('nvidia/cuda:{0}-{1}', matrix.cuda-version, 'devel-ubuntu22.04') }}
        options: --platform linux/${{ matrix.arch }} -w /src -v ${{ github.workspace }}:/src -e "COMPUTE_CAPABILITY"
        run: |
          apt-get update
          DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends cmake python3 python3-pip
          # install cmake, ninja to prepare cmake
          python3 -m pip install cmake==3.27.9 ninja

          for NO_CUBLASLT in OFF ON; do
            cmake -B build \
            -G Ninja \
            -DCMAKE_BUILD_TYPE=${{ matrix.build_type }} \
            -DCOMPUTE_CAPABILITY="${{ env.COMPUTE_CAPABILITY }}" \
            -DCOMPUTE_BACKEND=cuda \
            -DNO_CUBLASLT=$NO_CUBLASLT \
            -S .

            cmake --build build --config ${{ matrix.build_type }}
          done

    - name: Copy libraries
      shell: bash
      run: |
        mkdir -p output/${{ matrix.os }}/${{ matrix.arch }}
        echo " check lib files..."
        file bitsandbytes/libbitsandbytes*.*
        ( shopt -s nullglob && cp bitsandbytes/*.{so,dylib,dll} output/${{ matrix.os }}/${{ matrix.arch }}/ )
    - name: Upload build artifact
      uses: actions/upload-artifact@v4
      with:
        name: shared_library_cuda_${{ matrix.os }}_${{ matrix.arch }}_${{ matrix.cuda-version }}
        path: output/*

  build-wheels:
    needs:
    - build-shared-libs
    - build-shared-libs-cuda
    strategy:
      matrix:
        os: [ubuntu-latest, macos-latest, windows-latest]
        python-version: ["3.10"]
        arch: [x86_64, aarch64]
        exclude:
          - os: windows-latest # This probably requires arm64 Windows agents
            arch: aarch64
    runs-on: ${{ matrix.os }}
    steps:
      # Check out code
    - uses: actions/checkout@v4
      # Download shared libraries
    - name: Download build artifact
      uses: actions/download-artifact@v4
      with:
        merge-multiple: true
        pattern: "shared_library*_${{ matrix.os }}_${{ matrix.arch }}*"
        path: output/
    - name: Copy correct platform shared library
      shell: bash
      run: |
        ls -lR output/
        cp output/${{ matrix.os }}/${{ matrix.arch }}/* bitsandbytes/
      # Set up the Python version needed
    - name: Set up Python ${{ matrix.python-version }}
      uses: actions/setup-python@v5
      with:
        python-version: ${{ matrix.python-version }}
        cache: pip
    - name: Install build package
      shell: bash
      run: pip install build
    - name: Install Python test dependencies
      shell: bash
      run: pip install -r requirements-ci.txt
    # TODO: How to run CUDA tests on GitHub actions?
    #- name: Run unit tests
    #  if: ${{ matrix.arch == 'x86_64' }} # Tests are too slow to run in emulation. Wait for real aarch64 agents
    #  run: |
    #    PYTHONPATH=. pytest --log-cli-level=DEBUG tests
    - name: Build wheel
      shell: bash
      run: |
        python -m build . --wheel
        # fix wheel name
        if [ "${{ matrix.arch }}" = "aarch64" ]; then
            o=$(ls dist/*.whl)
            n=$(echo $o | sed 's@_x86_64@_aarch64@')
            [ "$n" != "$o" ] && mv $o $n
        fi
    - name: Upload build artifact
      uses: actions/upload-artifact@v4
      with:
        name: bdist_wheel_${{ matrix.os }}-${{ matrix.arch }}
        path: ${{ github.workspace }}/dist/bitsandbytes-*.whl
  publish:
    needs: build-wheels
    runs-on: ubuntu-latest
    steps:
    - uses: actions/checkout@v4
    - name: Download build artifact
      uses: actions/download-artifact@v4
      with:
        path: dist/
        merge-multiple: true
        pattern: "bdist_wheel_*"
    - run: |
        ls -lR dist/
    - name: Publish to PyPi
      if: startsWith(github.ref, 'refs/tags')
      uses: pypa/gh-action-pypi-publish@release/v1
      with:
        password: ${{ secrets.pypi }}