admin/matmul-analysis
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add BMM_Theory: bmm_theory/models.py

Browse Source
This commit is contained in:
admin
2026-09-03 08:09:26 +00:00
parent ca6af2789c
commit 0f259b81e1
1 changed files with 270 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

270
BMM/BMM_Theory/bmm_theory/models.py Normal file
View File

@@ -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