diff --git a/ab_rlc_test.sh b/ab_rlc_test.sh index 3c8152f..2448ace 100644 --- a/ab_rlc_test.sh +++ b/ab_rlc_test.sh @@ -39,6 +39,51 @@ set -uo pipefail cd "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" source "./_ir_fingerprint.sh" +# RESCORE=<已有的 root>:判据改了之后,用当前的 ir_fingerprint 对盘上已有的 ttgir +# 重新打分并改写该 root 的 screen.tsv,不重跑任何测量。 +# +# 需要它是因为判据确实改过:_ir_fingerprint.sh 起初把 index.tsv 也算进哈希,后来又 +# 发现要抹掉开关自己盖的 ttg.rlc-* 属性。同一批 24 个算子,结论从 23 个 IR_CHANGED +# 变成 16 个。重打分是秒级的,重跑那批筛选要 79 分钟。 +# +# 原表存为 screen.tsv.before-rescore,判定有变化的算子逐个打出来。 +RESCORE="${RESCORE:-}" +if [[ -n "$RESCORE" ]]; then + [[ -d "$RESCORE" ]] || { echo "!!! 目录不存在: $RESCORE" >&2; exit 2; } + TSV="$RESCORE/screen.tsv" + [[ -f "$TSV" ]] || { echo "!!! 没有 screen.tsv: $TSV" >&2; exit 2; } + cp "$TSV" "$TSV.before-rescore" + + hdr=$(head -1 "$TSV.before-rescore") + # 列名 -> 序号,不写死列序 + idx() { echo "$hdr" | tr '\t' '\n' | grep -nxF "$1" | cut -d: -f1; } + c_op=$(idx op); c_rb=$(idx rows_before); c_ra=$(idx rows_after); c_v=$(idx verdict) + [[ -n "$c_op" && -n "$c_v" ]] || { echo "!!! screen.tsv 缺 op/verdict 列" >&2; exit 2; } + + printf '%s\n' "$hdr" > "$TSV" + changed=0 + while IFS= read -r line; do + op=$(printf '%s' "$line" | cut -f"$c_op") + old=$(printf '%s' "$line" | cut -f"$c_v") + rb=$(printf '%s' "$line" | cut -f"${c_rb:-0}"); ra=$(printf '%s' "$line" | cut -f"${c_ra:-0}") + fa=$(ir_fingerprint "$RESCORE/${op}_before") || fa="" + fb=$(ir_fingerprint "$RESCORE/${op}_after") || fb="" + if [[ "${rb:-0}" == 0 || "${ra:-0}" == 0 ]]; then v=NO_ROWS + elif [[ -z "$fa" || -z "$fb" ]]; then v=NO_TRITON + elif [[ "$fa" == "$fb" ]]; then v=NO_IR_DIFF + else v=IR_CHANGED; fi + printf '%s\n' "$line" | awk -v FS='\t' -v OFS='\t' -v i="$c_v" -v v="$v" '{$i=v; print}' >> "$TSV" + if [[ "$old" != "$v" ]]; then + printf ' %-22s %-12s -> %s\n' "$op" "$old" "$v" >&2 + changed=$((changed+1)) + fi + done < <(tail -n +2 "$TSV.before-rescore") + + echo ">>> 重打分完成,$changed 个算子的判定变了;原表存为 $TSV.before-rescore" + awk -F'\t' 'NR>1{c[$NF]++} END{for (k in c) printf(" %-12s %d\n", k, c[k])}' "$TSV" + exit 0 +fi + VENDOR="${VENDOR:-MUSA}" ENHANCE_VAR="FLAGTREE_${VENDOR}_RLC_ENHANCE" diff --git a/partition.py b/partition.py new file mode 100644 index 0000000..8eda301 --- /dev/null +++ b/partition.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python3 +"""把筛选通过的算子按实测耗时分到各张卡上。 + +为什么不用 `run_batch.py` 的 hash 分片:它按算子名的 hash 取模,不看重量。S5000 上 +17 个算子的六轮 campaign,hash 分片关键路径 120 分钟、有一张卡空转;按耗时贪心是 +70 分钟。同一批活,差 50 分钟。 + +贪心(最长优先塞进当前最空的卡)在这里够用,也不可能更好:`topk` 一个算子就要 70.6 +分钟,而理想均分是 45.4——**关键路径的下界由最重的那个算子决定,加卡不解决**。分片 +只能保证除它之外的卡不空等。 + +权重来自 `screen.tsv` 的 `elapsed_before + elapsed_after`(**一对 = 2 格**的秒数)。 +N 轮是 2N 格,所以预估耗时是权重的 N 倍——不是 N/6 倍,这里错过一次,把 topk 的 70.6 +分钟报成了 11.8。缺列时退化成等权重,并在 stderr 说明——等权重的分片仍然 +比不分好,但别让人以为它按耗时排过。 + +用法: + python partition.py --screen runs/ab_rlc_XXX/screen.tsv --devices 0,1,2,3 + python partition.py --screen ... --devices 0,1 --verdict IR_CHANGED --format shell +""" + +from __future__ import annotations + +import argparse +import csv +import sys +from pathlib import Path +from typing import Dict, List, Optional, Sequence, Tuple + +Item = Tuple[str, float] + + +def plan(items: Sequence[Item], slots: Sequence[str]) -> Dict[str, List[Item]]: + """最长优先贪心:每次把最重的活放进当前总量最小的卡。 + + 卡少于 1 张是调用方的错,直接抛;活比卡少时后面的卡就是空的,这不是错误。 + """ + if not slots: + raise ValueError("需要至少一张卡") + out: Dict[str, List[Item]] = {s: [] for s in slots} + load: Dict[str, float] = {s: 0.0 for s in slots} + # 同重量时按名字排,保证同一份输入每次得到同一个方案 + for name, weight in sorted(items, key=lambda kv: (-kv[1], kv[0])): + target = min(slots, key=lambda s: (load[s], s)) + out[target].append((name, weight)) + load[target] += weight + return out + + +def imbalance(assignment: Dict[str, List[Item]]) -> float: + """最慢的卡 / 理想均分。1.0 是完美;下界由最重的单个算子决定。""" + loads = [sum(w for _, w in v) for v in assignment.values()] + total = sum(loads) + if total <= 0: + return 1.0 + return max(loads) / (total / len(loads)) + + +def read_screen(path: Path, verdict: str = "IR_CHANGED" + ) -> Tuple[List[Item], Dict[str, str], List[str]]: + """(算子, 一对耗时秒), op -> op_file, 被判定挡掉的算子。""" + with path.open(newline="", encoding="utf-8") as f: + rows = list(csv.DictReader(f, delimiter="\t")) + if not rows: + raise SystemExit(f"筛选表没有数据行: {path}") + items: List[Item] = [] + files: Dict[str, str] = {} + skipped: List[str] = [] + warned = False + for r in rows: + op = r.get("op", "") + if not op: + continue + if verdict and r.get("verdict") != verdict: + skipped.append(f"{op}({r.get('verdict', '?')})") + continue + try: + w = float(r["elapsed_before"]) + float(r["elapsed_after"]) + except (KeyError, TypeError, ValueError): + w = 1.0 + if not warned: + print("!!! 筛选表没有 elapsed_* 列,退化成等权重分片——" + "分片仍然比不分好,但它没有按耗时排过", file=sys.stderr) + warned = True + items.append((op, w)) + files[op] = r.get("op_file") or op + return items, files, skipped + + +def main(argv: Optional[List[str]] = None) -> int: + p = argparse.ArgumentParser(description=__doc__.splitlines()[0], + formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--screen", type=Path, required=True, help="ab_rlc_test.sh 的 screen.tsv") + p.add_argument("--devices", required=True, help="逗号分隔的设备列表,如 0,1,2,3") + p.add_argument("--verdict", default="IR_CHANGED", + help="只取这个判定的算子;空字符串表示全要") + # 权重是一对(2 格)的秒数;N 轮是 2N 格,所以预估耗时是权重的 N 倍 + p.add_argument("--rounds", type=int, default=6, help="仅用于打印预估耗时") + p.add_argument("--format", choices=("human", "shell"), default="human", + help="shell 输出每卡一行 `\\t`,供驱动脚本读") + args = p.parse_args(argv) + + if not args.screen.is_file(): + raise SystemExit(f"筛选表不存在: {args.screen}") + devices = [d.strip() for d in args.devices.split(",") if d.strip()] + if not devices: + raise SystemExit("--devices 为空") + + items, files, skipped = read_screen(args.screen, args.verdict) + if not items: + raise SystemExit(f"筛选表里没有 verdict={args.verdict} 的算子——" + "先确认筛选跑过,且判据没有把所有算子都挡掉") + + assignment = plan(items, devices) + if args.format == "shell": + for d in devices: + ops = " ".join(f"{op}:{files[op]}" for op, _ in assignment[d]) + print(f"{d}\t{ops}") + return 0 + + total = sum(w for _, w in items) * args.rounds + print(f">>> {len(items)} 个算子 / {len(devices)} 张卡 / {args.rounds} 轮") + if skipped: + print(f">>> 判定挡掉 {len(skipped)} 个: {', '.join(skipped)}") + for d in devices: + load = sum(w for _, w in assignment[d]) * args.rounds + names = " ".join(op for op, _ in assignment[d]) or "(空)" + print(f" 卡{d} {load/60:6.1f} 分钟 {names}") + crit = max(sum(w for _, w in v) for v in assignment.values()) * args.rounds + heaviest = max(items, key=lambda kv: kv[1]) + print(f">>> 单卡串行 {total/3600:.2f} 小时;{len(devices)} 卡关键路径 " + f"{crit/60:.1f} 分钟,失衡 {imbalance(assignment):.2f}x") + print(f">>> 下界由最重的算子决定:{heaviest[0]} 单独就要 " + f"{heaviest[1]*args.rounds/60:.1f} 分钟,加卡不解决") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/test_backends.py b/test_backends.py index d8bd9dc..6867f57 100644 --- a/test_backends.py +++ b/test_backends.py @@ -659,5 +659,80 @@ class ForeignVendorModuleTests(unittest.TestCase): self.assertEqual(foreign_vendor_modules("mthreads", mods), []) +class PartitionTests(unittest.TestCase): + """按耗时把算子分到各卡。hash 分片在这批活上比它慢 50 分钟。""" + + HEADER = "op\top_file\telapsed_before\telapsed_after\tverdict" + + def _load(self, assignment): + return sorted(sum(w for _, w in v) for v in assignment.values()) + + def _screen(self, tmp, rows, header=None): + p = Path(tmp, "screen.tsv") + p.write_text((header or self.HEADER) + "\n" + "\n".join(rows) + "\n", + encoding="utf-8") + return p + + def test_heavy_first_keeps_cards_even(self): + from partition import plan + items = [("a", 10), ("b", 9), ("c", 8), ("d", 7), ("e", 6), ("f", 5)] + self.assertEqual(self._load(plan(items, ["0", "1", "2"])), [15, 15, 15]) + + def test_one_item_heavier_than_the_ideal_share_sets_the_floor(self): + # topk 单个 70 分钟、理想均分 45——加卡不解决,分片只保证别人不空等 + from partition import imbalance, plan + items = [("topk", 70)] + [(f"op{i}", 5) for i in range(10)] + a = plan(items, [str(i) for i in range(8)]) + self.assertEqual(max(sum(w for _, w in v) for v in a.values()), 70) + self.assertGreater(imbalance(a), 1.0) + + def test_more_cards_than_operators_leaves_some_empty(self): + from partition import plan + a = plan([("a", 1), ("b", 1)], ["0", "1", "2", "3"]) + self.assertEqual(sum(1 for v in a.values() if not v), 2) + + def test_same_input_gives_the_same_plan(self): + from partition import plan + items = [("a", 5), ("b", 5), ("c", 5), ("d", 5)] + first = plan(items, ["0", "1"]) + for _ in range(5): + self.assertEqual(plan(list(reversed(items)), ["0", "1"]), first) + + def test_no_cards_is_the_callers_error(self): + from partition import plan + with self.assertRaises(ValueError): + plan([("a", 1)], []) + + def test_perfect_balance_is_one(self): + from partition import imbalance, plan + self.assertEqual(imbalance(plan([("a", 5), ("b", 5)], ["0", "1"])), 1.0) + + def test_verdict_filters_and_reports_what_it_dropped(self): + from partition import read_screen + with tempfile.TemporaryDirectory() as tmp: + f = self._screen(tmp, ["a\ta\t10\t10\tIR_CHANGED", + "b\tbb\t5\t5\tNO_IR_DIFF", + "c\tc\t1\t1\tNO_TRITON"]) + items, files, skipped = read_screen(f) + self.assertEqual(items, [("a", 20.0)]) + self.assertEqual(files["a"], "a") + self.assertEqual(len(skipped), 2) + + def test_op_file_defaults_to_op_name(self): + from partition import read_screen + with tempfile.TemporaryDirectory() as tmp: + f = self._screen(tmp, ["a\t\t1\t1\tIR_CHANGED"]) + _, files, _ = read_screen(f) + self.assertEqual(files["a"], "a") + + def test_missing_elapsed_falls_back_to_equal_weights(self): + from partition import read_screen + with tempfile.TemporaryDirectory() as tmp: + f = self._screen(tmp, ["a\ta\tIR_CHANGED", "b\tb\tIR_CHANGED"], + header="op\top_file\tverdict") + items, _, _ = read_screen(f) + self.assertEqual([w for _, w in items], [1.0, 1.0]) + + if __name__ == "__main__": unittest.main()