- evaluate: 每核命令数 = ceil(b_core/b0) * ceil(K/k_l1^m) (K截断退化为 b_core/b0, 数值不变; L1绑定消除 b0 倍多计 —— v1.1 §4.4 恒劣恒等式的 n_K 是未合并粒度, 误代入合并后段数会多计 b0 倍, 可把仲裁方向翻错) - beats_iterbatch: 泛化为实际命令数比较 (cmds_iter=b_core*ceil(K/k_l1_iter) vs cmds_mb=ceil(b_core/b0)*ceil(K/k_l1^m), 节省>T_cmd vs drain 惩罚); K截断时严格 退化为文档闭式 b_core > b0*(T_comp+T_write)/T_cmd; 截断判定改用合并口径 plan.k_l1>=K (未合并截断不代表合并后截断); 覆盖 dValue 512B cap 第三情形; T_cmd<=0 策略路径改为 cmds_mb<cmds_iter 判 MergeBatch 优先 - router: 仲裁文案 [裁决] 位打印最终胜者 (修复分界/时延不一致时的自相矛盾表述) - 用户 case 家族 B=128,M=1~16,N=128,K=512 修复后: m=1/2/4 -> MergeBatch, m=8/16 -> IterBatch (修复前全判 IterBatch; 交叉点 m≈4~8, 物理合理) - 测试: TestIssue35 回归 5 例 (命令数公式/K截断不变/口径一致/路由家族/裁决文案); test_beats_iterbatch_policy 的 (128,64,64,512) 期望 True->False (第三情形: 合并侧 dValue cap 截断, 命令数 4=4 打平, 恒劣 —— 原期望基于误分类) - docs/01_MergeBatch分支.md: 分界小节补第三情形行 + 命令数口径警示 + 泛化净收益式 - 验证: 68/68 unittest; examples 重生成可复现 0 diff (仅仲裁文案 + 16.0->16 格式, plans.csv 不变); 压力 10000 例 (seed7/6000+seed2024/4000): 0 崩溃/0 NaN/0 违规/ 0 不可行/0 GM<V_in, 七分支全覆盖
201 lines
11 KiB
Python
201 lines
11 KiB
Python
"""分支决策路由: 按决策树推导分支, 重叠区由端到端时延模型仲裁.
|
|
|
|
决策树 (v0.98 §3.3 + 各分支理论文档):
|
|
1. 前置归约: BatchA=1 或 BatchB=1 -> 转Matmul
|
|
K=0 / K=1 -> 特殊分支 (AIV 向量通路)
|
|
2. B >= C 且 BatchA==BatchB: 切B -> IterBatch 与 MergeBatch 仲裁
|
|
仲裁规则 (v1.1 §4.5 统一分界):
|
|
MergeBatch 最优 <=> K截断(k_L1=K) 且 b_core > b0*(T_comp+T_write)/T_cmd
|
|
L1 绑定时 MergeBatch 恒劣于 IterBatch;
|
|
两分支同时合法时用端到端时延模型 T_total 仲裁;
|
|
例外: T_cmd<=0 (命令时延不可量化/未标定) 时, 指令级收益未建模,
|
|
按既定策略: K截断即可优先 MergeBatch (覆盖时延模型仅来自 drain/冗余的差额).
|
|
3. StreamK 检查: P <= C/2 且满足切K条件 -> StreamK (B/M/N 买不满时买 K)
|
|
4. 兜底: ASW_Basic 切 M/N (含降核模式)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from .hardware import NpuSpec, ASCEND950PR
|
|
from .models import BmmCase, ImplPlan, HardwareTiming
|
|
from .branches.base import BranchResult
|
|
from .branches.merge_batch import MergeBatchBranch
|
|
from .branches.iter_batch import IterBatchBranch
|
|
from .branches.to_matmul import ToMatmulBranch
|
|
from .branches.special import SpecialBranch
|
|
from .branches.stream_k import StreamKBranch
|
|
from .branches.asw_basic import AswBasicBranch
|
|
|
|
|
|
class BranchRouter:
|
|
"""case -> 理论最优分支 + 方案 + 时延评估."""
|
|
|
|
def __init__(self, spec: NpuSpec = ASCEND950PR):
|
|
self.spec = spec
|
|
self.merge_batch = MergeBatchBranch(spec)
|
|
self.iter_batch = IterBatchBranch(spec)
|
|
self.to_matmul = ToMatmulBranch(spec)
|
|
self.special = SpecialBranch(spec)
|
|
self.stream_k = StreamKBranch(spec)
|
|
self.asw_basic = AswBasicBranch(spec)
|
|
|
|
# ------------------------------------------------------------------
|
|
def route(self, case: BmmCase) -> dict:
|
|
"""返回 {branch, plan, timing, arbitration, candidates}."""
|
|
s = self.spec
|
|
|
|
# 1) 前置归约: 转Matmul / 特殊分支
|
|
if case.batch_a == 1 or case.batch_b == 1:
|
|
r = self.to_matmul.analyze(case)
|
|
return self._wrap_checked(case, r, "BatchA=1或BatchB=1, 折叠转普通Matmul")
|
|
if case.k <= 1:
|
|
r = self.special.analyze(case)
|
|
note = "K=0纯写值" if case.k == 0 else "K=1逐元素乘, 走AIV向量通路"
|
|
# issue#4 P0: 特殊分支 capable=False (如 K=1 但 B<128 不满足 UB 乒乓) 时
|
|
# r.plan=None, 必须兜底而不能把 None 传给下游 —— 显式标注"该区域暂无理论方案"
|
|
if not r.capable or r.plan is None:
|
|
return self._no_plan(case, "特殊分支",
|
|
note + f"; 但进入条件不满足 ({r.failed_conditions()}), "
|
|
f"该区域暂无理论方案, 建议参考 Cube 兜底或 AIV 单缓冲")
|
|
return self._wrap_checked(case, r, note)
|
|
|
|
# 2) B >= C 且 BatchA==BatchB: IterBatch / MergeBatch
|
|
if case.batch_c >= s.aic_num and case.batch_a == case.batch_b:
|
|
return self._route_split_b(case)
|
|
|
|
# 3) StreamK: P <= C/2 且切K条件满足
|
|
sk = self.stream_k.analyze(case)
|
|
if sk.capable:
|
|
return self._wrap_checked(case, sk, f"P<=C/2, B/M/N并行度买不满, 切K (grid_K={sk.plan.grid_k})")
|
|
|
|
# 4) 兜底: ASW_Basic (含降核模式)
|
|
asw = self.asw_basic.analyze(case)
|
|
note = "ASW_Basic兜底"
|
|
if sk.checks and not sk.capable:
|
|
note += f" (StreamK未过: {sk.failed_conditions()})"
|
|
return self._wrap_checked(case, asw, note)
|
|
|
|
# ------------------------------------------------------------------
|
|
def _route_split_b(self, case: BmmCase) -> dict:
|
|
mb = self.merge_batch.analyze(case)
|
|
ib = self.iter_batch.analyze(case)
|
|
|
|
# 候选表: [(分支名, BranchResult)], 顺序 = 仲裁优先级
|
|
cand_map = {self.merge_batch.name: mb, self.iter_batch.name: ib}
|
|
capable = {n: r.capable for n, r in cand_map.items()}
|
|
|
|
arbitration = ""
|
|
if mb.capable and ib.capable:
|
|
# 统一分界条件 + 端到端时延仲裁双保险
|
|
mb_win, detail = self.merge_batch.beats_iterbatch(case)
|
|
t_mb = mb.timing.t_total
|
|
t_ib = ib.timing.t_total
|
|
lat_win = self.merge_batch.name if t_mb <= t_ib else self.iter_batch.name
|
|
win = self.merge_batch.name if mb_win else self.iter_batch.name
|
|
# 冲突解决: 默认"时延模型为最终裁决"; 例外是 T_cmd<=0 且分界条件判
|
|
# MergeBatch 胜 (K截断) 的情形 —— 此时时延模型不含指令级收益
|
|
# (MergeBatch 搬移命令数/主机指令数省 b0 倍, 未量化), 按既定策略
|
|
# 优先 MergeBatch (时延模型内的差额只是 drain 惩罚/冗余, 方向已知小量).
|
|
policy_merge = (self.spec.t_cmd <= 0 and mb_win and
|
|
win != lat_win and win == self.merge_batch.name)
|
|
if policy_merge:
|
|
arbitration = (
|
|
f"两分支均合法, 仲裁: "
|
|
f"[分界条件] MergeBatch最优={mb_win} ({detail}); "
|
|
f"[时延模型] T_MergeBatch={t_mb*1e6:.2f}us vs T_IterBatch={t_ib*1e6:.2f}us -> {lat_win}更优; "
|
|
f"[裁决] {win} (T_cmd<=0 策略: 命令/指令级收益未建模, 时延模型差异仅来自 "
|
|
f"drain/冗余, 以 MergeBatch 优先策略裁决)")
|
|
else:
|
|
arbitration = (
|
|
f"两分支均合法, 仲裁: "
|
|
f"[分界条件] MergeBatch最优={mb_win} ({detail}); "
|
|
f"[时延模型] T_MergeBatch={t_mb*1e6:.2f}us vs T_IterBatch={t_ib*1e6:.2f}us -> {lat_win}更优; "
|
|
f"[裁决] {lat_win}" + ("" if win == lat_win else
|
|
f" (分界条件判{win}, 与时延模型不一致, 以时延模型为准)")
|
|
)
|
|
win = lat_win # 时延模型为最终裁决
|
|
elif any(capable.values()):
|
|
win = next(n for n, v in capable.items() if v)
|
|
arbitration = f"仅 {win} 条件满足"
|
|
else:
|
|
# 切B分支都不满足, 尝试 StreamK 再回落 ASW
|
|
sk = self.stream_k.analyze(case)
|
|
if sk.capable:
|
|
return self._wrap_checked(case, sk, "切B分支条件不满足, 落 StreamK")
|
|
asw = self.asw_basic.analyze(case)
|
|
return self._wrap_checked(
|
|
case, asw,
|
|
f"IterBatch/MergeBatch 进入条件均不满足, 回落 ASW_Basic; "
|
|
f"IterBatch未过: {ib.failed_conditions()}; "
|
|
f"MergeBatch未过: {mb.failed_conditions()}")
|
|
|
|
# 可行性保障 (issue#13): 仲裁胜出方案必须通过约束自检, 否则按
|
|
# (另一切B候选 -> StreamK -> ASW_Basic) 顺序回退到首个可行方案.
|
|
from .constraints import check_plan_constraints
|
|
|
|
def _feasible(n):
|
|
r = cand_map[n]
|
|
return r.plan is not None and not check_plan_constraints(case, r.plan, self.spec)
|
|
|
|
if _feasible(win):
|
|
chosen = cand_map[win]
|
|
else:
|
|
loser = self.merge_batch.name if win == self.iter_batch.name else self.iter_batch.name
|
|
fallback_note = (f"; 但 {win} 方案自检违规: "
|
|
f"{'; '.join(check_plan_constraints(case, cand_map[win].plan, self.spec))}")
|
|
if capable.get(loser) and _feasible(loser):
|
|
chosen, win = cand_map[loser], loser
|
|
fallback_note += f", 回退可行候选 {loser}"
|
|
else:
|
|
sk = self.stream_k.analyze(case)
|
|
if sk.capable and sk.plan is not None and \
|
|
not check_plan_constraints(case, sk.plan, self.spec):
|
|
return self._wrap_checked(case, sk, arbitration + fallback_note + ", 落 StreamK")
|
|
asw = self.asw_basic.analyze(case)
|
|
if asw.plan is not None and not check_plan_constraints(case, asw.plan, self.spec):
|
|
return self._wrap_checked(case, asw, arbitration + fallback_note + ", 回落 ASW_Basic")
|
|
chosen, win = cand_map[win], win # 无可行方案: 保留原裁决, 由自检标注
|
|
arbitration += fallback_note
|
|
|
|
result = BranchResult(capable=True, plan=chosen.plan, timing=chosen.timing)
|
|
return self._wrap_checked(case, result, arbitration,
|
|
candidates=capable)
|
|
|
|
# ------------------------------------------------------------------
|
|
def _wrap_checked(self, case: BmmCase, result: BranchResult, note: str,
|
|
candidates: dict | None = None) -> dict:
|
|
"""生成后自检 (issue#5): 推荐方案必须通过统一约束源校验, 不可行则标注违规.
|
|
|
|
约束源与 evaluate 共用 constraints.check_plan_constraints, 保证
|
|
"推荐方案 vs 自带评估器" 口径一致, 不再出现生成说可行、校验说不可行的矛盾.
|
|
"""
|
|
from .constraints import check_plan_constraints
|
|
violations = check_plan_constraints(case, result.plan, self.spec) if result.plan else []
|
|
if violations:
|
|
note = (note + " [自检违规: " + "; ".join(violations) +
|
|
"] —— 方案生成存在缺陷, 需人工复核")
|
|
# issue#34: 兜底分支 (ASW) 效率下限不满足时降级标注 (warning), 不判违规
|
|
if result.plan is not None and "效率降级" in result.plan.note:
|
|
note += " [效率降级标注: 搬移效率低于模型假设, 时延可能低估, 见 plan.note]"
|
|
return {
|
|
"branch": result.plan.branch if result.plan else "未知",
|
|
"plan": result.plan,
|
|
"timing": result.timing,
|
|
"arbitration": note + (f" | {result.note}" if result.note else ""),
|
|
"candidates": candidates or {},
|
|
"self_check_violations": violations,
|
|
}
|
|
|
|
@staticmethod
|
|
def _no_plan(case: BmmCase, branch: str, note: str) -> dict:
|
|
"""兜底: 分支 capable=False 时给出最小占位方案, 保证下游不崩溃 (issue#4).
|
|
|
|
方案标注 used_core_num=0 + 分支名, arbitration 说明"该区域暂无理论方案",
|
|
不产生时延 (timing=None), 供上层跳过或人工处理.
|
|
"""
|
|
plan = ImplPlan(case_id=case.case_id, branch=branch,
|
|
used_core_num=0, note="该区域暂无理论方案(进入条件不满足)")
|
|
return {"branch": branch, "plan": plan, "timing": None,
|
|
"arbitration": "[无方案] " + note, "candidates": {},
|
|
"self_check_violations": []}
|