Update BMM_Theory: bmm_theory/router.py (fix review issues #4-#10)

BMM/BMM_Theory/bmm_theory/router.py

@@ -45,11 +45,17 @@ class BranchRouter:
         # 1) 前置归约: 转Matmul / 特殊分支
         if case.batch_a == 1 or case.batch_b == 1:
             r = self.to_matmul.analyze(case)
-            return self._wrap(r, "BatchA=1或BatchB=1, 折叠转普通Matmul")
+            return self._wrap_checked(case, r, "BatchA=1或BatchB=1, 折叠转普通Matmul")
         if case.k <= 1:
             r = self.special.analyze(case)
             note = "K=0纯写值" if case.k == 0 else "K=1逐元素乘, 走AIV向量通路"
-            return self._wrap(r, note)
+            # issue#4 P0: 特殊分支 capable=False (如 K=1 但 B<128 不满足 UB 乒乓) 时
+            # r.plan=None, 必须兜底而不能把 None 传给下游 —— 显式标注"该区域暂无理论方案"
+            if not r.capable or r.plan is None:
+                return self._no_plan(case, "特殊分支",
+                                     note + f"; 但进入条件不满足 ({r.failed_conditions()}), "
+                                            f"该区域暂无理论方案, 建议参考 Cube 兜底或 AIV 单缓冲")
+            return self._wrap_checked(case, r, note)
 
         # 2) B >= C 且 BatchA==BatchB: IterBatch / MergeBatch
         if case.batch_c >= s.aic_num and case.batch_a == case.batch_b:
@@ -58,14 +64,14 @@ class BranchRouter:
         # 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})")
+            return self._wrap_checked(case, 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)
+        return self._wrap_checked(case, asw, note)
 
     # ------------------------------------------------------------------
     def _route_split_b(self, case: BmmCase) -> dict:
@@ -101,24 +107,55 @@ class BranchRouter:
             # 切B分支都不满足, 尝试 StreamK 再回落 ASW
             sk = self.stream_k.analyze(case)
             if sk.capable:
-                return self._wrap(sk, "切B分支条件不满足, 落 StreamK")
+                return self._wrap_checked(case, sk, "切B分支条件不满足, 落 StreamK")
             asw = self.asw_basic.analyze(case)
-            return self._wrap(
+            return self._wrap_checked(
-                asw,
+                case, asw,
                 f"IterBatch/MergeBatch 进入条件均不满足, 回落 ASW_Basic; "
                 f"IterBatch未过: {ib.failed_conditions()}; "
                 f"MergeBatch未过: {mb.failed_conditions()}")
 
         chosen = mb if win == self.merge_batch.name else ib
+        result = BranchResult(capable=True, plan=chosen.plan, timing=chosen.timing)
+        return self._wrap_checked(case, result, arbitration,
+                                  candidates={n: r.capable for n, r in
+                                              [(self.merge_batch.name, mb), (self.iter_batch.name, ib)]})
+
+    # ------------------------------------------------------------------
+    def _wrap_checked(self, case: BmmCase, result: BranchResult, note: str,
+                      candidates: dict | None = None) -> dict:
+        """生成后自检 (issue#5): 推荐方案必须通过统一约束源校验, 不可行则标注违规.
+
+        约束源与 evaluate 共用 constraints.check_plan_constraints, 保证
+        "推荐方案 vs 自带评估器" 口径一致, 不再出现生成说可行、校验说不可行的矛盾.
+        """
+        from .constraints import check_plan_constraints
+        violations = check_plan_constraints(case, result.plan, self.spec) if result.plan else []
+        if violations:
+            note = (note + " [自检违规: " + "; ".join(violations) +
+                    "] —— 方案生成存在缺陷, 需人工复核")
         return {
-            "branch": win,
+            "branch": result.plan.branch if result.plan else "未知",
-            "plan": chosen.plan,
+            "plan": result.plan,
-            "timing": chosen.timing,
+            "timing": result.timing,
-            "arbitration": arbitration,
+            "arbitration": note + (f" | {result.note}" if result.note else ""),
-            "candidates": {n: r.capable for n, r in
+            "candidates": candidates or {},
-                           [(self.merge_batch.name, mb), (self.iter_batch.name, ib)]},
+            "self_check_violations": violations,
         }
 
+    @staticmethod
+    def _no_plan(case: BmmCase, branch: str, note: str) -> dict:
+        """兜底: 分支 capable=False 时给出最小占位方案, 保证下游不崩溃 (issue#4).
+
+        方案标注 used_core_num=0 + 分支名, arbitration 说明"该区域暂无理论方案",
+        不产生时延 (timing=None), 供上层跳过或人工处理.
+        """
+        plan = ImplPlan(case_id=case.case_id, branch=branch,
+                        used_core_num=0, note="该区域暂无理论方案(进入条件不满足)")
+        return {"branch": branch, "plan": plan, "timing": None,
+                "arbitration": "[无方案] " + note, "candidates": {},
+                "self_check_violations": []}
+
     # ------------------------------------------------------------------
     @staticmethod
     def _wrap(result: BranchResult, note: str) -> dict:
@@ -128,4 +165,5 @@ class BranchRouter:
             "timing": result.timing,
             "arbitration": note + (f" | {result.note}" if result.note else ""),
             "candidates": {},
+            "self_check_violations": [],
         }