92 lines
3.7 KiB
Python
92 lines
3.7 KiB
Python
"""特殊分支: K=0 (纯写值) / K=1 (逐元素乘), 走 AIV 向量通路.
|
|
|
|
理论依据: 《BMM算子优化分析 v0.98》§九 + docs/02_分支理论/04_特殊分支.md.
|
|
|
|
K 维是 Cube 存在的意义 (累加深度). K=0/1 时 Cube 的 16x16x16 粒度浪费,
|
|
走 AIV 向量通路 (GM->UB->算->GM) 优于 Cube 通路. 与切分正交的前置判断.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from ..hardware import NpuSpec, ASCEND950PR
|
|
from ..models import BmmCase, ImplPlan, HardwareTiming
|
|
from ..timing import assemble_timing
|
|
from .base import Branch, ConditionCheck
|
|
|
|
|
|
class SpecialBranch(Branch):
|
|
name = "特殊分支"
|
|
|
|
def __init__(self, spec: NpuSpec = ASCEND950PR):
|
|
super().__init__(spec)
|
|
|
|
def check_conditions(self, case: BmmCase) -> list:
|
|
c1 = case.k <= 1
|
|
checks = [ConditionCheck("1_K<=1 (Cube 无用)", c1, f"K={case.k}")]
|
|
if case.k == 1:
|
|
# K=1 触发 AIV 通路需 B >= 2*AIV核数 且单 batch 输入输出能驻留 UB
|
|
c2 = case.batch_c >= 2 * self.spec.aiv_num
|
|
checks.append(ConditionCheck(
|
|
"2_K=1的AIV触发: B >= 2*AIV核数 (开UB乒乓)",
|
|
c2, f"B={case.batch_c} vs {2*self.spec.aiv_num}"))
|
|
return checks
|
|
|
|
# ------------------------------------------------------------------
|
|
def make_plan(self, case: BmmCase) -> ImplPlan:
|
|
s = self.spec
|
|
if case.k == 0:
|
|
sub = "K=0纯写值"
|
|
note = "无任何计算, C=bias 或 0, 纯 AIV 写值; 按行均分到 AIV 核"
|
|
else:
|
|
sub = "K=1逐元素乘"
|
|
note = ("退化为 C=A⊙B 无累加深度, Cube 16x16x16 粒度浪费 15/16; "
|
|
"走 AIV 通路 GM->UB->Mul->GM, UB 乒乓")
|
|
return ImplPlan(
|
|
case_id=case.case_id, branch=self.name,
|
|
used_core_num=s.aiv_num, # 用 AIV 核
|
|
split_b=1, m_cnt=1, n_cnt=1, grid_k=1,
|
|
core_map="AIV 核间按行均分 (无 Cube tile 概念)",
|
|
b_core=0, merge_b0=1,
|
|
single_core_m=0, single_core_n=0, single_core_k=case.k,
|
|
k_l1=0, b_l1=1, l1_form="UB驻留(AIV)",
|
|
base_m=0, base_n=0, base_k=0,
|
|
l2_policy_in="allocate", l2_policy_out="direct_gm",
|
|
swizzle_w=0, workspace_bytes=0,
|
|
tail_strategy="不涉及(AIV逐元素)",
|
|
fixpipe_unitflag=False,
|
|
out_dtype_bytes=case.dtype_out_bytes,
|
|
note=f"{sub}: {note}",
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
def evaluate(self, case: BmmCase, plan: ImplPlan) -> HardwareTiming:
|
|
"""AIV 通路时延: 瓶颈在搬移 (AIV 算力远剩)."""
|
|
s = self.spec
|
|
b = case.batch_c
|
|
m, n, k = case.m, case.n, case.k
|
|
dt = case.dtype_in_bytes
|
|
out_b = case.dtype_out_bytes
|
|
|
|
if k == 0:
|
|
# 纯写值: 仅写出
|
|
t_in, t_compute, t_out = 0.0, 0.0, b * m * n * out_b / s.bw_gm
|
|
in_bytes, cube_flops = 0.0, 0.0
|
|
else:
|
|
# 逐元素乘: 搬入 A+B, 搬出 C, AIV 算力远剩
|
|
in_bytes = b * (m * k + k * n) * dt
|
|
out_bytes = b * m * n * out_b
|
|
t_in = in_bytes / s.bw_gm
|
|
t_out = out_bytes / s.bw_gm
|
|
# AIV 求积吞吐 (近似按 Q_AIV)
|
|
t_compute = b * m * n / s.q_aiv
|
|
cube_flops = float(b * m * n)
|
|
|
|
t_total = max(t_in, t_out, t_compute)
|
|
return assemble_timing(
|
|
t_mte2_gm=t_in, t_mte2_l2=0.0, t_dma_cmd=0.0,
|
|
t_mmad=t_compute, t_fixpipe=t_out, t_reduce=0.0, t_drain=0.0,
|
|
gm_read_bytes=in_bytes, l2_read_bytes=0.0, dma_cmd_count=0.0,
|
|
cube_flops=cube_flops,
|
|
fixpipe_bytes=b * m * n * out_b,
|
|
)
|