#!/bin/bash
# 第一步：筛掉不值得跑 A/B 的算子。
#
# 为什么要单独一步：S5000 实测，筛选一个算子约 0.8 分钟，六轮 A/B 一个算子
# 13~44 分钟——差 20~50 倍。21 个算子筛一遍 15 分钟，挡掉 13 个，把昂贵队列
# 从 7 小时压到 2.7 小时。筛选的全部价值就在这个比例上。
#
# 口径与 A/B 完全一致，因为筛选也是一次 A/B：A 侧 record，B 侧 REPLAY_FROM 它。
# 两侧各自 record 是不行的——各自 sweep 会选出不同 config，生成的 IR 自然不同，
# 于是每个算子都被判成"开关改变了生成代码"，筛选等于没筛。这条本仓踩过。
#
# 判据是 ttgir 指纹，不是 convert_layout 计数。cvt 数相同不等于生成代码相同，
# 而 A/B 要测的是"生成代码变了没有"。cvt 数照样记下来，它解释的是变化的性质。
#
# 用法：
#   SWITCH=FLAGTREE_MUSA_RLC_ENHANCE OPS="col2im:col2im topk:topk perf_fft:fft" \
#     ZL_BACKEND=musa FLAGGEMS_DIR=/path/to/FlagGems bash screen.sh
#
#   OPS 元素为 OP[:OP_FILE[:DTYPES]]，省略 OP_FILE 时等于 OP。
#   DTYPES 留空表示用该算子自己的全套——**不要全局写死一个 dtype**：
#   fft 只吃 complex64，被一刀切的 DTYPES=float16 直接判死过一次。
#
# 产物：runs/screen_<开关>_<时间戳>/，含 screen.tsv 与逐算子 run 目录。
set -uo pipefail

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

: "${SWITCH:?SWITCH is required, e.g. FLAGTREE_MUSA_RLC_ENHANCE}"
OPS="${OPS:-}"
: "${OPS:?OPS is required, e.g. \"topk:topk perf_fft:fft\"}"
CELL_TIMEOUT="${CELL_TIMEOUT:-2400}"

ROOT="runs/screen_${SWITCH}_$(date +%Y%m%d-%H%M%S)"
mkdir -p "$ROOT"
TSV="$ROOT/screen.tsv"
LOG="$ROOT/screen.log"
printf 'op\top_file\tdtypes\trows_a\trows_b\tkernels_a\tkernels_b\tcvt_a\tcvt_b\tverdict\n' > "$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 ">>> switch=$SWITCH  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

    echo "########## $op | A ($SWITCH=0, record) ##########"
    OP="$op" OP_FILE="$op_file" OUT_DIR="$ROOT/${op}_A" DTYPES="$dtypes" \
      timeout -k 30 "$CELL_TIMEOUT" env "$SWITCH=0" bash run_pytest.sh
    rc_a=$?

    echo "########## $op | B ($SWITCH=1, replay A) ##########"
    OP="$op" OP_FILE="$op_file" OUT_DIR="$ROOT/${op}_B" DTYPES="$dtypes" \
      REPLAY_FROM="$ROOT/${op}_A" \
      timeout -k 30 "$CELL_TIMEOUT" env "$SWITCH=1" bash run_pytest.sh
    rc_b=$?

    ra=$(count_rows "$ROOT/${op}_A"); rb=$(count_rows "$ROOT/${op}_B")
    ka=$(count_kernels "$ROOT/${op}_A"); kb=$(count_kernels "$ROOT/${op}_B")
    ca=$(count_cvt "$ROOT/${op}_A"); cb=$(count_cvt "$ROOT/${op}_B")
    fa=$(ir_fingerprint "$ROOT/${op}_A") || fa=""
    fb=$(ir_fingerprint "$ROOT/${op}_B") || fb=""

    if (( rc_a != 0 || rc_b != 0 )) || [[ "${ra:-0}" == 0 || "${rb:-0}" == 0 ]]; then
      verdict="NO_ROWS"
      note="跑不出测量行（rc=$rc_a/$rc_b）。先看 console，多半是 dtype 不被该算子接受，或该算子在本后端拿不到 config"
    elif [[ -z "$fa" || -z "$fb" ]]; then
      verdict="NO_KERNEL"
      note="没编出 .ttgir：该算子在本后端没走 Triton。编译器开关对它无从谈起——这不是'开关无作用'"
    elif [[ "$fa" == "$fb" ]]; then
      verdict="NO_IR_DIFF"
      note="两侧 ttgir 逐字节相同（config 已锁死）：开关没改变该算子的生成代码，跑 A/B 只会测到噪声"
    else
      verdict="IR_CHANGED"
      note="生成代码有变（cvt ${ca}→${cb}，指纹 ${fa:0:8}/${fb: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>}" "${ra:-0}" "${rb:-0}" \
      "$ka" "$kb" "${ca:-0}" "${cb:-0}" "$verdict" >> "$TSV"
    echo ">>> $op: $verdict — $note"
    rm -rf "$ROOT/${op}_A/.triton_cache" "$ROOT/${op}_B/.triton_cache"
  done

  echo "########## 汇总 ##########"
  awk -F'\t' 'NR>1{print "  " $10 "\t" $1 "\t(cvt " $8 "→" $9 ", kernels " $6 "/" $7 ", rows " $4 "/" $5 ")"}' "$TSV" | sort
  echo
  echo ">>> 只把 IR_CHANGED 的算子送进 ab_repeat.sh。"
  echo ">>> NO_IR_DIFF 跑六轮只会测出噪声并报成一个'无差异'结论。"
  echo ">>> NO_KERNEL / NO_ROWS 要先修可运行性，不要当成开关的结论。"
} 2>&1 | tee "$LOG"

sed -i -E $'s/\x1b\\[[0-9;]*[A-Za-z]//g' "$LOG"
echo ">>> 结果表: $(pwd)/$TSV"
echo ">>> 总 log:  $(pwd)/$LOG"