name: 'Build and Test on Linux'

inputs:
  toolkit:
    description: 'The toolkit to build with'
    required: false
    default: 'cpu'

runs:
  using: "composite"
  steps:
    - name: Set cmake args
      shell: bash
      env:
        CMAKE_ARGS_CPU: >
          -DCMAKE_COMPILE_WARNING_AS_ERROR=ON
          -DMLX_BUILD_CUDA=OFF
        CMAKE_ARGS_CUDA: >
          -DCMAKE_COMPILE_WARNING_AS_ERROR=ON
          -DCMAKE_CUDA_COMPILER=/usr/local/${{ inputs.toolkit }}/bin/nvcc
          -DMLX_BUILD_CUDA=ON
      run: |
        CMAKE_ARGS="$CMAKE_ARGS_${{ startsWith(inputs.toolkit, 'cuda') && 'CUDA' || 'CPU' }}"
        # FIXME: Unable to build tests without libnccl on arm64.
        if ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }} ; then
          CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF"
        fi
        echo CMAKE_ARGS="$CMAKE_ARGS" > $GITHUB_ENV

    - name: Install Python package
      shell: sh
      env:
        DEBUG: 1
      run: pip install --no-build-isolation -e ".[dev]" -v

    - name: Generate package stubs
      # FIXME: Unable to generate stubs without libnccl on arm64.
      if: ${{ !(startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64') }}
      shell: sh
      run: |
        pip install typing_extensions
        python setup.py generate_stubs

    - name: Build CPP only
      shell: bash
      run: |
        cmake . -B build -DCMAKE_BUILD_TYPE=Debug ${{ env.CMAKE_ARGS }}
        cmake --build build -j $(nproc)