blob: 21a796896ac12d83ade2cb5f8c9e3c1471e95806 [file] [log] [blame]
#!/bin/bash
# Copyright 2022 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
#
# Runs SHARK tank using both SHARK-Runtime and IREE-Runtime, producing benchmark
# numbers.
#
# Usage:
# ./run_shark.sh \
# --sha <SHA of https://github.com/nod-ai/SHARK.git to pin to> \
# --pytest-regex e.g. "cpu", "cuda", "cuda and torch". \
# --driver e.g. "cpu", "cuda", "vulkan" \
# --save-version-info <true|false> \
# --iree-source-dir <source directory IREE if a local version is used> \
# --output-dir <local output dir>
set -xeuo pipefail
export SAVE_VERSION_INFO=false
export BENCHMARK_REGEX="cpu"
export DRIVER="cpu"
export IREE_SOURCE_DIR=""
while [[ "$#" -gt 0 ]]; do
case $1 in
--sha)
SHARK_SHA="$2"
shift;;
--pytest-regex)
BENCHMARK_REGEX="$2"
shift;;
-d|--driver)
DRIVER="$2"
shift;;
--save-version-info)
SAVE_VERSION_INFO="$2"
shift;;
-o|--output-dir)
SHARK_OUTPUT_DIR="$(realpath $2)"
shift;;
--iree-source-dir)
if [[ ! -z "$2" ]]; then
IREE_SOURCE_DIR="$(realpath $2)"
fi
shift;;
*)
echo "Unknown parameter passed: $1"
exit 1;;
esac
shift
done
mkdir "${SHARK_OUTPUT_DIR}"
git clone https://github.com/nod-ai/SHARK.git
pushd SHARK
git reset --hard ${SHARK_SHA}
# Remove existing data.
rm -rf ./shark_tmp
rm -rf ~/.local/shark_tank
declare -a args=(
--benchmark
--update_tank
--maxfail=500
-k "${BENCHMARK_REGEX}"
)
if [[ ${DRIVER} == "cuda" ]]; then
args+=(--tf32)
fi
# Run with SHARK-Runtime.
PYTHON=python3.11 VENV_DIR=shark.venv BENCHMARK=1 IMPORTER=1 ./setup_venv.sh
source shark.venv/bin/activate
export SHARK_VERSION=`pip show iree-compiler | grep Version | sed -e "s/^Version: \(.*\)$/\1/g"`
pytest "${args[@]}" tank/test_models.py || true
echo "######################################################"
echo "Benchmarks for SHARK-Runtime Complete"
cat bench_results.csv
mv bench_results.csv "${SHARK_OUTPUT_DIR}/${DRIVER}_shark_${SHARK_VERSION}.csv"
echo "######################################################"
deactivate
# Remove existing data.
rm -rf ./shark_tmp
rm -rf ~/.local/shark_tank
# Run with IREE.
PYTHON=python3.11 VENV_DIR=iree.venv BENCHMARK=1 IMPORTER=1 USE_IREE=1 ./setup_venv.sh
source iree.venv/bin/activate
export IREE_VERSION=$(pip show iree-compiler | grep Version | sed -e "s/^Version: \(.*\)$/\1/g")
if [[ ! -z "${IREE_SOURCE_DIR}" ]]; then
# `setup_venv.sh` would have installed a nightly release of IREE so we uninstall here.
pip uninstall -y iree-compiler iree-runtime
pushd "${IREE_SOURCE_DIR}"
git submodule update --init
pip install -r ./runtime/bindings/python/iree/runtime/build_requirements.txt
export IREE_VERSION="sha_$(git rev-parse --short=10 HEAD)"
# We build using the same Python Virtual Environment as SHARK to ensure compatibility.
cmake -GNinja -B iree-build -S . \
-DCMAKE_BUILD_TYPE=RelWithDebInfo \
-DIREE_ENABLE_ASSERTIONS=ON \
-DCMAKE_C_COMPILER=clang \
-DCMAKE_CXX_COMPILER=clang++ \
-DIREE_ENABLE_LLD=ON \
-DIREE_BUILD_PYTHON_BINDINGS=ON \
-DPython3_EXECUTABLE="$(which python)" \
-DIREE_HAL_DRIVER_CUDA=ON \
-DIREE_TARGET_BACKEND_CUDA=ON
cmake --build iree-build
# Install Python bindings.
pushd iree-build
pip install -e compiler/
pip install -e runtime/
popd # iree-build
popd # ${IREE_SOURCE_DIR}
fi
pytest "${args[@]}" tank/test_models.py || true
echo "######################################################"
echo "Benchmarks for IREE Complete"
cat bench_results.csv
mv bench_results.csv "${SHARK_OUTPUT_DIR}/${DRIVER}_iree_${IREE_VERSION}.csv"
echo "######################################################"
# Save version info.
if [[ "${SAVE_VERSION_INFO}" == true ]]; then
touch "${SHARK_OUTPUT_DIR}/version_info.txt"
echo "iree=${IREE_VERSION}" >> "${SHARK_OUTPUT_DIR}/version_info.txt"
echo "shark=${SHARK_VERSION}" >> "${SHARK_OUTPUT_DIR}/version_info.txt"
export TF_VERSION=$(pip show tensorflow | grep "Version" | awk '{print $NF}')
echo "tf=Tensorflow XLA ${TF_VERSION}" >> "${SHARK_OUTPUT_DIR}/version_info.txt"
export TORCH_VERSION=$(pip show torch | grep "Version" | awk '{print $NF}')
echo "torch=Torch Inductor ${TORCH_VERSION}" >> "${SHARK_OUTPUT_DIR}/version_info.txt"
fi
deactivate