diff --git a/BMM/BMM_Theory/bmm_theory/models.py b/BMM/BMM_Theory/bmm_theory/models.py new file mode 100644 index 0000000..e327e75 --- /dev/null +++ b/BMM/BMM_Theory/bmm_theory/models.py @@ -0,0 +1,270 @@ +"""核心数据模型: Case 输入 / 实现方案结构体 / 评估结果. + +约定: +- 所有"时延"内部统一用秒 (float), 输出 csv 时转 us; +- 所有"字节数"用 int (Byte); +- dtype 一律归一化为小写字符串, 如 "bf16"/"fp16"/"fp8"/"fp32". +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field, asdict +from typing import Optional + + +# --------------------------------------------------------------------------- +# dtype 工具 +# --------------------------------------------------------------------------- + +DTYPE_BYTES = { + "fp32": 4, + "f32": 4, + "tf32": 4, + "fp16": 2, + "f16": 2, + "bf16": 2, + "fp8": 1, + "fp8_e4m3": 1, + "fp8_e5m2": 1, + "int8": 1, +} + +# Cube 累加器 (L0C) 中元素字节数: 16bit 输入 -> fp32 累加; fp8 输入 -> fp32 累加 +L0C_DTYPE_BYTES = 4 + + +def dtype_bytes(dtype: str) -> int: + key = dtype.strip().lower() + if key not in DTYPE_BYTES: + raise ValueError(f"不支持的 dtype: {dtype!r}, 支持 {sorted(DTYPE_BYTES)}") + return DTYPE_BYTES[key] + + +def ceil_div(a: int, b: int) -> int: + return -(-a // b) + + +def align_up(x: int, align: int) -> int: + return ceil_div(x, align) * align + + +def align_down(x: int, align: int) -> int: + return (x // align) * align + + +# --------------------------------------------------------------------------- +# Case 输入 +# --------------------------------------------------------------------------- + +@dataclass +class BmmCase: + """一个 batch_mat_mul_v3 case 的输入描述. + + 对应算子接口: + A: [batch_a, M, K] (可带转置: trans_a=True 表示 [batch_a, K, M]) + B: [batch_b, K, N] (可带转置: trans_b=True 表示 [batch_b, N, K]) + bias: [B, 1, N] 可选 + C: [batch_c, M, N], batch_c = broadcast(batch_a, batch_b) + """ + + case_id: str = "" + batch_a: int = 1 + batch_b: int = 1 + m: int = 1 + n: int = 1 + k: int = 1 + dtype_a: str = "bf16" + dtype_b: str = "bf16" + dtype_c: str = "bf16" # 输出 C 的 dtype (fp16/fp8 时 fixpipe 随路转换) + trans_a: bool = False + trans_b: bool = False + has_bias: bool = False + out_nd: bool = True # 输出是否 ND 格式 (StreamK 要求 ND) + deterministic_level: int = 0 # 确定性等级, >=2 禁用 StreamK + + # ---- 派生属性 ---- + @property + def batch_c(self) -> int: + return max(self.batch_a, self.batch_b) + + @property + def dtype_in_bytes(self) -> int: + # A/B 输入元素字节数 (要求 A/B 同 dtype, 不一致时取较大者并在校验中报 warning) + return max(dtype_bytes(self.dtype_a), dtype_bytes(self.dtype_b)) + + @property + def dtype_out_bytes(self) -> int: + return dtype_bytes(self.dtype_c) + + @property + def flops(self) -> float: + """总计算量 (乘加各计一次).""" + return 2.0 * self.batch_c * self.m * self.n * self.k + + @property + def input_bytes(self) -> float: + """输入总数据量 (按广播前实际存储计).""" + dt = self.dtype_in_bytes + return self.batch_a * self.m * self.k * dt + self.batch_b * self.k * self.n * dt + + @property + def output_bytes(self) -> float: + return self.batch_c * self.m * self.n * self.dtype_out_bytes + + @property + def ai_full(self) -> float: + """case 固有全量算存比 AI_full = 2MNK / (MK + KN + MN) (单 batch).""" + m, n, k = self.m, self.n, self.k + denom = m * k + k * n + m * n + return (2.0 * m * n * k / denom) if denom > 0 else 0.0 + + def to_row(self) -> dict: + d = asdict(self) + return d + + +# --------------------------------------------------------------------------- +# 实现方案结构体 (标准结构体定义) +# --------------------------------------------------------------------------- + +@dataclass +class ImplPlan: + """BMM 实现方案 (理论分析输出 / 用户评估输入 共用的标准结构体). + + 字段分四组: + 1) 分支与核间切分; 2) 核内 tiling; 3) 存储/Cache 策略; 4) 尾轮与流水策略. + + 说明: 本结构体对齐 batch_mat_mul_v3 tiling 的概念层级 + (核间 grid 划分 -> singleCoreM/N/K -> L1 tile -> L0 tile), + 字段名采用理论文档符号, 便于与文档公式直接对照. + """ + + # --- 0) 基本信息 --- + case_id: str = "" + branch: str = "" # 转Matmul / 特殊分支 / MergeBatch / IterBatch / + # StreamK / ASW_Basic / ASW_Basic_降核 + npu: str = "Ascend950PR" + op: str = "batch_mat_mul_v3" + + # --- 1) 核间切分 (grid 级) --- + used_core_num: int = 0 # 实际使用 AIC 核数 (降核时 < C) + split_b: int = 1 # 核间 B 维切分数 + m_cnt: int = 1 # 核间 M 维切分数 + n_cnt: int = 1 # 核间 N 维切分数 + grid_k: int = 1 # 核间 K 维切分数 (StreamK > 1) + core_map: str = "" # 核间分配策略, 如 "B->M->N线性映射+ASW滑窗蛇形(W=4)" + + # --- 2) 核内 tiling --- + b_core: int = 0 # 每核 batch 数 (切 B 分支) + merge_b0: int = 1 # MergeBatch 合并数 (IterBatch/其他 = 1) + single_core_m: int = 0 # 每核输出 tile M + single_core_n: int = 0 # 每核输出 tile N + single_core_k: int = 0 # 每核 K 段长度 (核间切 K 时 < K) + k_l1: int = 0 # GM->L1 的 K 向粒度 + b_l1: int = 1 # L1 内驻留 batch 数 (MergeBatch/IterBatch) + l1_form: str = "" # IterBatch L1 形态: a/b/c/d + base_m: int = 0 # L0 级 tile + base_n: int = 0 + base_k: int = 0 + + # --- 3) 存储/Cache 策略 --- + l2_policy_in: str = "" # 输入 L2 策略: allocate(随路驻留) / non_allocate + l2_policy_out: str = "" # 输出: resident(驻留L2异步回写) / direct_gm(直写GM) + swizzle_w: int = 0 # ASW 滑窗宽度 (0 = 不用) + workspace_bytes: int = 0 # StreamK 中间结果 workspace (驻留 L2) + + # --- 4) 尾轮与流水策略 --- + tail_strategy: str = "" # A0 / A1a / A1b / 方案B (仅切 M/N 类分支) + fixpipe_unitflag: bool = True # fixpipe 开 unitflag 随路搬出 + out_dtype_bytes: int = 2 # fixpipe 写出元素字节数 (C 矩阵 dtype) + + # --- 备注 --- + note: str = "" + + def to_row(self) -> dict: + return asdict(self) + + @staticmethod + def csv_fields() -> list: + return list(ImplPlan.__dataclass_fields__.keys()) + + @staticmethod + def from_row(row: dict) -> "ImplPlan": + """从 csv 行 (字符串字典) 恢复 ImplPlan.""" + kw = {} + for name, f in ImplPlan.__dataclass_fields__.items(): + if name not in row or row[name] is None or str(row[name]).strip() == "": + continue + v = row[name] + if f.type == "int": + kw[name] = int(float(v)) + elif f.type == "bool": + kw[name] = str(v).strip().lower() in ("1", "true", "yes", "y") + else: + kw[name] = v + return ImplPlan(**kw) + + +# --------------------------------------------------------------------------- +# 评估结果 +# --------------------------------------------------------------------------- + +@dataclass +class HardwareTiming: + """各硬件流水级时延 (秒) 与数据量明细.""" + + # 搬入 (MTE2) + gm_read_bytes: float = 0.0 # GM->L1 直读数据量 (不驻留/未命中 L2 的部分) + l2_read_bytes: float = 0.0 # L2->L1 数据量 (驻留 L2 后重复读命中部分) + t_mte2_gm: float = 0.0 # GM->L1 时延 (按 GM 带宽, 不累加 L2->L1) + t_mte2_l2: float = 0.0 # L2->L1 时延 (按 L2 带宽) + t_mte2: float = 0.0 # 搬入合计 = t_mte2_gm + t_mte2_l2 (两者发生在不同数据上) + dma_cmd_count: float = 0.0 # GM->L1 DMA 命令次数 (T_cmd 分析用) + t_dma_cmd: float = 0.0 # DMA 命令固定开销合计 + + # 计算 (Cube MMAD) + cube_flops: float = 0.0 # Cube 实际计算量 (MergeBatch 含冗余) + t_mmad: float = 0.0 + + # 搬出 (Fixpipe) + fixpipe_bytes: float = 0.0 # 写出数据量 (按 C 矩阵 dtype / StreamK 临时矩阵按 4B) + t_fixpipe: float = 0.0 + + # 归约 (StreamK 专用) + t_reduce: float = 0.0 + + # 汇总 + t_steady: float = 0.0 # 稳态流水时延 = max(各级) + t_drain: float = 0.0 # 流水排空暴露 + t_total: float = 0.0 # 端到端时延 + bottleneck: str = "" # 瓶颈级: MTE2_GM / MTE2_L2 / MMAD / FIXPIPE / REDUCE + + def to_row(self, prefix: str = "") -> dict: + return {prefix + k: v for k, v in asdict(self).items()} + + +@dataclass +class EvalResult: + """单个 case 的完整评估输出 (csv 一行的内容).""" + + case: BmmCase = field(default_factory=BmmCase) + plan: ImplPlan = field(default_factory=ImplPlan) + timing: HardwareTiming = field(default_factory=HardwareTiming) + feasible: bool = True # 方案是否满足硬件约束 + violations: str = "" # 违反的约束列表 (";" 分隔) + bound_type: str = "" # 计算Bound / 访存Bound / 写出Bound / 归约Bound + advice: str = "" # 瓶颈分析与优化建议 + + def to_row(self) -> dict: + row = {} + row.update(self.case.to_row()) + row.update({"plan_" + k: v for k, v in self.plan.to_row().items()}) + row.update(self.timing.to_row()) + row.update({ + "feasible": self.feasible, + "violations": self.violations, + "bound_type": self.bound_type, + "advice": self.advice, + }) + return row