理论最优实现方案是什么 —— 该走哪个分支 (转Matmul/特殊分支/MergeBatch/IterBatch/StreamK/ASW_Basic), 核间怎么切, 核内 tile 怎么取, Cache/流水怎么配;
这个实现方案在 NPU 上跑得怎么样 —— 各级硬件 (MTE2 搬入 / Cube 计算 / Fixpipe 搬出 / StreamK 归约) 时延多少, 谁是瓶颈, 怎么优化.

与 ops-nn 源码的关系: 源码是"当前实现", 本软件是"理论最优"——用于评估源码方案离理论上限有多远、瓶颈在哪里.

快速上手

# 无需安装, 在项目根目录直接运行 (Python >= 3.8, 仅标准库)

# 模式 1: 方案推荐 —— 输入 case, 输出理论最优方案 + 时延评估
python -m bmm_theory recommend examples/cases_demo.csv -o result.csv --plans plans.csv -v

# 模式 2: 方案评估 —— 输入 case + 自己的实现方案, 评估硬件表现与瓶颈
python -m bmm_theory evaluate examples/cases_demo.csv examples/plans.csv -o eval.csv -v

# 跑单元测试
python -m unittest discover -s tests -v

case 输入 csv 格式

列名	含义	缺省
case_id	标识	case_行号
batch_a, batch_b	A/B 的 batch 维	1
m, n, k	矩阵维度	1
dtype_a, dtype_b, dtype_c	数据类型 (bf16/fp16/fp8/fp32)	bf16
trans_a, trans_b	是否转置	0
has_bias	是否有 bias	0
deterministic_level	确定性等级 (>=2 禁用 StreamK)	0

examples/cases_demo.csv 里有一份覆盖各分支的示例.

输出

result.csv: 每行一个 case —— case 参数 + plan_* 方案字段 + 时延评估 (t_mte2/t_mmad/t_fixpipe/t_total 等, 单位秒) + bottleneck 瓶颈级 + bound_type Bound 类型 + advice 仲裁过程与优化建议;
plans.csv (可选): 纯方案表, 标准结构体 ImplPlan 的全部字段, 可直接作为模式 2 的输入.

时延模型要点 (读结果前必看)

搬入 MTE2 分两段: t_mte2_gm (GM→L1 直读, 按 GM 带宽 1.6TB/s; 无论是否随路驻留 L2 都不再累加 L2→L1) 与 t_mte2_l2 (数据驻留 L2 后重复读命中, 按 L2 带宽 5.2TB/s);
Cube 计算 t_mmad: 单核算力 ≈15.2 TFLOPS, MergeBatch 的冗余计算计入;
Fixpipe 搬出 t_fixpipe: 数据量按 C 矩阵 dtype 计 (fp16/fp8 时随路转换、写出量减半); StreamK 中间部分和为防精度丢失按 4B (L0C dtype) 计;
总时延 t_total = max(各级) + t_drain (双缓冲稳态取最大 + 末级排空暴露);
瓶颈交换: 搬移瓶颈可牺牲算力换搬移效率 (MergeBatch), 计算瓶颈可牺牲搬移换计算效率 (ASW_Basic 切 M/N).

目录结构

BMM_Theory/
├── README.md                  # 本文档 (用法)
├── bmm_theory/                # 软件包
│   ├── __main__.py            # CLI 入口 (recommend / evaluate)
│   ├── models.py              # 数据模型: BmmCase / ImplPlan(标准结构体) / HardwareTiming
│   ├── hardware/ascend950pr.py# 950PR 硬件参数表 (换芯片只换这份)
│   ├── timing.py              # 时延评估引擎 (MTE2/Cube/Fixpipe/Reduce 模型)
│   ├── router.py              # 分支决策路由 + 重叠区仲裁
│   ├── evaluator.py           # 方案评估器 (约束校验 + 瓶颈分析)
│   ├── io_csv.py              # csv 输入输出
│   └── branches/              # 各分支理论脚本 (本期: MergeBatch + IterBatch)
│       ├── merge_batch.py
│       └── iter_batch.py
├── docs/                      # 文档 (架构 + 理论梳理)
│   ├── 01_软件架构.md
│   └── 02_分支理论/
│       ├── 00_总纲_分支决策树.md
│       ├── 01_MergeBatch分支.md
│       └── 02_IterBatch分支.md
├── examples/                  # 示例输入输出
└── tests/                     # 单元测试 (固化文档边界 case)

当前进度与路线图

分支	进入条件	方案生成	时延评估	状态
MergeBatch	✅	✅	✅	本期完成
IterBatch (a/b/c/d 四形态)	✅	✅	✅	本期完成
转Matmul	标注	—	—	路由标注, 详实现待迭代
特殊分支 (K=0/1)	标注	—	—	路由标注, 详实现待迭代
ASW_Basic (含降核/swizzle/L2分组/尾轮A0/A1a/A1b/方案B)	—	—	—	下一迭代
StreamK	—	—	—	下一迭代

设计原则

以理论分析文档为准——源码实现仅作对照, 不作为依据;
参数表与逻辑分离——换芯片只换 hardware/ 下的参数表, 分支逻辑不动;
结构体标准化——ImplPlan 是推荐输出与评估输入的统一格式, 字段与理论文档符号一一对应;
可验证——文档中的典型边界 case 全部固化进 tests/, 改动不破结论.