test.yml

name: Test

on:
  pull_request:
    branches:
      - main
  push:
    branches:
      - main
  workflow_dispatch:
    branches:
      - main
  schedule:
    # At minute 0 past hour 6. (see https://crontab.guru)
    - cron: '00 06 * * *'

concurrency:
  # A PR number if a pull request and otherwise the commit hash. This cancels
  # queued and in-progress runs for the same PR (presubmit) or commit
  # (postsubmit).
  group: ci-build-test-${{ github.event.number || github.sha }}
  cancel-in-progress: true

env:
  SYSTEM_VERSION_COMPAT: 0

jobs:

  test-mlir-bindings:

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

    strategy:
      fail-fast: false
      matrix:
        os: [ ubuntu-22.04, macos-13, macos-14, windows-2022 ]
        py_version: [ "3.8", "3.9", "3.10", "3.11", "3.12" ]

        exclude:
          - os: macos-13
            py_version: "3.8"

          - os: macos-13
            py_version: "3.9"

          - os: macos-14
            py_version: "3.8"

          - os: macos-14
            py_version: "3.9"

    steps:
      - name: Checkout
        uses: actions/checkout@v2

      - name: Setup Python
        uses: actions/setup-python@v4
        with:
          python-version: ${{ matrix.py_version }}
          allow-prereleases: true

      - name: Free disk space
        if: contains(matrix.os, 'ubuntu')
        uses: descriptinc/free-disk-space@main
        with:
          tool-cache: true
          android: true
          dotnet: true
          haskell: true
          large-packages: true
          swap-storage: false # This frees space on the wrong partition.

      - uses: Jimver/cuda-toolkit@v0.2.15
        if: contains(matrix.os, 'ubuntu')
        id: cuda-toolkit
        with:
          cuda: '12.4.1'
          linux-local-args: '["--toolkit"]'
          log-file-suffix: "${{ matrix.os }}-${{ matrix.py_version }}.txt"

      - name: Install and configure
        shell: bash
        run: |
          pip install .[test,mlir] -v -f https://makslevental.github.io/wheels 

      - name: Test
        shell: bash
        run: |
          if [ ${{ matrix.os }} == 'windows-2022' ]; then
            pytest -s tests
          else
            pytest --capture=tee-sys tests
          fi

      - name: Test mwe
        shell: bash
        run: |
          python examples/mwe.py

  test-other-host-bindings:

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

    strategy:
      fail-fast: false
      matrix:
        os: [ ubuntu-22.04, macos-12, macos-14, windows-2022 ]
        py_version: [ "3.10", "3.11" ]

    steps:
      - name: Checkout
        uses: actions/checkout@v2

      - name: Setup Python
        uses: actions/setup-python@v4
        with:
          python-version: ${{ matrix.py_version }}
          allow-prereleases: true

      - name: Install and configure
        shell: bash
        run: |
          export PIP_FIND_LINKS=https://makslevental.github.io/wheels
          HOST_MLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir pip install .[test,jax] -v

      - name: Test
        shell: bash
        run: |
          if [ ${{ matrix.os }} == 'windows-2022' ]; then
            pytest -s tests/test_other_hosts.py
          else
            pytest --capture=tee-sys tests/test_other_hosts.py
          fi

  test-jupyter:

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

    strategy:
      fail-fast: false
      matrix:
        os: [ ubuntu-22.04 ]
        py_version: [ "3.10" ]

    steps:
      - name: Checkout
        uses: actions/checkout@v2

      - name: Setup Python
        uses: actions/setup-python@v4
        with:
          python-version: ${{ matrix.py_version }}
          allow-prereleases: true

      - name: Run notebook
        shell: bash
        run: |
          
          pip install jupyter
          pip install -q mlir-python-bindings -f https://makslevental.github.io/wheels
          pip install -q .
          
          sed -i.bak 's/OUTPUT_TIMEOUT = 10/OUTPUT_TIMEOUT = 100/g' \
            $(python -c 'import site; print(site.getsitepackages()[0])')/jupyter_client/runapp.py
          
          jupyter execute examples/mlir_python_extras.ipynb
          jupyter execute examples/vectorization_e2e.ipynb

  mlir-bindings-aarch64:

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

    strategy:
      fail-fast: false
      matrix:
        os: [ ubuntu-22.04 ]
        py_version: [ "3.8", "3.9", "3.10", "3.11", "3.12" ]

    steps:
      - name: Checkout
        uses: actions/checkout@v2

      - name: Install, configure, and test
        uses: uraimo/run-on-arch-action@v2
        with:
          arch: aarch64
          distro: ubuntu20.04
          dockerRunArgs: --volume "${{ github.workspace }}:/workspace"
          install: |
            
            apt-get update -q -y
            apt-get install -y wget build-essential
            
            mkdir -p ~/miniconda3
            wget -q https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-aarch64.sh -O miniconda.sh
            bash miniconda.sh -b -u -p /root/miniconda3
            eval "$(/root/miniconda3/bin/conda shell.bash hook)"
            conda init

          run: |
            
            eval "$(/root/miniconda3/bin/conda shell.bash hook)"
            conda create -n env -q -y -c conda-forge/label/python_rc python=${{ matrix.py_version }} 
            conda activate env
            
            cd /workspace
            
            pip install -q .[test,mlir] -f https://makslevental.github.io/wheels
            
            pytest --capture=tee-sys tests
            python examples/mwe.py