Add BMM_Theory: bmm_theory/branches/special.py

BMM/BMM_Theory/bmm_theory/branches/special.py

@@ -0,0 +1,91 @@
+"""特殊分支: 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,
+        )