From 5bee7315702f5fd614d3c47265429cd529f41c96 Mon Sep 17 00:00:00 2001 From: admin Date: Thu, 3 Sep 2026 09:25:38 +0000 Subject: [PATCH] Update BMM_Theory: bmm_theory/router.py --- BMM/BMM_Theory/bmm_theory/router.py | 87 ++++++++++++++++++----------- 1 file changed, 53 insertions(+), 34 deletions(-) diff --git a/BMM/BMM_Theory/bmm_theory/router.py b/BMM/BMM_Theory/bmm_theory/router.py index 07e4d44..9f27d54 100644 --- a/BMM/BMM_Theory/bmm_theory/router.py +++ b/BMM/BMM_Theory/bmm_theory/router.py @@ -1,15 +1,15 @@ """分支决策路由: 按决策树推导分支, 重叠区由端到端时延模型仲裁. -决策树 (v0.98 §3.3): - 1. 前置归约: BatchA=1 或 BatchB=1 -> 转Matmul (本期仅标注, 详实现后续迭代) - K=0 / K=1 -> 特殊分支 (本期仅标注) - 2. B >= C: 切B -> IterBatch 与 MergeBatch 仲裁 +决策树 (v0.98 §3.3 + 各分支理论文档): + 1. 前置归约: BatchA=1 或 BatchB=1 -> 转Matmul + K=0 / K=1 -> 特殊分支 (AIV 向量通路) + 2. B >= C 且 BatchA==BatchB: 切B -> IterBatch 与 MergeBatch 仲裁 仲裁规则 (v1.1 §4.5 统一分界): MergeBatch 最优 <=> K截断(k_L1=K) 且 b_core > b0*(T_comp+T_write)/T_cmd L1 绑定时 MergeBatch 恒劣于 IterBatch; 两分支同时合法时用端到端时延模型 T_total 仲裁. - 3. B < C: 切 M/N -> ASW_Basic (后续迭代) - 4. B/M/N 都买不满: 切 K -> StreamK (后续迭代) + 3. StreamK 检查: P <= C/2 且满足切K条件 -> StreamK (B/M/N 买不满时买 K) + 4. 兜底: ASW_Basic 切 M/N (含降核模式) """ from __future__ import annotations @@ -19,11 +19,10 @@ from .models import BmmCase, ImplPlan, HardwareTiming from .branches.base import BranchResult from .branches.merge_batch import MergeBatchBranch from .branches.iter_batch import IterBatchBranch - -BRANCH_TO_MATMUL = "转Matmul" -BRANCH_SPECIAL = "特殊分支" -BRANCH_ASW = "ASW_Basic" -BRANCH_STREAMK = "StreamK" +from .branches.to_matmul import ToMatmulBranch +from .branches.special import SpecialBranch +from .branches.stream_k import StreamKBranch +from .branches.asw_basic import AswBasicBranch class BranchRouter: @@ -33,27 +32,40 @@ class BranchRouter: self.spec = spec self.merge_batch = MergeBatchBranch(spec) self.iter_batch = IterBatchBranch(spec) + self.to_matmul = ToMatmulBranch(spec) + self.special = SpecialBranch(spec) + self.stream_k = StreamKBranch(spec) + self.asw_basic = AswBasicBranch(spec) # ------------------------------------------------------------------ def route(self, case: BmmCase) -> dict: """返回 {branch, plan, timing, arbitration, candidates}.""" s = self.spec - # 1) 前置归约 + # 1) 前置归约: 转Matmul / 特殊分支 if case.batch_a == 1 or case.batch_b == 1: - return self._stub(case, BRANCH_TO_MATMUL, - "BatchA=1或BatchB=1, 折叠转普通Matmul (该分支详实现待后续迭代)") + r = self.to_matmul.analyze(case) + return self._wrap(r, "BatchA=1或BatchB=1, 折叠转普通Matmul") if case.k <= 1: - return self._stub(case, BRANCH_SPECIAL, - f"K={case.k}, Cube 无用, 走 AIV 向量通路 (该分支详实现待后续迭代)") + r = self.special.analyze(case) + note = "K=0纯写值" if case.k == 0 else "K=1逐元素乘, 走AIV向量通路" + return self._wrap(r, note) - # 2) B >= C: IterBatch / MergeBatch + # 2) B >= C 且 BatchA==BatchB: IterBatch / MergeBatch if case.batch_c >= s.aic_num and case.batch_a == case.batch_b: return self._route_split_b(case) - # 3/4) 兜底标注 - return self._stub(case, BRANCH_ASW, - "B ASW_Basic/StreamK (该分支详实现待后续迭代)") + # 3) StreamK: P <= C/2 且切K条件满足 + sk = self.stream_k.analyze(case) + if sk.capable: + return self._wrap(sk, f"P<=C/2, B/M/N并行度买不满, 切K (grid_K={sk.plan.grid_k})") + + # 4) 兜底: ASW_Basic (含降核模式) + asw = self.asw_basic.analyze(case) + note = "ASW_Basic兜底" + if sk.checks and not sk.capable: + note += f" (StreamK未过: {sk.failed_conditions()})" + return self._wrap(asw, note) # ------------------------------------------------------------------ def _route_split_b(self, case: BmmCase) -> dict: @@ -75,11 +87,10 @@ class BranchRouter: lat_win = self.merge_batch.name if t_mb <= t_ib else self.iter_batch.name win = self.merge_batch.name if mb_win else self.iter_batch.name arbitration = ( - f"两分支均合法, 仲裁:\n" - f" [分界条件] MergeBatch最优={mb_win} ({detail})\n" - f" [时延模型] T_MergeBatch={t_mb*1e6:.2f}us vs T_IterBatch={t_ib*1e6:.2f}us " - f"-> {lat_win}更优\n" - f" [裁决] {win}" + ("" if win == lat_win else f" (分界条件与时延模型不一致, 以时延模型为准: {lat_win})") + f"两分支均合法, 仲裁: " + f"[分界条件] MergeBatch最优={mb_win} ({detail}); " + f"[时延模型] T_MergeBatch={t_mb*1e6:.2f}us vs T_IterBatch={t_ib*1e6:.2f}us -> {lat_win}更优; " + f"[裁决] {win}" + ("" if win == lat_win else f" (分界条件与时延模型不一致, 以时延模型为准: {lat_win})") ) if win != lat_win: win = lat_win # 时延模型为最终裁决 @@ -87,10 +98,14 @@ class BranchRouter: win = candidates[0][0] arbitration = f"仅 {win} 条件满足" else: - return self._stub( - case, BRANCH_ASW, - "IterBatch/MergeBatch 进入条件均不满足 (如负载均衡/搬移效率不达标), " - "回落 ASW_Basic (该分支详实现待后续迭代); " + # 切B分支都不满足, 尝试 StreamK 再回落 ASW + sk = self.stream_k.analyze(case) + if sk.capable: + return self._wrap(sk, "切B分支条件不满足, 落 StreamK") + asw = self.asw_basic.analyze(case) + return self._wrap( + asw, + f"IterBatch/MergeBatch 进入条件均不满足, 回落 ASW_Basic; " f"IterBatch未过: {ib.failed_conditions()}; " f"MergeBatch未过: {mb.failed_conditions()}") @@ -105,8 +120,12 @@ class BranchRouter: } # ------------------------------------------------------------------ - def _stub(self, case: BmmCase, branch: str, note: str) -> dict: - plan = ImplPlan(case_id=case.case_id, branch=branch, - used_core_num=self.spec.aic_num, note=note) - return {"branch": branch, "plan": plan, "timing": None, - "arbitration": note, "candidates": {}} + @staticmethod + def _wrap(result: BranchResult, note: str) -> dict: + return { + "branch": result.plan.branch if result.plan else "未知", + "plan": result.plan, + "timing": result.timing, + "arbitration": note + (f" | {result.note}" if result.note else ""), + "candidates": {}, + }