Files
matmul-analysis/BMM/BMM_Theory/docs/01_软件架构.md

7.3 KiB
Raw Blame History

BMM_Theory 软件架构

1. 分层视图

┌─────────────────────────────────────────────────────────┐
│ CLI (__main__.py)   recommend / evaluate 两种模式         │
├─────────────────────────────────────────────────────────┤
│ 应用层                                                   │
│  ├─ router.BranchRouter   case -> 分支决策 + 方案 + 时延  │
│  └─ evaluator.PlanEvaluator  (case, 方案) -> 校验 + 时延  │
├─────────────────────────────────────────────────────────┤
│ 分支层 (branches/)  —— 每个分支一个文件, 三个接口          │
│  ├─ merge_batch.py   MergeBatchBranch                    │
│  ├─ iter_batch.py    IterBatchBranch                     │
│  ├─ to_matmul.py     ToMatmulBranch                      │
│  ├─ special.py       SpecialBranch                       │
│  ├─ stream_k.py      StreamKBranch                       │
│  └─ asw_basic.py     AswBasicBranch                      │
│      统一接口: check_conditions / make_plan / evaluate    │
├─────────────────────────────────────────────────────────┤
│ 模型层                                                   │
│  ├─ models.py       BmmCase / ImplPlan / HardwareTiming  │
│  ├─ timing.py       MTE2 / MMAD / Fixpipe / Reduce 时延引擎│
│  └─ constraints.py  单一约束源 (生成与校验共用)            │
├─────────────────────────────────────────────────────────┤
│ 硬件层 (hardware/)                                       │
│  └─ ascend950pr.py  NpuSpec 参数表 (换芯片只换这份)        │
├─────────────────────────────────────────────────────────┤
│ IO 层 (io_csv.py)  case/plan/result 的 csv 读写           │
└─────────────────────────────────────────────────────────┘

2. 两种工作模式的数据流

模式 1: 方案推荐 (recommend)

cases.csv ──> load_cases ──> BranchRouter.route(case)
                                  │ 决策树: 前置归约 -> B>=C? -> IterBatch/MergeBatch 仲裁
                                  ▼
                        ImplPlan (标准结构体) + HardwareTiming
                                  ▼
                 result.csv (case + plan_* + 时延 + 瓶颈 + 建议)
                 plans.csv  (纯方案表, 可作模式 2 输入)

模式 2: 方案评估 (evaluate)

cases.csv + plans.csv ──> PlanEvaluator.evaluate(case, plan)
                               │ 1. 硬件约束校验 (L0C/L0A/L0B/L1 容量, dValue, 核数)
                               │ 2. 按方案分支调用对应时延模型
                               │ 3. 瓶颈分析 + 优化建议
                               ▼
                    eval.csv (feasible / violations / 时延 / advice)

3. 关键设计决策

3.1 ImplPlan 是系统的"通用货币"

推荐模式的输出与评估模式的输入共用同一个结构体 ImplPlanmodels.py),字段分四组:

字段 说明
分支与核间切分 branch, used_core_num, split_b, m_cnt, n_cnt, grid_k, core_map 第一性问题: 核间怎么分 B/M/N/K
核内 tiling b_core, merge_b0, single_core_m/n/k, k_l1, b_l1, l1_form, base_m/n/k 第二性问题: 核内 tile
Cache 策略 l2_policy_in, l2_policy_out, swizzle_w, workspace_bytes L2 驻留/直写、滑窗
流水策略 tail_strategy, fixpipe_unitflag, out_dtype_bytes 尾轮 A0/A1a/A1b/方案B、unitflag

to_row() / from_row() 负责与 csv 的双向转换,字段名即 csv 列名。

3.2 分支 = 自包含插件

每个分支类实现三个方法:

class Branch:
    def check_conditions(case) -> list[ConditionCheck]  # 逐条进入条件判定 (可解释)
    def make_plan(case) -> ImplPlan                     # 理论最优方案生成
    def evaluate(case, plan) -> HardwareTiming          # 时延评估

新增分支 = 在 branches/ 下加一个文件 + 在 router.py 注册。互不依赖,可独立迭代。

3.3 硬件参数与逻辑分离

hardware/ascend950pr.pyNpuSpec 集中所有芯片常数(核数/算力/各级容量/带宽/搬移效率经验值/T_cmd。所有分支通过 self.spec 取参数——换芯片时新增一份参数表即可,分支逻辑零改动。

3.4 时延引擎统一在 timing.py

分支不各自造轮子,统一调用 timing.py 的四个函数:

  • eval_mte2(move) —— GM/L2 两段搬入 + DMA 命令开销
  • eval_mmad(flops) —— Cube 计算
  • eval_fixpipe(bytes, to_l2) —— 搬出(直写 GM 或驻留 L2
  • eval_streamk_reduce(...) —— StreamK 归约(部分和 4B 驻留 L2、AIV 求和)

assemble_timing(...) 汇总并判瓶颈。归约计账约定issue#9REDUCE 默认串行追加(体现在 t_drain),不进稳态 max();仅当显式声明可流水掩盖(reduce_serial=False)时才进 max()。带宽模型:每核独立 DMA 引擎、带宽按核数配平(尾轮文档 §2.4),活跃核数 < C 时聚合带宽按比例下降。

3.5 单一约束源constraints.py

约束知识L0C/L0A/L0B/L1/dValue/min_TileSize/核数)只有一个事实来源constraints.check_plan_constraints(case, plan)。生成侧recommend 的 _wrap_checked 自检与评估侧evaluate 的 _check_constraints)共用同一函数,保证"推荐方案 vs 自带评估器"口径一致,不再出现生成说可行、校验说不可行的矛盾。

生成侧的 tile 收敛辅助(clamp_base_k / clamp_base_mn_l0c)也在此模块,各分支 make_plan 调用同源函数保证生成的 tile 不越界。dValue 口径裁定issue#6仅当 K 被切分成段k_l1 < K、K 段为搬移连续维)时 128B 下限才生效K 整驻留k_l1 ≥ K时连续维是 M/N豁免 K 向下限。

4. 扩展指南(后续迭代)

当前六分支已全部实现转Matmul / 特殊 / MergeBatch / IterBatch / StreamK / ASW_Basic。后续方向

  1. 转Matmul 精切: 当前折叠后只粗估(to_matmul.py::evaluate 按 Matmul 总量),接入 MM 理论体系做精确切分。
  2. 换芯片: 复制 hardware/ascend950pr.py 改常数,NpuSpec 接口不变。
  3. 标定 T_cmd: 当前取 50ns 估计值,实测后改 t_cmd_ns 一处即可。
  4. 新分支: 在 branches/ 下加一个文件实现三接口 + 在 router.py 注册 + 在 evaluator._BRANCH_EVAL 注册;约束一律走 constraints.py,不在分支内另造规则。