name: Build cmake
inputs:
  opt_level:
    description: 'Compile options / optimization level.'
    required: false
    default: generic
  gpu:
    description: 'Enable GPU support.'
    required: false
    default: OFF
  raft:
    description: 'Enable RAFT support.'
    required: false
    default: OFF
  rocm:
    description: 'Enable ROCm support.'
    required: false
    default: OFF
runs:
  using: composite
  steps:
    - name: Setup miniconda
      uses: conda-incubator/setup-miniconda@v3
      with:
        python-version: '3.11'
        miniconda-version: latest
    - name: Configure build environment
      shell: bash
      run: |
        # initialize Conda
        conda config --set solver libmamba
        conda update -y -q conda
        echo "$CONDA/bin" >> $GITHUB_PATH

        conda install -y -q python=3.11 cmake make swig numpy scipy pytest

        # install base packages for ARM64
        if [ "${{ runner.arch }}" = "ARM64" ]; then
          conda install -y -q -c conda-forge openblas gxx_linux-aarch64 sysroot_linux-aarch64
        fi

        # install base packages for X86_64
        if [ "${{ runner.arch }}" = "X64" ]; then
          # TODO: merge this with ARM64
          conda install -y -q -c conda-forge gxx_linux-64 sysroot_linux-64
          conda install -y -q mkl=2023 mkl-devel=2023
        fi

        # install CUDA packages
        if [ "${{ inputs.gpu }}" = "ON" ] && [ "${{ inputs.raft }}" = "OFF" ]; then
          conda install -y -q cuda-toolkit -c "nvidia/label/cuda-12.4.0"
        fi

        # install RAFT packages
        if [ "${{ inputs.raft }}" = "ON" ]; then
          conda install -y -q libraft cuda-version=12.4 cuda-toolkit -c rapidsai-nightly -c "nvidia/label/cuda-12.4.0" -c conda-forge
        fi

        # install test packages
        conda install -y pytest
        if [ "${{ inputs.rocm }}" = "ON" ]; then
          : # skip torch install via conda, we need to install via pip to get
            #  ROCm-enabled version until it's supported in conda by PyTorch
        elif [ "${{ inputs.gpu }}" = "ON" ]; then
          conda install -y -q pytorch pytorch-cuda=12.4 -c pytorch -c nvidia/label/cuda-12.4.0
        else
          conda install -y -q pytorch -c pytorch
        fi
    - name: ROCm - Install dependencies
      if: inputs.rocm == 'ON'
      shell: bash
      run: |
        # Update repos and install kmod, wget, gpg
        sudo apt-get update
        sudo apt-get install -y kmod wget gpg

        # Get UBUNTU version name
        UBUNTU_VERSION_NAME=`cat /etc/os-release | grep UBUNTU_CODENAME | awk -F= '{print $2}'`

        # Set ROCm version
        ROCM_VERSION="6.1.1"

        # Download, prepare, and install the package signing key
        mkdir --parents --mode=0755 /etc/apt/keyrings
        wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null

        # Add rocm repository
        wget -qO - http://repo.radeon.com/rocm/rocm.gpg.key | sudo apt-key add -
        rocm_baseurl="http://repo.radeon.com/rocm/apt/${ROCM_VERSION}"
        echo "deb [arch=amd64] ${rocm_baseurl} ${UBUNTU_VERSION_NAME} main" | sudo tee /etc/apt/sources.list.d/rocm.list
        sudo apt-get update --allow-insecure-repositories
        sudo apt-get install -y --allow-unauthenticated \
            "rocm-dev${ROCM_VERSION}" "rocm-utils${ROCM_VERSION}" "rocm-libs${ROCM_VERSION}"

        # Fake presence of MI200-class accelerators
        echo "gfx90a" | sudo tee /opt/rocm/bin/target.lst

        # Cleanup
        sudo apt-get autoclean && sudo apt-get clean
        sudo rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
    - name: ROCm - Hipify
      if: inputs.rocm == 'ON'
      shell: bash
      run: ./faiss/gpu/hipify.sh
    - name: Symblink system dependencies
      if: inputs.raft == 'ON' || inputs.rocm == 'ON'
      shell: bash
      run: |
        # symblink system libraries for HIP compiler
        sudo ln -s /lib/x86_64-linux-gnu/libc.so.6 /lib64/libc.so.6
        sudo ln -s /lib/x86_64-linux-gnu/libc_nonshared.a /usr/lib64/libc_nonshared.a
        sudo ln -s /usr/lib/x86_64-linux-gnu/libpthread.so.0 /lib64/libpthread.so.0
        sudo ln -s $HOME/miniconda3/x86_64-conda-linux-gnu/sysroot/usr/lib64/libpthread_nonshared.a /usr/lib64/libpthread_nonshared.a
    - name: Build all targets
      shell: bash
      run: |
        eval "$(conda shell.bash hook)"
        conda activate
        cmake -B build \
              -DBUILD_TESTING=ON \
              -DBUILD_SHARED_LIBS=ON \
              -DFAISS_ENABLE_GPU=${{ inputs.gpu }} \
              -DFAISS_ENABLE_RAFT=${{ inputs.raft }} \
              -DFAISS_ENABLE_ROCM=${{ inputs.rocm }} \
              -DFAISS_OPT_LEVEL=${{ inputs.opt_level }} \
              -DFAISS_ENABLE_C_API=ON \
              -DPYTHON_EXECUTABLE=$CONDA/bin/python \
              -DCMAKE_BUILD_TYPE=Release \
              -DBLA_VENDOR=${{ runner.arch == 'X64' && 'Intel10_64_dyn' || '' }} \
              -DCMAKE_CUDA_FLAGS=${{ runner.arch == 'X64' && '"-gencode arch=compute_75,code=sm_75"' || '' }} \
              .
        make -k -C build -j$(nproc)
    - name: C++ tests
      shell: bash
      run: |
        export GTEST_OUTPUT="xml:$(realpath .)/test-results/googletest/"
        make -C build test
    - name: Install Python extension
      shell: bash
      working-directory: build/faiss/python
      run: |
        $CONDA/bin/python setup.py install
    - name: ROCm - install ROCm-enabled torch via pip
      if: inputs.rocm == 'ON'
      shell: bash
      run: |
        pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.1
    - name: Python tests (CPU only)
      if: inputs.gpu == 'OFF'
      shell: bash
      run: |
        pytest --junitxml=test-results/pytest/results.xml tests/test_*.py
        pytest --junitxml=test-results/pytest/results-torch.xml tests/torch_*.py
    - name: Python tests (CPU + GPU)
      if: inputs.gpu == 'ON'
      shell: bash
      run: |
        pytest --junitxml=test-results/pytest/results.xml tests/test_*.py
        pytest --junitxml=test-results/pytest/results-torch.xml tests/torch_*.py
        cp tests/common_faiss_tests.py faiss/gpu/test
        pytest --junitxml=test-results/pytest/results-gpu.xml faiss/gpu/test/test_*.py
        pytest --junitxml=test-results/pytest/results-gpu-torch.xml faiss/gpu/test/torch_*.py
    - name: Test avx2 loading
      if: inputs.opt_level == 'avx2'
      shell: bash
      run: |
        FAISS_DISABLE_CPU_FEATURES=AVX2 LD_DEBUG=libs $CONDA/bin/python -c "import faiss" 2>&1 | grep faiss.so
        LD_DEBUG=libs $CONDA/bin/python -c "import faiss" 2>&1 | grep faiss_avx2.so
    - name: Upload test results
      uses: actions/upload-artifact@v4
      with:
        name: test-results-arch=${{ runner.arch }}-opt=${{ inputs.opt_level }}-gpu=${{ inputs.gpu }}-raft=${{ inputs.raft }}-rocm=${{ inputs.rocm }}
        path: test-results
