Update BMM_Theory: bmm_theory/evaluator.py

BMM/BMM_Theory/bmm_theory/evaluator.py

@@ -11,10 +11,19 @@ from .models import BmmCase, ImplPlan, EvalResult
 from .timing import bound_type_of
 from .branches.merge_batch import MergeBatchBranch
 from .branches.iter_batch import IterBatchBranch
+from .branches.to_matmul import ToMatmulBranch
+from .branches.special import SpecialBranch
+from .branches.stream_k import StreamKBranch
+from .branches.asw_basic import AswBasicBranch
 
 _BRANCH_EVAL = {
     "MergeBatch": MergeBatchBranch,
     "IterBatch": IterBatchBranch,
+    "转Matmul": ToMatmulBranch,
+    "特殊分支": SpecialBranch,
+    "StreamK": StreamKBranch,
+    "ASW_Basic": AswBasicBranch,
+    "ASW_Basic_降核": AswBasicBranch,
 }
 
 
@@ -48,10 +57,22 @@ class PlanEvaluator:
     def _check_constraints(self, case: BmmCase, plan: ImplPlan) -> list:
         s = self.spec
         v = []
+        # 特殊分支走 AIV 通路, 无 Cube tile 概念, 跳过 Cube 侧约束
+        if plan.branch == "特殊分支":
+            if plan.used_core_num > s.aiv_num:
+                v.append(f"used_core_num={plan.used_core_num} 超 AIV 核数 {s.aiv_num}")
+            return v
+
+        # 转Matmul 折叠后由 Matmul 体系承接, 此处不校验 Cube tile
+        if plan.branch == "转Matmul":
+            return v
+
         if plan.used_core_num > s.aic_num:
             v.append(f"used_core_num={plan.used_core_num} 超 AIC 核数 {s.aic_num}")
-        if plan.base_m * plan.base_n * 4 * 2 > s.l0c_bytes:
-            v.append(f"L0C tile 超容量: BaseM*BaseN*4B*2={plan.base_m*plan.base_n*8}B > {s.l0c_bytes}B")
+        # ASW 降核模式每核单份 L0C (无双缓冲), 其他分支双缓冲
+        l0c_factor = 1 if plan.branch == "ASW_Basic_降核" else 2
+        if plan.base_m * plan.base_n * 4 * l0c_factor > s.l0c_bytes:
+            v.append(f"L0C tile 超容量: BaseM*BaseN*4B*{l0c_factor}={plan.base_m*plan.base_n*4*l0c_factor}B > {s.l0c_bytes}B")
         if plan.base_m * plan.base_k * case.dtype_in_bytes * 2 > s.l0a_bytes:
             v.append("L0A tile 超容量")
         if plan.base_n * plan.base_k * case.dtype_in_bytes * 2 > s.l0b_bytes:
@@ -59,9 +80,10 @@ class PlanEvaluator:
         l1_need = 2 * plan.k_l1 * (plan.single_core_m + plan.single_core_n) * case.dtype_in_bytes
         if l1_need > s.l1_bytes and plan.branch in ("MergeBatch", "IterBatch"):
             v.append(f"L1 tile 超容量: 2*k_L1*(sM+sN)*dtype={l1_need/1024:.0f}KB > {s.l1_bytes/1024:.0f}KB")
-        if plan.k_l1 * case.dtype_in_bytes < s.dvalue_min:
+        if plan.k_l1 > 0 and plan.k_l1 * case.dtype_in_bytes < s.dvalue_min:
             v.append(f"dValue={plan.k_l1*case.dtype_in_bytes}B < 下限 {s.dvalue_min}B, 搬移带宽利用率崩塌")
-        if plan.out_dtype_bytes != case.dtype_out_bytes:
+        # StreamK 中间部分和按 4B (L0C dtype) 防精度丢失, 是正确行为, 不算违规
+        if plan.branch != "StreamK" and plan.out_dtype_bytes != case.dtype_out_bytes:
             v.append(f"方案写出 dtype ({plan.out_dtype_bytes}B) 与 case C 矩阵 dtype ({case.dtype_c}) 不一致")
         return v