#!/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 区间。
#
# 这一轮同时充当 ab_repeat.sh 的筛选前置：一对的代价约 0.8 分钟，六轮 A/B 是
# 13~44 分钟（S5000 实测），差 20~50 倍。所以先用它过一遍，只把 IR_CHANGED 的
# 算子送进 ab_repeat.sh。判定写在 <root>/screen.tsv，四种之一：
#
#   IR_CHANGED  两侧 ttgir 不同        -> 排六轮
#   NO_IR_DIFF  逐字节相同（config 已锁死）-> 跳过，跑六轮只会测到噪声
#   NO_KERNEL   跑出行但没编出 .ttgir  -> 该算子在本后端没走 Triton，
#                                        编译器开关无从谈起。不是"开关无作用"
#   NO_ROWS     没有测量行            -> 先修可运行性，多半是 dtype
#
# 判据是 ttgir 指纹而不是 convert_layout 计数：cvt 数相同不等于生成代码相同。
# cvt 数照样记，它解释的是变化的性质。
#
# 用法：
#   OPS="softmax gelu_backward" bash ab_rlc_test.sh
#   OPS="sort:perf_sort topk" bash ab_rlc_test.sh          # OP:OP_FILE
#   OPS="perf_fft:fft:complex64 topk:topk" bash ab_rlc_test.sh   # 再加 DTYPES
#
# DTYPES 留空表示用该算子自己的全套。**不要全局写死一个 dtype**：fft 只吃
# complex64，被一刀切的 DTYPES=float16 判死过一次——那是 harness 的问题，
# 却被记成了算子的问题。
#
# 注意：RLC 开关名带厂商前缀，用 VENDOR 指定（默认 MUSA）。开关名拼错时两侧跑
# 同一份代码，此时判定会是 NO_IR_DIFF——先查拼写，再下"开关无作用"的结论。
set -uo pipefail

cd "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "./_ir_fingerprint.sh"

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

# 显式 OUT_DIR，不再用 `ls -dt runs/<op>_*` 猜"最新那个就是我刚跑的"。串行下
# 那个假设成立，多卡并行共用 runs/ 就不成立了。
ROOT="runs/ab_rlc_$(date +%Y%m%d-%H%M%S)"
mkdir -p "$ROOT"
LOG="$ROOT/summary.log"
TSV="$ROOT/screen.tsv"
printf 'op	op_file	dtypes	rows_before	rows_after	kernels_before	kernels_after	cvt_before	cvt_after	verdict
' > "$TSV"

count_cvt()     { find "$1/ttgir" -name '*.ttgir' -type f -exec cat {} + 2>/dev/null | grep -c convert_layout || true; }
count_kernels() { find "$1/ttgir" -name '*.ttgir' -type f 2>/dev/null | wc -l; }
count_rows()    { grep -c '^SUCCESS' "$1/run.log" 2>/dev/null || true; }

{
  echo ">>> vendor=$VENDOR  switch=$ENHANCE_VAR  ops=$OPS"
  for entry in $OPS; do
    op="${entry%%:*}"; rest="${entry#*:}"
    if [[ "$rest" == "$entry" ]]; then op_file="$op"; dtypes=""
    else
      op_file="${rest%%:*}"; d="${rest#*:}"
      [[ "$d" == "$rest" ]] && dtypes="" || dtypes="$d"
    fi
    BEFORE="$ROOT/${op}_before"; AFTER="$ROOT/${op}_after"

    echo "########## $op | $ENHANCE_VAR=0 before (record) ##########"
    OP="$op" OP_FILE="$op_file" OUT_DIR="$BEFORE" DTYPES="$dtypes" \
      env "$ENHANCE_VAR=0" bash run_pytest.sh
    rc_before=$?

    echo "########## $op | $ENHANCE_VAR=1 after (replay $BEFORE) ##########"
    OP="$op" OP_FILE="$op_file" OUT_DIR="$AFTER" DTYPES="$dtypes" \
      REPLAY_FROM="$BEFORE" env "$ENHANCE_VAR=1" bash run_pytest.sh
    rc_after=$?

    rb=$(count_rows "$BEFORE");     ra=$(count_rows "$AFTER")
    kb=$(count_kernels "$BEFORE");  ka=$(count_kernels "$AFTER")
    cb=$(count_cvt "$BEFORE");      ca=$(count_cvt "$AFTER")
    fb=$(ir_fingerprint "$BEFORE") || fb=""
    fa=$(ir_fingerprint "$AFTER")  || fa=""

    if (( rc_before != 0 || rc_after != 0 )) || [[ "${rb:-0}" == 0 || "${ra:-0}" == 0 ]]; then
      verdict=NO_ROWS
      note="跑不出测量行（rc=$rc_before/$rc_after）。看 run.log——多半是 dtype 不被该算子接受，或它在本后端拿不到 config"
    elif [[ -z "$fb" || -z "$fa" ]]; then
      verdict=NO_KERNEL
      note="没编出 .ttgir：该算子在本后端没走 Triton，编译器开关对它无从谈起。这不是'开关无作用'"
    elif [[ "$fb" == "$fa" ]]; then
      verdict=NO_IR_DIFF
      note="两侧 ttgir 逐字节相同（config 已锁死）：开关没改变生成代码。先查开关名拼写，再下结论"
    else
      verdict=IR_CHANGED
      note="生成代码有变（cvt ${cb}→${ca}，指纹 ${fb:0:8}/${fa:0:8}）：值得排六轮 A/B"
    fi
    printf '%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n' \
      "$op" "$op_file" "${dtypes:-<all>}" "${rb:-0}" "${ra:-0}" \
      "$kb" "$ka" "${cb:-0}" "${ca:-0}" "$verdict" >> "$TSV"
    echo ">>> $op: $verdict — $note"

    if [[ "$verdict" == NO_ROWS ]]; then
      echo "!!! $op 没有可对比的测量，不出 speedup 区间"
      continue
    fi
    echo "########## $op | summary ##########"
    python batch_summary.py --ab "$BEFORE/run.log" "$AFTER/run.log" --label "$op" \
      || echo "!!! summary failed: $op"
  done

  echo "########## 筛选汇总 ##########"
  awk -F'\t' 'NR>1{printf "  %-12s %-22s cvt %s->%s  kernels %s/%s  rows %s/%s\n", \
       $10,$1,$8,$9,$6,$7,$4,$5}' "$TSV" | sort
  echo ">>> 只把 IR_CHANGED 的算子送进 ab_repeat.sh。"
  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 ">>> 筛选表: $(pwd)/$TSV"
echo ">>> 每算子两份 run.log 在 $ROOT/<op>_{before,after}/ 下"