Add BMM_Theory: bmm_theory/timing.py

BMM/BMM_Theory/bmm_theory/timing.py

@@ -0,0 +1,132 @@
+"""时延评估引擎: 按 NPU 硬件流水模型计算各级时延.
+
+模型依据 (v0.98 §3.1 + 用户补充说明):
+
+    T_total = max(T_MMAD, T_MTE2, T_MTE1, T_Fixpipe [, T_Reduce]) + T_drain
+
+搬入 (MTE2) 的两条路径:
+  1. GM->L1 直读: 数据量按 GM 带宽计时.
+     算子软件可配置是否随路驻留 L2 (Cache Hint allocate);
+     无论是否驻留, 该次搬入时延都按 GM 带宽计算, 不累加后续 L2->L1.
+     驻留 L2 的意义只在于: 后续重复读可以走路径 2.
+  2. L2->L1: 数据已驻留 L2 后的再次/重复读取, 按 L2 带宽 (5.2TB/s) 计时.
+
+Cube (MMAD): 按单核算力 Q16 x 使用核数计时. MergeBatch 的冗余计算计入 cube_flops.
+
+Fixpipe 搬出: 写出 dtype = C 矩阵 dtype (fp16/fp8 时随路转换, 数据量按 C dtype 计);
+  StreamK 的中间部分和为防精度丢失按 L0C dtype (4B) 计, 且需再次读入 AIV 归约.
+
+流水掩盖: 双缓冲下稳态时延取各级最大; T_drain 为末级排空暴露, 由各分支模型给出.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .hardware import NpuSpec, ASCEND950PR
+from .models import HardwareTiming
+
+
+@dataclass
+class MoveInPlan:
+    """一次稳态搬入的描述 (单核视角, 单位 Byte)."""
+    gm_bytes: float = 0.0    # GM->L1 直读量
+    l2_bytes: float = 0.0    # L2->L1 命中量
+    dma_cmds: float = 0.0    # GM->L1 DMA 命令次数
+
+
+def eval_mte2(move: MoveInPlan, spec: NpuSpec = ASCEND950PR,
+              active_cores: int | None = None) -> tuple[float, float, float, float]:
+    """MTE2 搬入时延.
+
+    返回 (t_mte2_gm, t_mte2_l2, t_mte2_total, t_dma_cmd).
+
+    带宽模型: 每核独立 DMA 引擎, 单核带宽份额 = 聚合带宽/C (尾轮文档 §2.4).
+    active_cores < C 时聚合带宽 = active_cores x 单核份额 (尾轮不加速).
+    """
+    cores = active_cores or spec.aic_num
+    bw_gm_eff = spec.bw_pc * cores if cores < spec.aic_num else spec.bw_gm
+    bw_l2_eff = spec.bw_l2_pc * cores if cores < spec.aic_num else spec.bw_l2
+
+    # move 为单核数据量; 聚合时延 = 单核量 / 单核带宽份额 (等价 总量/聚合带宽)
+    t_gm = move.gm_bytes / spec.bw_pc if move.gm_bytes > 0 else 0.0
+    t_l2 = move.l2_bytes / spec.bw_l2_pc if move.l2_bytes > 0 else 0.0
+    t_cmd = move.dma_cmds * spec.t_cmd
+    # GM 直读与 L2 命中是不同数据的两段搬入, 时延相加; DMA 开销叠加在 GM 路径上
+    return t_gm, t_l2, t_gm + t_l2 + t_cmd, t_cmd
+
+
+def eval_mmad(flops_per_core: float, spec: NpuSpec = ASCEND950PR) -> float:
+    """Cube 计算时延: 单核计算量 / 单核算力."""
+    return flops_per_core / spec.q16 if flops_per_core > 0 else 0.0
+
+
+def eval_fixpipe(bytes_per_core: float, to_l2: bool,
+                 spec: NpuSpec = ASCEND950PR) -> float:
+    """Fixpipe 搬出时延.
+
+    to_l2=True: 写出驻留 L2 (5.2TB/s 写口, 异步回写 GM 不占算子内时延);
+    to_l2=False: 直写 GM, 与读共享 1.6TB/s 总线.
+    """
+    if bytes_per_core <= 0:
+        return 0.0
+    bw = spec.bw_l2_pc if to_l2 else spec.bw_pc
+    return bytes_per_core / bw
+
+
+def eval_streamk_reduce(tile_elems: float, grid_k: int, out_dtype_bytes: int,
+                        spec: NpuSpec = ASCEND950PR) -> float:
+    """StreamK 单 tile 归约时延 (v0.98 §七).
+
+    部分和 dtype = L0C dtype (4B, 防精度丢失), 驻留 L2, AIV 归约:
+      AIC 写部分和 grid_k x tile x 4B / W_L2
+      AIV 读回     grid_k x tile x 4B / W_L2
+      AIV 求和     grid_k x tile / Q_AIV
+      写回         tile x outB / W_L2
+    """
+    b4 = 4
+    w_l2 = spec.bw_l2
+    t_write_partial = grid_k * tile_elems * b4 / w_l2
+    t_read_back = grid_k * tile_elems * b4 / w_l2
+    t_sum = grid_k * tile_elems / spec.q_aiv
+    t_write_out = tile_elems * out_dtype_bytes / w_l2
+    return t_write_partial + t_read_back + t_sum + t_write_out
+
+
+def assemble_timing(t_mte2_gm: float, t_mte2_l2: float, t_dma_cmd: float,
+                    t_mmad: float, t_fixpipe: float, t_reduce: float,
+                    t_drain: float,
+                    gm_read_bytes: float, l2_read_bytes: float,
+                    dma_cmd_count: float, cube_flops: float,
+                    fixpipe_bytes: float) -> HardwareTiming:
+    """汇总各级时延, 判定瓶颈."""
+    t_mte2 = t_mte2_gm + t_mte2_l2 + t_dma_cmd
+    stages = {
+        "MTE2_GM": t_mte2_gm + t_dma_cmd,
+        "MTE2_L2": t_mte2_l2,
+        "MMAD": t_mmad,
+        "FIXPIPE": t_fixpipe,
+        "REDUCE": t_reduce,
+    }
+    bottleneck = max(stages, key=stages.get)
+    t_steady = max(stages.values())
+    return HardwareTiming(
+        gm_read_bytes=gm_read_bytes, l2_read_bytes=l2_read_bytes,
+        t_mte2_gm=t_mte2_gm, t_mte2_l2=t_mte2_l2, t_mte2=t_mte2,
+        dma_cmd_count=dma_cmd_count, t_dma_cmd=t_dma_cmd,
+        cube_flops=cube_flops, t_mmad=t_mmad,
+        fixpipe_bytes=fixpipe_bytes, t_fixpipe=t_fixpipe,
+        t_reduce=t_reduce,
+        t_steady=t_steady, t_drain=t_drain, t_total=t_steady + t_drain,
+        bottleneck=bottleneck,
+    )
+
+
+def bound_type_of(bottleneck: str) -> str:
+    return {
+        "MMAD": "计算Bound",
+        "MTE2_GM": "访存Bound(GM)",
+        "MTE2_L2": "访存Bound(L2)",
+        "FIXPIPE": "写出Bound",
+        "REDUCE": "归约Bound",
+    }.get(bottleneck, "")