112 lines
5.3 KiB
Python
112 lines
5.3 KiB
Python
"""方案评估器: 对 (case, plan) 做硬件约束校验 + 时延评估 + 瓶颈分析.
|
|
|
|
评估模式入口: 用户自带实现方案 (ImplPlan), 软件评估其在 NPU 上的
|
|
各级硬件时延、流水情况与瓶颈, 并做可行性校验.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from .hardware import NpuSpec, ASCEND950PR
|
|
from .models import BmmCase, ImplPlan, EvalResult
|
|
from .timing import bound_type_of
|
|
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
|
|
|
|
_BRANCH_EVAL = {
|
|
"MergeBatch": MergeBatchBranch,
|
|
"IterBatch": IterBatchBranch,
|
|
"转Matmul": ToMatmulBranch,
|
|
"特殊分支": SpecialBranch,
|
|
"StreamK": StreamKBranch,
|
|
"ASW_Basic": AswBasicBranch,
|
|
"ASW_Basic_降核": AswBasicBranch,
|
|
}
|
|
|
|
|
|
class PlanEvaluator:
|
|
def __init__(self, spec: NpuSpec = ASCEND950PR):
|
|
self.spec = spec
|
|
|
|
def evaluate(self, case: BmmCase, plan: ImplPlan) -> EvalResult:
|
|
res = EvalResult(case=case, plan=plan)
|
|
|
|
# 1) 硬件约束校验
|
|
violations = self._check_constraints(case, plan)
|
|
res.feasible = not violations
|
|
res.violations = "; ".join(violations)
|
|
|
|
# 2) 时延评估 (按方案分支调用对应模型)
|
|
branch_cls = _BRANCH_EVAL.get(plan.branch)
|
|
if branch_cls is None:
|
|
res.advice = (f"分支 {plan.branch!r} 的评估模型待后续迭代; "
|
|
f"当前支持: {sorted(_BRANCH_EVAL)}")
|
|
return res
|
|
timing = branch_cls(self.spec).evaluate(case, plan)
|
|
res.timing = timing
|
|
res.bound_type = bound_type_of(timing.bottleneck)
|
|
|
|
# 3) 瓶颈分析建议
|
|
res.advice = self._advice(case, plan, res)
|
|
return res
|
|
|
|
# ------------------------------------------------------------------
|
|
def _check_constraints(self, case: BmmCase, plan: ImplPlan) -> list:
|
|
s = self.spec
|
|
v = []
|
|
# 特殊分支走 AIV 通路, 无 Cube tile 概念, 跳过 Cube 侧约束
|
|
if plan.branch == "特殊分支":
|
|
if plan.used_core_num > s.aiv_num:
|
|
v.append(f"used_core_num={plan.used_core_num} 超 AIV 核数 {s.aiv_num}")
|
|
return v
|
|
|
|
# 转Matmul 折叠后由 Matmul 体系承接, 此处不校验 Cube tile
|
|
if plan.branch == "转Matmul":
|
|
return v
|
|
|
|
if plan.used_core_num > s.aic_num:
|
|
v.append(f"used_core_num={plan.used_core_num} 超 AIC 核数 {s.aic_num}")
|
|
# ASW 降核模式每核单份 L0C (无双缓冲), 其他分支双缓冲
|
|
l0c_factor = 1 if plan.branch == "ASW_Basic_降核" else 2
|
|
if plan.base_m * plan.base_n * 4 * l0c_factor > s.l0c_bytes:
|
|
v.append(f"L0C tile 超容量: BaseM*BaseN*4B*{l0c_factor}={plan.base_m*plan.base_n*4*l0c_factor}B > {s.l0c_bytes}B")
|
|
if plan.base_m * plan.base_k * case.dtype_in_bytes * 2 > s.l0a_bytes:
|
|
v.append("L0A tile 超容量")
|
|
if plan.base_n * plan.base_k * case.dtype_in_bytes * 2 > s.l0b_bytes:
|
|
v.append("L0B tile 超容量")
|
|
l1_need = 2 * plan.k_l1 * (plan.single_core_m + plan.single_core_n) * case.dtype_in_bytes
|
|
if l1_need > s.l1_bytes and plan.branch in ("MergeBatch", "IterBatch"):
|
|
v.append(f"L1 tile 超容量: 2*k_L1*(sM+sN)*dtype={l1_need/1024:.0f}KB > {s.l1_bytes/1024:.0f}KB")
|
|
if plan.k_l1 > 0 and plan.k_l1 * case.dtype_in_bytes < s.dvalue_min:
|
|
v.append(f"dValue={plan.k_l1*case.dtype_in_bytes}B < 下限 {s.dvalue_min}B, 搬移带宽利用率崩塌")
|
|
# StreamK 中间部分和按 4B (L0C dtype) 防精度丢失, 是正确行为, 不算违规
|
|
if plan.branch != "StreamK" and plan.out_dtype_bytes != case.dtype_out_bytes:
|
|
v.append(f"方案写出 dtype ({plan.out_dtype_bytes}B) 与 case C 矩阵 dtype ({case.dtype_c}) 不一致")
|
|
return v
|
|
|
|
# ------------------------------------------------------------------
|
|
def _advice(self, case: BmmCase, plan: ImplPlan, res: EvalResult) -> str:
|
|
t = res.timing
|
|
tips = []
|
|
if not res.feasible:
|
|
tips.append("方案违反硬件约束, 需先修正: " + res.violations)
|
|
bn = t.bottleneck
|
|
if bn == "MTE2_GM":
|
|
tips.append("瓶颈在 GM 搬入: 可考虑增大 tile 提升 dValue/单核搬移量, "
|
|
"或利用 L2 驻留吸收重复读 (MergeBatch/ASW swizzle 方向)")
|
|
elif bn == "MTE2_L2":
|
|
tips.append("瓶颈在 L2 重复读: 优化核间分配/swizzle 窗口压低活跃工作集")
|
|
elif bn == "MMAD":
|
|
tips.append("瓶颈在 Cube 计算: 已接近理论算力上限, 检查是否有冗余计算 "
|
|
"(MergeBatch 交叉项) 可消除")
|
|
elif bn == "FIXPIPE":
|
|
tips.append("瓶颈在 Fixpipe 写出: 检查输出 dtype (fp16/fp8 可减半写出量), "
|
|
"或评估输出驻留 L2 异步回写策略")
|
|
if plan.branch == "MergeBatch" and plan.k_l1 < case.k:
|
|
tips.append("警告: MergeBatch 处于 L1 绑定情形 (k_L1<K), 理论证明其恒劣于 "
|
|
"IterBatch (v1.1 §4.4), 建议改用 IterBatch")
|
|
return " | ".join(tips) if tips else "方案合理, 流水均衡"
|