132 lines
5.6 KiB
Python
132 lines
5.6 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 仲裁.
|
|
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(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向量通路"
|
|
return self._wrap(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(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(asw, note)
|
|
|
|
# ------------------------------------------------------------------
|
|
def _route_split_b(self, case: BmmCase) -> dict:
|
|
mb = self.merge_batch.analyze(case)
|
|
ib = self.iter_batch.analyze(case)
|
|
|
|
candidates = []
|
|
if mb.capable:
|
|
candidates.append((self.merge_batch.name, mb))
|
|
if ib.capable:
|
|
candidates.append((self.iter_batch.name, ib))
|
|
|
|
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
|
|
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}" + ("" if win == lat_win else f" (分界条件与时延模型不一致, 以时延模型为准: {lat_win})")
|
|
)
|
|
if win != lat_win:
|
|
win = lat_win # 时延模型为最终裁决
|
|
elif candidates:
|
|
win = candidates[0][0]
|
|
arbitration = f"仅 {win} 条件满足"
|
|
else:
|
|
# 切B分支都不满足, 尝试 StreamK 再回落 ASW
|
|
sk = self.stream_k.analyze(case)
|
|
if sk.capable:
|
|
return self._wrap(sk, "切B分支条件不满足, 落 StreamK")
|
|
asw = self.asw_basic.analyze(case)
|
|
return self._wrap(
|
|
asw,
|
|
f"IterBatch/MergeBatch 进入条件均不满足, 回落 ASW_Basic; "
|
|
f"IterBatch未过: {ib.failed_conditions()}; "
|
|
f"MergeBatch未过: {mb.failed_conditions()}")
|
|
|
|
chosen = mb if win == self.merge_batch.name else ib
|
|
return {
|
|
"branch": win,
|
|
"plan": chosen.plan,
|
|
"timing": chosen.timing,
|
|
"arbitration": arbitration,
|
|
"candidates": {n: r.capable for n, r in
|
|
[(self.merge_batch.name, mb), (self.iter_batch.name, ib)]},
|
|
}
|
|
|
|
# ------------------------------------------------------------------
|
|
@staticmethod
|
|
def _wrap(result: BranchResult, note: str) -> dict:
|
|
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": {},
|
|
}
|