diff --git a/ab_rlc_test.sh b/ab_rlc_test.sh new file mode 100644 index 0000000..270e9f3 --- /dev/null +++ b/ab_rlc_test.sh @@ -0,0 +1,68 @@ +#!/bin/bash +# A/B 对比:FlagTree RemoveLayoutConversions 增强关 vs 开,逐算子两轮。 +# +# 口径:每个算子先跑关(record autotune config),再用 REPLAY_FROM 回放同一套 +# config 跑开,两侧只差 FLAGTREE__RLC_ENHANCE 这一个开关。产物沿用 +# run_pytest.sh 的目录结构,所以每个算子恰好留下两份 run.log: +# +# runs/_<时间戳>/run.log 关(before) +# runs/_<时间戳>/run.log 开(after) +# +# 跑完用 batch_summary.py --ab 得到该算子的 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/_<时间戳>/ 下,按时间先后即 before / after" diff --git a/batch_summary.py b/batch_summary.py index 3092ecf..6d0d9b6 100644 --- a/batch_summary.py +++ b/batch_summary.py @@ -165,12 +165,82 @@ def summarize(batch_dir: Path, out_path: Optional[Path] = None) -> Path: return out_path +def ab_rows(before_text: str, after_text: str) -> List[Dict[str, str]]: + """把两次 run.log 的测量行按 (op, dtype, size_detail) 配对。 + + speedup 列是 gems 对 torch,不是前后对比,所以这里用 gems_lat 自己算 + before/after - 1。只有两侧都 SUCCESS 且延迟为正的行才计入;对不上的行 + 单独返回,避免用不完整的配对报一个看起来完整的区间。 + """ + def index(text: str) -> Dict[Tuple[str, str, str], float]: + out: Dict[Tuple[str, str, str], float] = {} + for row in parse_run_log(text): + if row.get("row_status") != "SUCCESS": + continue + try: + lat = float(row.get("gems_lat", "")) + except ValueError: + continue + if lat > 0: + out[(row.get("fg_op_name", ""), row.get("dtype", ""), + row.get("size_detail", ""))] = lat + return out + + before, after = index(before_text), index(after_text) + paired: List[Dict[str, str]] = [] + for key in sorted(before.keys() & after.keys()): + op, dtype, size = key + b, a = before[key], after[key] + paired.append({ + "fg_op_name": op, "dtype": dtype, "size_detail": size, + "before_lat": f"{b:.6f}", "after_lat": f"{a:.6f}", + "speedup_percent": f"{(b / a - 1.0) * 100.0:.2f}", + }) + return paired + + +def report_ab(before: Path, after: Path, label: str = "") -> List[Dict[str, str]]: + """打印一个算子的前后对比:逐行 speedup 与整体区间。""" + before_text = before.read_text(encoding="utf-8", errors="replace") + after_text = after.read_text(encoding="utf-8", errors="replace") + rows = ab_rows(before_text, after_text) + name = label or (rows[0]["fg_op_name"] if rows else "?") + if not rows: + print(f"[ab] {name}: 没有可配对的测量行" + f"(before={before} after={after})") + return rows + values = [float(r["speedup_percent"]) for r in rows] + for r in rows: + print(f" {r['dtype']:<16} {float(r['speedup_percent']):+7.2f}% " + f"{r['before_lat']} -> {r['after_lat']} {r['size_detail']}") + fallback = "" + if DIRECT_NO_CUDAGRAPH_RE.search(before_text) or \ + DIRECT_NO_CUDAGRAPH_RE.search(after_text): + # 回退行用的是 do_bench 口径,与 cudagraph 行不同表可比性 + fallback = " [含 BENCHMARK_DIRECT_NO_CUDAGRAPH 回退行]" + print(f"[ab] {name}: {len(rows)} 行配对,speedup " + f"{min(values):+.2f}% ~ {max(values):+.2f}%{fallback}") + return rows + + def main() -> None: p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) - p.add_argument("batch_dir", type=Path, help="run_batch.py 的产物目录") + p.add_argument("batch_dir", type=Path, nargs="?", + help="run_batch.py 的产物目录") p.add_argument("-o", "--output", type=Path, default=None, help="输出 CSV(默认 /summary.csv)") + p.add_argument("--ab", nargs=2, type=Path, metavar=("BEFORE", "AFTER"), + help="两份 run.log,按 shape 配对输出该算子的 speedup 区间") + p.add_argument("--label", default="", help="--ab 报告里的算子名") args = p.parse_args() + if args.ab: + for path in args.ab: + if not path.is_file(): + raise SystemExit(f"run.log 不存在: {path}") + report_ab(args.ab[0], args.ab[1], args.label) + return + if args.batch_dir is None: + raise SystemExit("需要 batch_dir,或用 --ab BEFORE AFTER") if not args.batch_dir.is_dir(): raise SystemExit(f"目录不存在: {args.batch_dir}") out = summarize(args.batch_dir, args.output)