Add BMM_Theory: bmm_theory/branches/to_matmul.py

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,
+        )