ci-gpu.yaml

# Copyright 2024 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

name: "Wave GPU CI"

on:
  workflow_dispatch:
  pull_request:
      types: [opened, synchronize, ready, ready_for_review, converted_to_draft]
  push:
    branches:
      - main

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). The workflow name is prepended to avoid conflicts between
  # different workflows.
  group: ${{ github.workflow }}-${{ github.event.number || github.sha }}
  cancel-in-progress: true

env:
  LLVM_SHA_FILE: llvm-sha.txt
  LLVM_CACHE_NUMBER: 2  # Increase to reset cache

jobs:
  # Water uses its own LLVM version independent of the one used by IREE and
  # potentially can have different compilation flags/enabled features.
  # We build and cache it here to avoid rebuilding it for every PR.
  # Water links LLVM statically, so it won't conflict with IREE's one.
  build_llvm_linux:
    name: Builds Linux LLVM/MLIR
    strategy:
      fail-fast: false
      matrix:
        version: [3.11]
        runs-on: [linux-mi325-1gpu-ossci-iree-org]
    runs-on: ${{matrix.runs-on}}
    timeout-minutes: 240 # Building LLVM can take multiple hours on public GH runners
    steps:
      - name: Checkout repo
        uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2

      - name: Setup Cache Vars
        run: |
          echo "LLVM_SHA=$(cat $GITHUB_WORKSPACE/water/$LLVM_SHA_FILE)" >> $GITHUB_ENV

      - name: Cache LLVM-MLIR
        id: cache-llvm-mlir
        uses: actions/cache@668228422ae6a00e4ad889ee87cd7109ec5666a7 # v5.0.4
        with:
          path: llvm-mlir/_mlir_install/**
          key: ${{ runner.os }}-build-llvm-${{ env.LLVM_CACHE_NUMBER }}-${{ env.LLVM_SHA }}

      - name: Setup env
        if: steps.cache-llvm-mlir.outputs.cache-hit != 'true'
        run: |
          sudo apt-get update
          sudo apt-get install -y ninja-build cmake clang lld dwarfdump

      - name: "Setting up Python"
        id: setup_python
        uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
        with:
          python-version: ${{matrix.version}}
          pip-install: -r water/requirements-dev.txt

      - name: Checkout LLVM
        if: steps.cache-llvm-mlir.outputs.cache-hit != 'true'
        uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
        with:
          repository: llvm/llvm-project
          ref: ${{ env.LLVM_SHA }}
          path: llvm-mlir/llvm-project

      - name: Build LLVM-MLIR
        if: steps.cache-llvm-mlir.outputs.cache-hit != 'true'
        run: |
          pushd ${GITHUB_WORKSPACE}/llvm-mlir
          echo "INFO: Need to rebuild LLVM-MLIR. Previous installation for MLIR not found"
          np=`nproc`
          echo "INFO: nproc $np"
          mkdir _build
          cd _build
          export CC=clang
          export CXX=clang++
          cmake ../llvm-project/llvm                                           \
            -GNinja                                                            \
            -DCMAKE_BUILD_TYPE=Release                                         \
            -DLLVM_ENABLE_PROJECTS="mlir;llvm;lld;clang"                       \
            -DLLVM_ENABLE_ASSERTIONS=ON                                        \
            -DLLVM_INSTALL_UTILS=ON                                            \
            -DLLVM_TARGETS_TO_BUILD="X86;AMDGPU"                               \
            -DLLVM_ENABLE_BINDINGS=OFF                                         \
            -DLLVM_ENABLE_ZSTD=OFF                                             \
            -DMLIR_INCLUDE_TESTS=OFF                                           \
            -DLLVM_USE_LINKER=lld                                              \
            -DLLVM_DISTRIBUTION_COMPONENTS="llvm-headers;llvm-libraries;cmake-exports;FileCheck;count;not;mlir-headers;mlir-libraries;mlir-cmake-exports;mlir-tblgen;mlir-python-sources;lld;clang;clang-resource-headers" \
            -DMLIR_ENABLE_BINDINGS_PYTHON=ON                                   \
            -DCMAKE_INSTALL_PREFIX=${GITHUB_WORKSPACE}/llvm-mlir/_mlir_install
          echo "INFO: working around a missing dependency on stubgen"
          ninja MLIRPythonModules.extension._mlir.dso._mlir.type_stubs
          ninja install-distribution-stripped
          popd


  test:
    name: "${{ contains( matrix.os, 'self-hosted') && matrix.os[0] || matrix.os }} :: ${{ matrix.version }} :: Unit Tests and Type Checking"
    strategy:
      fail-fast: false
      matrix:
        version: [3.11]
        os: [ubuntu-22.04, linux-mi325-1gpu-ossci-iree-org, linux-mi35x-1gpu-ossci-iree-org, [rdna4, self-hosted, Linux, X64, shark49]] # nodai-amdgpu-mi250-x86-64
    runs-on: ${{matrix.os}}
    timeout-minutes: 60
    needs: build_llvm_linux
    if: github.event_name != 'pull_request' || github.event.pull_request.draft == false
    env:
      VENV_DIR: ${{ github.workspace }}/.wave-venv
      IS_CDNA3: ${{ contains(matrix.os, 'mi325') }}
      IS_CDNA4: ${{ contains(matrix.os, 'mi35x') }}
      IS_RDNA4: ${{ contains(matrix.os, 'rdna4') }}
      HAS_GPU: ${{ contains(matrix.os, 'rdna4') || contains(matrix.os, 'mi325') ||  contains(matrix.os, 'mi35x') }}

    steps:
      - name: Checkout repo
        uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
        with:
          fetch-depth: 0

      - name: Print env
        run: |
          echo "IS_CDNA3=$IS_CDNA3"
          echo "IS_CDNA4=$IS_CDNA4"
          echo "IS_RDNA4=$IS_RDNA4"
          echo "HAS_GPU=$HAS_GPU"

      - name: Setup Cache Vars
        if: ${{ env.IS_CDNA3 == 'true' || env.IS_CDNA4 == 'true' }}
        run: |
          echo "LLVM_SHA=$(cat $GITHUB_WORKSPACE/water/$LLVM_SHA_FILE)" >> $GITHUB_ENV
          echo "WAVE_BUILD_WATER=1" >> $GITHUB_ENV
          echo "WAVE_BUILD_WAVEASM=1" >> $GITHUB_ENV
          echo "WAVE_LLVM_DIR=${GITHUB_WORKSPACE}/llvm-mlir/_mlir_install" >> $GITHUB_ENV

      - name: Cache LLVM-MLIR
        id: cache-llvm-mlir
        uses: actions/cache@668228422ae6a00e4ad889ee87cd7109ec5666a7 # v5.0.4
        if: ${{ env.IS_CDNA3 == 'true' || env.IS_CDNA4 == 'true' }}
        with:
          path: llvm-mlir/_mlir_install/**
          key: ${{ runner.os }}-build-llvm-${{ env.LLVM_CACHE_NUMBER }}-${{ env.LLVM_SHA }}

      - name: "Setting up Python"
        id: setup_python
        uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
        with:
          python-version: ${{matrix.version}}

      - name: Create Python venv
        run: |
          python -m venv ${VENV_DIR}
          source ${VENV_DIR}/bin/activate
          echo VIRTUAL_ENV=$VIRTUAL_ENV >> "$GITHUB_ENV"
          echo "$VENV_DIR/bin" >> "$GITHUB_PATH"

      - name: "Setting up Rust"
        uses: actions-rust-lang/setup-rust-toolchain@150fca883cd4034361b621bd4e6a9d34e5143606 # v1.15.4
        with:
          toolchain: stable

      - name: "Install dwarfdump"
        if: ${{ env.IS_RDNA4 == 'false' }}
        run: |
          sudo apt-get update
          sudo apt-get install -y dwarfdump

      - name: Install pip deps
        if: ${{ env.HAS_GPU == 'false' }}
        run: |
          python -m pip install --no-compile --upgrade pip
          # Note: We install in three steps in order to satisfy requirements
          # from non default locations first. Installing the PyTorch CPU
          # wheels saves multiple minutes and a lot of bandwidth on runner setup.
          pip install --no-compile -r pytorch-cpu-requirements.txt
          pip install --no-cache-dir -r requirements-iree-pinned.txt --upgrade
          pip install -e ".[dev]"

      - name: Install pip deps (CDNA4)
        if: ${{ env.IS_CDNA4 == 'true' }}
        run: |
          # Install TheRock
          python -m pip install --upgrade pip
          python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx950-dcgpu/ "rocm[libraries,devel]" "torch>=2.6,<2.9"

          # Set environment variables
          # TODO: Why do we even need this for mi350?
          export ROCM_PATH=$(python -m rocm_sdk path --root)
          export LD_LIBRARY_PATH="$ROCM_PATH/lib":$LD_LIBRARY_PATH
          echo "ROCM_PATH=$ROCM_PATH" >> $GITHUB_ENV
          echo "LD_LIBRARY_PATH=$LD_LIBRARY_PATH" >> $GITHUB_ENV

          pip install --no-cache-dir -r requirements-iree-pinned.txt --upgrade
          pip install -e ".[dev]"

      - name: Install pip deps (CDNA3/RDNA4)
        if: ${{ env.HAS_GPU == 'true' && env.IS_CDNA4 == 'false' }}
        run: |
          python -m pip install --upgrade pip
          ./gen-pytorch-rocm-requirements.py -o /tmp/pytorch-rocm-requirements.txt
          pip install -r /tmp/pytorch-rocm-requirements.txt
          pip install --no-cache-dir -r requirements-iree-pinned.txt --upgrade
          pip install -e ".[dev]"

      - name: Run unit tests
        run: |
          pytest -n 4 --capture=tee-sys -vv ./tests/unittests/
          pytest -n 4 --capture=tee-sys -vv ./tests/mlir_wave_iface

      - name: Run LIT tests
        env:
           WAVE_TEST_WATER: ${{ (env.IS_CDNA3 == 'true' || env.is_CDNA4 == 'true')  && '1' || '0'  }}
           WAVE_TEST_DWARFDUMP: ${{ env.IS_RDNA4 == 'false' && '1' || '0' }}
        run: |
          # TODO: can't sudo to install dwarfdump on rdna4
          echo "WAVE_TEST_WATER=$WAVE_TEST_WATER"
          echo "WAVE_TEST_DWARFDUMP=$WAVE_TEST_DWARFDUMP"
          lit lit_tests/ -vv

      - name: Test TKW runtime related stack on amdgpu
        if: ${{ env.HAS_GPU == 'true' }}
        run: |
          python -c "import torch; print(torch.cuda.get_device_properties().gcnArchName if torch.cuda.is_available() else 'cpu')"
          export WAVE_CACHE_DIR=$PWD/.wave
          rm -rf ./.wave
          nproc
          WAVE_CACHE_ON=1 pytest --timeout=300 --capture=tee-sys -vv --run-e2e --durations=100 ./tests/kernel/runtime

      - name: Run e2e tests on AMD GPU
        if: ${{ env.HAS_GPU == 'true' && (github.event_name == 'pull_request') }}
        run: |
          WAVE_CACHE_ON=0 pytest -n 4 --timeout=300 --capture=tee-sys -vv --run-e2e --durations=100 ./tests/kernel/

      - name: Run expensive e2e tests on AMD GPU
        if: ${{ env.HAS_GPU == 'true' && (github.event_name != 'pull_request') }}
        run: |
          WAVE_CACHE_ON=0 pytest -n 4 --timeout=600 --capture=tee-sys -vv --run-e2e --run-expensive-tests --durations=100 ./tests/kernel/

      - name: MyPy Type Checking
        run: |
          mypy


  test_linux_water:
    name: "${{ matrix.name }} :: SHARED_LIBS ${{ matrix.shared_libs }} :: Run water tests"
    strategy:
      fail-fast: false
      matrix:
        version: [3.11]
        name: [linux-mi325-1gpu-ossci-iree-org]
        shared_libs: ["ON", "OFF"]
    runs-on: ${{ matrix.name }}
    timeout-minutes: 60
    needs: build_llvm_linux

    steps:
      - name: Checkout repo
        uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
        with:
          fetch-depth: 0

      - name: Setup env
        run: |
          sudo apt-get update
          sudo apt-get install -y ninja-build cmake clang lld

      - name: "Setting up Python"
        id: setup_python
        uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
        with:
          python-version: ${{matrix.version}}
          pip-install: -r water/requirements-dev.txt

      - name: Setup Cache Vars
        run: |
          echo "LLVM_SHA=$(cat $GITHUB_WORKSPACE/water/$LLVM_SHA_FILE)" >> $GITHUB_ENV

      - name: Cache LLVM-MLIR
        id: cache-llvm-mlir
        uses: actions/cache@668228422ae6a00e4ad889ee87cd7109ec5666a7 # v5.0.4
        with:
          path: |
            llvm-mlir/_mlir_install/**
          key: ${{ runner.os }}-build-llvm-${{ env.LLVM_CACHE_NUMBER }}-${{ env.LLVM_SHA }}


      - name: Build water
        run: |
          export EXTERNAL_LIT=${GITHUB_WORKSPACE}/water/scripts/runlit.py
          export LLVM_DIR=${GITHUB_WORKSPACE}/llvm-mlir/_mlir_install
          mkdir -p cmake_build
          cd cmake_build
          export CC=clang
          export CXX=clang++
          cmake ${GITHUB_WORKSPACE}/water                                      \
            -GNinja                                                            \
            -DCMAKE_BUILD_TYPE=Release                                         \
            -DLLVM_DIR=${LLVM_DIR}/lib/cmake/llvm                              \
            -DMLIR_DIR=${LLVM_DIR}/lib/cmake/mlir                              \
            -DBUILD_SHARED_LIBS=${{ matrix.shared_libs }}                      \
            -DLLVM_EXTERNAL_LIT=${EXTERNAL_LIT}                                \
            -DMLIR_ENABLE_BINDINGS_PYTHON=ON                                   \
            -DWATER_ENABLE_PYTHON=ON
          cmake --build .

      - name: Test water
        # As we are mixing static LLVM build with shared Water we will get a
        # usual double option registration LLVM issue. Only compile Water in
        # this case to look for missing libraries.
        if: ${{ matrix.shared_libs == 'OFF' }}
        run: |
          cd cmake_build
          cmake --build . --target check-water

      - name: Build docs
        run: |
          cd cmake_build
          cmake --build . --target mlir-doc


  check_water_changes:
    name: Check for water/ changes
    if: github.event_name == 'pull_request'
    runs-on: ubuntu-latest
    outputs:
      changed: ${{ steps.check.outputs.changed }}
    steps:
      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
        with:
          fetch-depth: 0
      - id: check
        run: |
          CHANGED=$(git diff --name-only ${{ github.event.pull_request.base.sha }}...${{ github.sha }} -- water/)
          if [ -n "$CHANGED" ]; then
            echo "changed=true" >> "$GITHUB_OUTPUT"
          else
            echo "changed=false" >> "$GITHUB_OUTPUT"
          fi

  coverage_water:
    name: "CDNA4 :: Water Code Coverage"
    runs-on: linux-mi35x-1gpu-ossci-iree-org
    timeout-minutes: 60
    needs: [build_llvm_linux, check_water_changes]
    if: needs.check_water_changes.outputs.changed == 'true'
    permissions:
      pull-requests: write
    env:
      LLVM_DIR: ${{ github.workspace }}/llvm-mlir/_mlir_install

    steps:
      - name: Checkout repo
        uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
        with:
          fetch-depth: 0

      - name: Setup env
        run: |
          sudo apt-get update
          sudo apt-get install -y ninja-build cmake clang lld llvm libclang-rt-dev

      - name: "Setting up Python"
        id: setup_python
        uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
        with:
          python-version: "3.11"
          pip-install: -r water/requirements-dev.txt

      - name: Setup Cache Vars
        run: |
          echo "LLVM_SHA=$(cat $GITHUB_WORKSPACE/water/$LLVM_SHA_FILE)" >> $GITHUB_ENV

      - name: Cache LLVM-MLIR
        id: cache-llvm-mlir
        uses: actions/cache@668228422ae6a00e4ad889ee87cd7109ec5666a7 # v5.0.4
        with:
          path: llvm-mlir/_mlir_install/**
          key: ${{ runner.os }}-build-llvm-${{ env.LLVM_CACHE_NUMBER }}-${{ env.LLVM_SHA }}

      - name: Build water with coverage
        run: |
          export EXTERNAL_LIT=${GITHUB_WORKSPACE}/water/scripts/runlit.py
          mkdir -p cmake_build
          cd cmake_build
          export CC=clang
          export CXX=clang++
          cmake ${GITHUB_WORKSPACE}/water                                      \
            -GNinja                                                            \
            -DCMAKE_BUILD_TYPE=Release                                         \
            -DLLVM_DIR=${LLVM_DIR}/lib/cmake/llvm                              \
            -DMLIR_DIR=${LLVM_DIR}/lib/cmake/mlir                              \
            -DBUILD_SHARED_LIBS=OFF                                            \
            -DLLVM_EXTERNAL_LIT=${EXTERNAL_LIT}                                \
            -DMLIR_ENABLE_BINDINGS_PYTHON=ON                                   \
            -DWATER_ENABLE_PYTHON=ON                                           \
            -DWATER_ENABLE_COVERAGE=ON
          cmake --build .

      - name: Run tests with coverage
        run: |
          cd cmake_build
          mkdir -p coverage
          export LLVM_PROFILE_FILE="$PWD/coverage/default-%p-%m.profraw"
          cmake --build . --target check-water

      - name: Process coverage data
        run: |
          cd cmake_build
          llvm-profdata merge -sparse coverage/*.profraw -o coverage.profdata
          llvm-cov report ./bin/water-opt                                      \
            -show-region-summary=false                                         \
            -instr-profile=coverage.profdata                                   \
            -sources ${GITHUB_WORKSPACE}/water/lib                             \
                     ${GITHUB_WORKSPACE}/water/include                         \
            > coverage-summary.txt
          llvm-cov show ./bin/water-opt                                        \
            -instr-profile=coverage.profdata                                   \
            -sources ${GITHUB_WORKSPACE}/water/lib                             \
                     ${GITHUB_WORKSPACE}/water/include                         \
            --format=html -output-dir=coverage-html
          cat coverage-summary.txt

      - name: Upload coverage HTML report
        id: upload-coverage
        uses: actions/upload-artifact@bbbca2ddaa5d8feaa63e36b76fdaad77386f024f # v7.0.0
        with:
          name: water-coverage-html
          path: cmake_build/coverage-html

      - name: Post coverage comment on PR
        continue-on-error: true
        env:
          GITHUB_TOKEN: ${{ github.token }}
          ARTIFACT_URL: https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}/artifacts/${{ steps.upload-coverage.outputs.artifact-id }}
        run: |
          cd cmake_build
          {
            echo '<!-- water-coverage -->'
            echo '## Water Code Coverage'
            echo ''
            echo '```'
            cat coverage-summary.txt
            echo '```'
            echo ''
            echo "[Download full HTML report]($ARTIFACT_URL)"
          } > comment-body.txt

          API_URL="https://api.github.com/repos/${{ github.repository }}"
          AUTH_HEADER="Authorization: token $GITHUB_TOKEN"

          EXISTING=$(curl -s -H "$AUTH_HEADER" \
            "$API_URL/issues/${{ github.event.number }}/comments" | \
            jq '.[] | select(.body | contains("<!-- water-coverage -->")) | .id' | head -1)

          PAYLOAD=$(jq -Rs '{body: .}' comment-body.txt)

          if [ -n "$EXISTING" ]; then
            curl -s -X PATCH -H "$AUTH_HEADER" -H "Content-Type: application/json" \
              -d "$PAYLOAD" "$API_URL/issues/comments/$EXISTING"
          else
            curl -s -X POST -H "$AUTH_HEADER" -H "Content-Type: application/json" \
              -d "$PAYLOAD" "$API_URL/issues/${{ github.event.number }}/comments"
          fi


  test_linux_waveasm:
    name: "${{ matrix.name }} :: SHARED_LIBS ${{ matrix.shared_libs }} :: Run waveasm tests"
    strategy:
      fail-fast: false
      matrix:
        version: [3.11]
        name: [linux-mi325-1gpu-ossci-iree-org]
        shared_libs: ["ON", "OFF"]
    runs-on: ${{ matrix.name }}
    timeout-minutes: 60
    needs: build_llvm_linux

    steps:
      - name: Checkout repo
        uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
        with:
          fetch-depth: 0

      - name: Setup env
        run: |
          sudo apt-get update
          sudo apt-get install -y ninja-build cmake clang lld

      - name: "Setting up Python"
        id: setup_python
        uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
        with:
          python-version: ${{matrix.version}}
          pip-install: -r water/requirements-dev.txt

      - name: Setup Cache Vars
        run: |
          echo "LLVM_SHA=$(cat $GITHUB_WORKSPACE/water/$LLVM_SHA_FILE)" >> $GITHUB_ENV

      - name: Cache LLVM-MLIR
        id: cache-llvm-mlir
        uses: actions/cache@668228422ae6a00e4ad889ee87cd7109ec5666a7 # v5.0.4
        with:
          path: |
            llvm-mlir/_mlir_install/**
          key: ${{ runner.os }}-build-llvm-${{ env.LLVM_CACHE_NUMBER }}-${{ env.LLVM_SHA }}

      - name: Build waveasm
        run: |
          export EXTERNAL_LIT=${GITHUB_WORKSPACE}/water/scripts/runlit.py
          export LLVM_DIR=${GITHUB_WORKSPACE}/llvm-mlir/_mlir_install
          mkdir -p cmake_build_waveasm
          cd cmake_build_waveasm
          export CC=clang
          export CXX=clang++
          cmake ${GITHUB_WORKSPACE}/waveasm                                    \
            -GNinja                                                            \
            -DCMAKE_BUILD_TYPE=Release                                         \
            -DLLVM_DIR=${LLVM_DIR}/lib/cmake/llvm                              \
            -DMLIR_DIR=${LLVM_DIR}/lib/cmake/mlir                              \
            -DBUILD_SHARED_LIBS=${{ matrix.shared_libs }}                      \
            -DLLVM_EXTERNAL_LIT=${EXTERNAL_LIT}
          cmake --build .

      - name: Test waveasm
        # As we are mixing static LLVM build with shared WaveASM we will get a
        # usual double option registration LLVM issue. Only compile WaveASM in
        # this case to look for missing libraries.
        if: ${{ matrix.shared_libs == 'OFF' }}
        run: |
          cd cmake_build_waveasm
          cmake --build . --target check-waveasm