[water] Add ScaledMMAOp #4841

Workflow file for this run

.github/workflows/ci-gpu.yaml at 5af85bd

	# 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@cdf6c1fa76f9f475f3d7449005a359c84ca0f306 # v5.0.3
	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@cdf6c1fa76f9f475f3d7449005a359c84ca0f306 # v5.0.3
	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@a0b538fa0b742a6aa35d6e2c169b4bd06d225a98 # v1.15.3
	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@cdf6c1fa76f9f475f3d7449005a359c84ca0f306 # v5.0.3
	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@cdf6c1fa76f9f475f3d7449005a359c84ca0f306 # v5.0.3
	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@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
	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@cdf6c1fa76f9f475f3d7449005a359c84ca0f306 # v5.0.3
	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

Navigation Menu

Uh oh!

[water] Add ScaledMMAOp #4841

Workflow file

[water] Add ScaledMMAOp #4841

Uh oh!

Workflow file for this run