#!/bin/bash
set -e

# export MIOPEN_ENABLE_LOGGING_CMD=1
# export MIOPEN_ENABLE_LOGGING=1
# export MIOPEN_LOG_LEVEL=6
# export ROCBLAS_LAYER=4
export HIP_VISIBLE_DEVICES=1
export MODEL_NAME=evo2_7b
export MODEL_PATH=/models/arcinstitute/evo2_7b/evo2_7b.pt

EVO_CMD="numactl -m 1 -N 1 \
            python -m evo2.test.test_evo2_generation_batched \
            --model_name ${MODEL_NAME} \
            --local_path ${MODEL_PATH} \
            --n_tokens 500 \
            --n_warmups 1 \
            --prompt_stretch \
            --trace_gzip \
            --trace_logdir ./log/pt-trace/stretch"

for batch_size in 1 2; do
    echo "================================================"
    echo "Running all tests for batch size ${batch_size}"
    echo "================================================"
    mkdir -p log &>/dev/null

    echo "==== Normal run ===="
    ${EVO_CMD} --batch_size ${batch_size}

    for step in 0 1; do
        echo "==== Torch profiler trace for prompt ${step} ===="
        ${EVO_CMD} --batch_size ${batch_size} --trace --trace_step ${step} \
            --trace_file_prefix evo2-bw1000-bs${batch_size}-s${step}
    done

    # echo "==== Hipprof trace ===="
    # hipprof --hip-trace -o log/trace-bs${batch_size} \
    #     ${EVO_CMD} --batch_size ${batch_size}

    # echo "==== Nsight-systems trace ===="
    # nsys profile --force-overwrite=true \
    #     --stats=true --trace=cuda \
    #     -o log/trace-bs${batch_size} \
    #     ${EVO_CMD} --batch_size ${batch_size}
done