bench: partition operators across cards by measured cost, and rescore in place

Two pieces the parallel campaign needs.

partition.py assigns screened operators to cards longest-first. run_batch.py's
existing hash partition ignores weight: on this batch it puts the critical path
at 120 minutes and leaves a card idle, where greedy gets 70. The floor is set
by the single heaviest operator — topk alone is 70.6 minutes against an ideal
share of 45.4 — so partitioning cannot beat it, only keep the other cards from
waiting. Weights come from screen.tsv's elapsed_before + elapsed_after, which
is one pair, i.e. 2 cells; N rounds is 2N cells, so the estimate is N times the
weight. Getting that wrong reported topk as 11.8 minutes.

RESCORE=<root> re-judges an existing screening root against the current
fingerprint and rewrites its screen.tsv, keeping the original alongside. The
criterion has already changed twice — index.tsv used to be hashed, and the
switch's own ttg.rlc-* annotations had to be stripped — and the second change
moved 24 operators from 23 IR_CHANGED to 16. Rescoring the artifacts already on
disk takes seconds; re-running that screen takes 79 minutes.
This commit is contained in:
Aharrypotter committed 2026-09-10 19:25:34 +08:00
1 parent c6a7f834cb
commit 03fdcbab28
3 files changed
+259

No files matched your search

+139
View File
@@ -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 输出每卡一行 `<device>\\t<OPS>`,供驱动脚本读")
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())