#!/bin/bash
# A/B 对比：FlagTree RemoveLayoutConversions 增强关 vs 开，逐算子两轮。
#
# 口径：每个算子先跑关（record autotune config），再用 REPLAY_FROM 回放同一套
# config 跑开，两侧只差 FLAGTREE_<VENDOR>_RLC_ENHANCE 这一个开关。产物沿用
# run_pytest.sh 的目录结构，所以每个算子恰好留下两份 run.log：
#
#   runs/<op>_<时间戳>/run.log   关（before）
#   runs/<op>_<时间戳>/run.log   开（after）
#
# 跑完用 batch_summary.py --ab <before> <after> 得到该算子的 speedup 区间。
#
# 用法：
#   OPS="softmax gelu_backward" bash ab_rlc_test.sh
#   OPS="sort:perf_sort topk" bash ab_rlc_test.sh     # 冒号分隔 OP:OP_FILE
#
# 注意：RLC 开关名带厂商前缀，用 VENDOR 指定（默认 MUSA）。套用前先确认开关有效：
# grep 该开关名于 FlagTree 源码应有命中，否则两侧跑的是同一份代码。
set -uo pipefail

cd "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

VENDOR="${VENDOR:-MUSA}"
ENHANCE_VAR="FLAGTREE_${VENDOR}_RLC_ENHANCE"

# 空格分隔的算子表。元素为 OP 或 OP:OP_FILE（省略时 OP_FILE=OP）。
OPS="${OPS:-softmax}"

if [[ -t 1 && -z "${NO_COLOR:-}" ]]; then
  export FLAGGEMS_PERF_COLOR=always
fi

LOG="runs/ab_rlc_$(date +%Y%m%d-%H%M%S).log"
mkdir -p runs

{
  echo ">>> vendor=$VENDOR  switch=$ENHANCE_VAR  ops=$OPS"
  for entry in $OPS; do
    op="${entry%%:*}"
    op_file="${entry#*:}"
    [[ "$op_file" == "$entry" ]] && op_file="$op"

    echo "########## $op | $ENHANCE_VAR=0 before (record) ##########"
    if ! OP="$op" OP_FILE="$op_file" env "$ENHANCE_VAR=0" bash run_pytest.sh; then
      echo "!!! FAILED before: $op"
      continue
    fi
    BEFORE=$(ls -dt "runs/${op}_"* | head -1)

    echo "########## $op | $ENHANCE_VAR=1 after (replay $BEFORE) ##########"
    if ! OP="$op" OP_FILE="$op_file" REPLAY_FROM="$BEFORE" env "$ENHANCE_VAR=1" bash run_pytest.sh; then
      echo "!!! FAILED after: $op"
      continue
    fi
    AFTER=$(ls -dt "runs/${op}_"* | head -1)

    echo "########## $op | summary ##########"
    python batch_summary.py --ab "$BEFORE/run.log" "$AFTER/run.log" --label "$op" \
      || echo "!!! summary failed: $op"
  done
  echo "########## ALL DONE ##########"
} 2>&1 | tee "$LOG"

# 去掉 ANSI 转义以便 grep/diff（终端输出保留颜色）
sed -i -E $'s/\x1b\\[[0-9;]*[A-Za-z]//g' "$LOG"

echo ">>> 总 log: $(pwd)/$LOG"
echo ">>> 每算子两份 run.log 在 runs/<op>_<时间戳>/ 下，按时间先后即 before / after"