# SPDX-License-Identifier: MIT
# Copyright (C) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
#!/bin/bash

if [ $# -ge 1 ] ; then
    FMA_API=$1  # build fwd/bwd/fwd_v3/bwd_v3
else
    FMA_API=""  # build all
fi

echo "######## building mha kernel $FMA_API"
python3 compile.py --api=$FMA_API

SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
TOP_DIR=$(dirname "$SCRIPT_DIR")/../../

if [ x"$FMA_API" = x"fwd" ] || [ x"$FMA_API" = x"fwd_v3" ] || [ x"$FMA_API" = x"" ] ; then
echo "######## linking mha fwd"

[ x"$FMA_API" = x"fwd_v3" ] && splitkv_api=0 || splitkv_api=1

/opt/rocm/bin/hipcc  -I$TOP_DIR/3rdparty/composable_kernel/include \
                     -I$TOP_DIR/3rdparty/composable_kernel/example/ck_tile/01_fmha/ \
                     -I$TOP_DIR/csrc/include \
                     -std=c++20 -O3 \
                     -DUSE_ROCM=1 \
                     -DCK_TILE_FMHA_FWD_SPLITKV_API=$splitkv_api \
                     --offload-arch=native \
                     -L $SCRIPT_DIR -lmha_fwd \
                     $SCRIPT_DIR/benchmark_mha_fwd.cpp -o fwd.exe
fi

if [ x"$FMA_API" = x"bwd" ] || [ x"$FMA_API" = x"bwd_v3" ] || [ x"$FMA_API" = x"" ] ; then
echo "######## linking mha bwd"
/opt/rocm/bin/hipcc  -I$TOP_DIR/3rdparty/composable_kernel/include \
                     -I$TOP_DIR/3rdparty/composable_kernel/example/ck_tile/01_fmha/ \
                     -I$TOP_DIR/csrc/include \
                     -std=c++20 -O3 \
                     -DUSE_ROCM=1 \
                     --offload-arch=native \
                     -L $SCRIPT_DIR -lmha_bwd \
                     $SCRIPT_DIR/benchmark_mha_bwd.cpp -o bwd.exe
fi
