test.yml

name: Test

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

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
  PIP_FIND_LINKS: "https://github.com/makslevental/mlir-wheels/releases/expanded_assets/latest"

jobs:

  test-mlir-bindings:

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

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

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

          - os: macos-14
            py_version: "3.9"

          # <frozen importlib._bootstrap>:491: Warning: Numpy built with MINGW-W64 on Windows 64 bits is experimental, and only available for testing. You are advised not to use it for production.
          - os: windows-2022
            py_version: "3.14"

          - os: windows-2022
            py_version: "3.14t"

    name: "${{ matrix.os }}-${{ matrix.py_version }}"

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

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

      - name: Install and configure
        shell: bash
        run: |
        
            pip install .[test,mlir] -v
    
      - name: Test
        shell: bash
        run: |
          
          # so that loaded modules isn't reused between modules
          for f in tests/test*.py; do
            echo "Testing $f"
            pytest -v $f
          done

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

  test-mlir-bindings-pyodide:

    runs-on: ubuntu-latest
    name: test pyodide

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

      - name: Setup Python
        uses: actions/setup-python@v5.6.0
        with:
          python-version: 3.12

      - name: Install pyodide and test
        shell: bash
        run: |
        
          pip install pyodide-build
          pyodide venv venv
          . venv/bin/activate
          
          pip install pytest astpretty
          pip install .[mlir] -v
          
          pytest -v tests

          python examples/mwe.py
          python examples/flash_attention.py

  test-other-host-bindings:

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

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

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

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

      - name: Install and configure
        shell: bash
        run: |
          pip install .[test,mlir] -v
          HOST_MLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir pip install .[test,jax] -v

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

  test-examples:

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

    strategy:
      fail-fast: false
      matrix:
        os: [ ubuntu-22.04, macos-13, macos-14, windows-2022, ubuntu-22.04-arm ]
        py_version: [ "3.12", "3.13", "3.14", "3.14t" ]

        # <frozen importlib._bootstrap>:491: Warning: Numpy built with MINGW-W64 on Windows 64 bits is experimental, and only available for testing. You are advised not to use it for production.
        exclude:
          - os: windows-2022
            py_version: "3.14"

          - os: windows-2022
            py_version: "3.14t"

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

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

      - name: Install and configure
        shell: bash
        run: |
          
          pip install .[test,mlir] -v

      - name: Test gpu examples
        shell: bash
        run: |
          
          python examples/rdna_matmul_opt.py
          python examples/cuda_matmul_opt.py

  test-jupyter-notebooks:

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

    strategy:
      fail-fast: false
      matrix:
        os: [ ubuntu-22.04, macos-13, macos-14, windows-2022, ubuntu-22.04-arm ]
        py_version: [ "3.12", "3.13", "3.14", "3.14t" ]

        # <frozen importlib._bootstrap>:491: Warning: Numpy built with MINGW-W64 on Windows 64 bits is experimental, and only available for testing. You are advised not to use it for production.
        exclude:
          - os: windows-2022
            py_version: "3.14"

          - os: windows-2022
            py_version: "3.14t"

    steps:

      - name: Checkout
        uses: actions/checkout@v2

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

      - name: Test jupyter notebooks
        # sed: can't read C:\hostedtoolcache\windows\Python\3.12.10\x64/jupyter_client/runapp.py: No such file or directory
        if: matrix.os != 'windows-2022'
        shell: bash
        env:
          BRANCH: ${{ github.head_ref || github.ref_name }}
        run: |
          
          pip install -q jupyter
          
          sed -i.bak 's/OUTPUT_TIMEOUT = 10/OUTPUT_TIMEOUT = 1000/g' \
            $(python -c 'import site; print(site.getsitepackages()[0])')/jupyter_client/runapp.py
          
          jupyter execute examples/mlir_python_extras.ipynb --output=mlir_python_extras_output
          cat examples/mlir_python_extras_output.ipynb | jq '.cells[].outputs | select(length > 0) | .[0] | .text'
          jupyter execute examples/vectorization_e2e.ipynb --output=vectorization_e2e_output
          cat examples/vectorization_e2e_output.ipynb | jq '.cells[].outputs | select(length > 0) | .[0] | .text'
          if [ ${{ matrix.os }} == 'ubuntu-22.04' ]; then
            jupyter execute examples/cuda_e2e.ipynb --output=cuda_e2e_output
            cat examples/cuda_e2e_output.ipynb | jq '.cells[].outputs | select(length > 0) | .[0] | .text'
          fi