From 149de60c3d4568a27bcd60bd3c73e42cb4c48f8d Mon Sep 17 00:00:00 2001 From: admin Date: Thu, 3 Sep 2026 09:25:44 +0000 Subject: [PATCH] Add BMM_Theory: bmm_theory/branches/to_matmul.py --- .../bmm_theory/branches/to_matmul.py | 96 +++++++++++++++++++ 1 file changed, 96 insertions(+) create mode 100644 BMM/BMM_Theory/bmm_theory/branches/to_matmul.py diff --git a/BMM/BMM_Theory/bmm_theory/branches/to_matmul.py b/BMM/BMM_Theory/bmm_theory/branches/to_matmul.py new file mode 100644 index 0000000..6a6e28a --- /dev/null +++ b/BMM/BMM_Theory/bmm_theory/branches/to_matmul.py @@ -0,0 +1,96 @@ +"""转Matmul 分支: 单边 batch=1 时折叠转普通 Matmul. + +理论依据: 《BMM算子优化分析 v0.98》§四 + docs/02_分支理论/03_转Matmul分支.md. + +核心: 单边 batch=1 的 BMM 与 Matmul 只差一个维度标签, 折叠后复用 Matmul 优化体系. + - BatchB=1: 左矩阵 [B,M,K] batch 维与 M 维 ND 下内存相邻紧排, 直接视图 [B*M,K], 零重排; + - BatchA=1: 右矩阵折叠 [K,B*N] 需真实转置重排 + 输出 scatter, 有代价; + A 小 (MK*dt<=L1) 时优先留 BMM 分支内做广播友好形态 (A 常驻 L1). +""" + +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 ToMatmulBranch(Branch): + name = "转Matmul" + + def __init__(self, spec: NpuSpec = ASCEND950PR): + super().__init__(spec) + + def check_conditions(self, case: BmmCase) -> list: + c1 = (case.batch_a == 1) or (case.batch_b == 1) + return [ConditionCheck( + "1_单边batch=1: BatchA=1 或 BatchB=1", + c1, f"batchA={case.batch_a}, batchB={case.batch_b}")] + + # ------------------------------------------------------------------ + def make_plan(self, case: BmmCase) -> ImplPlan: + s = self.spec + m, n, k = case.m, case.n, case.k + dt = case.dtype_in_bytes + + if case.batch_b == 1: + # BatchB=1: 免费折叠 [B,M,K] -> [B*M, K] + sub = "BatchB=1免费折叠" + fold_m, fold_n = case.batch_a * m, n + note = f"左矩阵 [{case.batch_a},{m},{k}] 视图折叠为 [{fold_m},{k}], 零重排零 split" + else: + # BatchA=1: 看 A 是否可常驻 L1 + a_bytes = m * k * dt + if a_bytes <= s.l1_bytes: + sub = "BatchA=1广播友好(A常驻L1)" + fold_m, fold_n = m, case.batch_b * n + note = (f"A[{m},{k}]={a_bytes/1024:.0f}KB <= L1, 留 BMM 分支内做广播: " + f"A 常驻 L1 逐 batch 复用, 避免转置重排") + else: + sub = "BatchA=1重排转Matmul" + fold_m, fold_n = m, case.batch_b * n + note = (f"A[{m},{k}]={a_bytes/1024:.0f}KB > L1, 广播扩展后比较 " + f"'BMM 分支' vs '重排+转Matmul' 时延择优 (重排代价 O(B*K*N))") + + return ImplPlan( + case_id=case.case_id, branch=self.name, + used_core_num=s.aic_num, + split_b=1, m_cnt=0, n_cnt=0, grid_k=1, # 折叠后走 Matmul 切分, 此处不展开 + core_map=f"折叠为 Matmul [{fold_m},{k}]x[{k},{fold_n}], 复用 Matmul 切分体系", + b_core=0, merge_b0=1, + single_core_m=0, single_core_n=0, single_core_k=k, + k_l1=0, b_l1=1, l1_form="", + base_m=0, base_n=0, base_k=0, + l2_policy_in="", l2_policy_out="", swizzle_w=0, workspace_bytes=0, + tail_strategy="转Matmul后由 Matmul 体系决定", + fixpipe_unitflag=True, + out_dtype_bytes=case.dtype_out_bytes, + note=f"{sub}: {note}", + ) + + # ------------------------------------------------------------------ + def evaluate(self, case: BmmCase, plan: ImplPlan) -> HardwareTiming: + """折叠后按 Matmul 粗估 (详细切分待 MM 理论体系打通后接入).""" + s = self.spec + if case.batch_b == 1: + fold_m, fold_n, kk = case.batch_a * case.m, case.n, case.k + else: + fold_m, fold_n, kk = case.m, case.batch_b * case.n, case.k + dt = case.dtype_in_bytes + out_b = case.dtype_out_bytes + + flops = 2.0 * fold_m * fold_n * kk + t_mmad = flops / (s.aic_num * s.q16) + in_bytes = (fold_m * kk + kk * fold_n) * dt + out_bytes = fold_m * fold_n * out_b + t_mte2 = in_bytes / s.bw_gm + t_fix = out_bytes / s.bw_gm + + return assemble_timing( + t_mte2_gm=t_mte2, t_mte2_l2=0.0, t_dma_cmd=0.0, + t_mmad=t_mmad, t_fixpipe=t_fix, 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=flops, fixpipe_bytes=out_bytes, + )