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

Fix issues #23-#25 (+#26 标注): GM读写共享总线累加计时 / ASW-切M/N 共享块 GM首读1次+L2重复读(n-1)次、场景按单batch判定、分组预算不再除B / StreamK 按plan实际tile芯片口径评估 / 降核线性带宽假设文档标注

Browse Source
This commit is contained in:
admin
2026-09-04 11:43:23 +08:00
parent f4b23d9f05
commit 4843053ad3
10 changed files with 248 additions and 180 deletions
Download Patch File Download Diff File Expand all files Collapse all files

94
BMM/BMM_Theory/bmm_theory/branches/asw_basic.py
View File

@@ -80,21 +80,22 @@ class AswBasicBranch(Branch):
# Step 4: swizzle 窗口 W = max{d | d|C, d <= floor(sqrt(C))}
swizzle_w = self._swizzle_w()
# Step 5: L2 分组判断
s_in = b * (m * k + k * n) * dt
s_out = b * m * n * case.dtype_out_bytes
if s_in + s_out <= s.l2_bytes:
l2_scene = "A_全驻留"
# Step 5: L2 分组判断 (issue#24: 按单 batch 工作集判定, 同一时刻单 batch 激活)
a_b = m * k * dt
bb_b = k * n * dt
out_batch = m * n * case.dtype_out_bytes
if a_b + bb_b + out_batch <= s.l2_bytes:
l2_scene = "A_单batch全驻留(输入+输出)"
l2_out = "resident(输出驻留L2异步回写)"
r_in = 1.0
elif s_in <= s.l2_bytes:
l2_scene = "B_输入驻留输出直写GM"
elif a_b + bb_b <= s.l2_bytes:
l2_scene = "B_单batch输入驻留输出直写GM"
l2_out = "direct_gm(输出直写GM不占L2)"
r_in = 1.0
else:
l2_scene = "C_输入超L2分组执行"
l2_scene = "C_单batch输入超L2分组执行"
l2_out = "direct_gm(输出直写GM不占L2)"
r_in = self._l2_group_r_in(case, single_m, single_n)
r_in = self._l2_group_r_gm(case, single_m, single_n)
# Step 6: k_l1 (GM->L1 K 向粒度, 须 >= dValue 下限; K 小于下限时整 K 一次搬入不切)
dv_min_elems = max(s.dvalue_hw_min // dt, 1)
@@ -124,7 +125,7 @@ class AswBasicBranch(Branch):
tail_block_cnt=tail["r"], tail_wave_num=tail["n_wave"],
fixpipe_unitflag=True,
out_dtype_bytes=case.dtype_out_bytes,
note=f"L2场景{l2_scene}, r_in={r_in:.2f}; 尾轮: {tail['reason']}",
note=f"L2场景{l2_scene}, GM首读倍率={r_in:.2f}; 尾轮: {tail['reason']}",
)
# ------------------------------------------------------------------
@@ -189,12 +190,18 @@ class AswBasicBranch(Branch):
w = d
return w
def _l2_group_r_in(self, case, sm, sn) -> float:
"""场景 C: L2 分组的重复读倍率 r_in = (n_grp*M + m_grp*N)/(M+N)."""
def _l2_group_r_gm(self, case, sm, sn) -> float:
"""场景 C (单 batch 输入仍超 L2): 分组执行的 GM 重复读倍率.
每组 (m_grp x n_grp 个 tile) 输入工作集 <= L2; 组间共享块复用落空回 GM.
预算 D 按**单 batch** 计 (issue#24: 线性映射同一时刻只激活 1 个 batch,
不再除以总 batch 数).
r_gm = (ceil(n_cnt/n_grp)*M + ceil(m_cnt/m_grp)*N)/(M+N).
"""
s = self.spec
m, n, k, b = case.m, case.n, case.k, case.batch_c
m, n, k = case.m, case.n, case.k
dt = case.dtype_in_bytes
d = s.l2_bytes / (b * k * dt)
d = s.l2_bytes / (k * dt)
m_grp = max(1, int(d / (2 * sm)))
n_grp = max(1, int(d / (2 * sn)))
m_cnt = ceil_div(m, sm)
@@ -261,37 +268,64 @@ class AswBasicBranch(Branch):
# ------------------------------------------------------------------
def evaluate(self, case: BmmCase, plan: ImplPlan) -> HardwareTiming:
"""MTE2 两段块级模型 (issue#24, 用户澄清):
- 首读走 GM (按 GM 带宽, 不叠加 L2); 共享块 (A 行块被 n_cnt 个 tile 读、
B 列块被 m_cnt 个 tile 读) 驻留 L2 后其余 (n_cnt-1)/(m_cnt-1) 次读走
L2 读口 (5.2TB/s 独享) —— 场景 A/B (单 batch 工作集可驻留);
- 单 batch 输入超 L2 -> 场景 C 分组执行, GM 重复按组间落空计 (r_gm), 窗口内
复用未另计 (保守);
- 输出: 场景 A 驻留 L2 (5.2 写口) / 场景 B,C 直写 GM —— GM 读写共享总线
累加由 assemble 的 MTE2 链处理 (issue#23).
"""
s = self.spec
b = case.batch_c
m, n, k = case.m, case.n, case.k
dt = case.dtype_in_bytes
out_b = case.dtype_out_bytes
used = max(plan.used_core_num, 1)
flops = 2.0 * b * m * n * k
t_mmad = flops / (plan.used_core_num * s.q16)
t_mmad = flops / (used * s.q16)
in_bytes = b * (m * k + k * n) * dt
out_bytes = b * m * n * out_b
# r_in 从 note 里解析困难, 重新计算
s_in = in_bytes
s_out = out_bytes
if s_in + s_out <= s.l2_bytes or s_in <= s.l2_bytes:
r_in = 1.0
# ---- 字节量 (每 batch 口径) ----
a_b = m * k * dt # 每 batch A 字节
bb_b = k * n * dt # 每 batch B 字节
out_all = b * m * n * out_b # 输出总字节
m_cnt = max(plan.m_cnt, 1)
n_cnt = max(plan.n_cnt, 1)
# ---- 场景判定: 按单 batch 工作集 (同一时刻单 batch 激活) ----
if a_b + bb_b + m * n * out_b <= s.l2_bytes:
scene, to_l2, r_gm = "A_单batch全驻留(输入+输出)", True, 1.0
elif a_b + bb_b <= s.l2_bytes:
scene, to_l2, r_gm = "B_单batch输入驻留,输出直写GM", False, 1.0
else:
r_in = self._l2_group_r_in(case, plan.single_core_m, plan.single_core_n)
scene, to_l2, r_gm = "C_单batch输入超L2分组执行", False, \
self._l2_group_r_gm(case, plan.single_core_m, plan.single_core_n)
t_mte2 = r_in * in_bytes / (plan.used_core_num * s.bw_pc)
to_l2 = "resident" in plan.l2_policy_out
t_fix = out_bytes / (plan.used_core_num * (s.bw_l2_pc if to_l2 else s.bw_pc))
# ---- MTE2: GM 段 (首读, 每字节一次) + L2 段 (共享块重复读) ----
gm_read = r_gm * b * (a_b + bb_b)
if scene.startswith(("A_", "B_")):
# 驻留命中: A 行块被 n_cnt 个 tile 复用 -> 其余 (n_cnt-1) 次走 L2 读口
l2_read = b * ((n_cnt - 1) * a_b + (m_cnt - 1) * bb_b)
else:
l2_read = 0.0 # 场景 C: 组间复用落空(已计 GM), 窗口内复用保守不计
t_gm = gm_read / (used * s.bw_pc)
t_l2 = l2_read / (used * s.bw_l2_pc)
# ---- Fixpipe ----
t_fix = out_all / (used * (s.bw_l2_pc if to_l2 else s.bw_pc))
# drain: 尾轮暴露 (方案 B 已均匀重切, drain 小; A1b 尾轮凑满, drain 小; A0 尾轮 r 核空转)
t_drain = 0.0
if plan.tail_strategy == "A0" and plan.tail_block_cnt > 0:
t_drain = max(t_mmad, t_mte2, t_fix) # 尾轮空转一个整块
t_drain = max(t_mmad, t_gm + t_l2, t_fix) # 尾轮空转一个整块
return assemble_timing(
t_mte2_gm=t_mte2, t_mte2_l2=0.0, t_dma_cmd=0.0,
t_mte2_gm=t_gm, t_mte2_l2=t_l2, t_dma_cmd=0.0,
t_mmad=t_mmad, t_fixpipe=t_fix, t_reduce=0.0, t_drain=t_drain,
gm_read_bytes=r_in * in_bytes, l2_read_bytes=0.0, dma_cmd_count=0.0,
cube_flops=flops, fixpipe_bytes=out_bytes,
gm_read_bytes=gm_read, l2_read_bytes=l2_read, dma_cmd_count=0.0,
cube_flops=flops, fixpipe_bytes=out_all,
fixpipe_to_gm=(not to_l2),
)

63
BMM/BMM_Theory/bmm_theory/branches/stream_k.py
View File

@@ -129,40 +129,59 @@ class StreamKBranch(Branch):
# ------------------------------------------------------------------
def evaluate(self, case: BmmCase, plan: ImplPlan) -> HardwareTiming:
"""芯片口径评估, 按 plan 实际 tile 布局 (issue#25).
切 K 组 (tile = single_m x single_n) 由 grid_k 核协作: 组内 K 段零重复读
(GM 一次); 组间共享 (A 行块被 n_cnt 个 tile 用) 在单 batch 工作集可驻留时
走 L2 读口 (与 ASW 同规则). 全部 tile-组在 C 核上分波执行:
waves = ceil(b*m_cnt*n_cnt*grid_k / C).
- GM 读: 首读总量 = b*(m*k + k*n)*dt (组间共享命中 L2 时), 或按 r_gm 放大
(工作集超 L2, 保守同 ASW 场景 C 公式);
- MMAD: 芯片总量 2*b*m*n*k / (C核) —— 全核忙稳态;
- 归约 (部分和 4B 写 L2/AIV 读回求和/最终写回): 每 tile-组一次、组间串行
追加为 drain: t_drain = waves * eval_streamk_reduce(tile, grid_k, out).
"""
s = self.spec
b = case.batch_c
m, n, k = case.m, case.n, case.k
dt = case.dtype_in_bytes
out_b = case.dtype_out_bytes
grid_k = plan.grid_k
grid_k = max(plan.grid_k, 1)
m_cnt = max(plan.m_cnt, 1)
n_cnt = max(plan.n_cnt, 1)
tile_m = max(plan.single_core_m, 1)
tile_n = max(plan.single_core_n, 1)
tile = tile_m * tile_n # 每核实际 tile 元素数
# 单 tile (L0C 满载基本块)
tile_elems = s.l0c_elems # 65536
t_mmad_tile = 2.0 * k * s.l0c_bytes / 4 / s.q16 # 2K/Q16 * L0C/4B
t_mte2_tile = k * (2 * math.sqrt(tile_elems)) * dt / s.bw_pc # 近似方形 tile
# ---- MTE2: GM 段 + L2 段 (同 ASW 块级规则) ----
a_b = m * k * dt
bb_b = k * n * dt
if a_b + bb_b <= s.l2_bytes:
gm_read = b * (a_b + bb_b) # 组间共享命中 L2: GM 每字节一次
l2_read = b * ((n_cnt - 1) * a_b + (m_cnt - 1) * bb_b)
else:
gm_read = b * (a_b + bb_b) # 保守: 共享跨组回落 GM 未另计
l2_read = 0.0
t_gm = gm_read / s.bw_gm # 芯片口径 (全核并发)
t_l2 = l2_read / s.bw_l2
t_mte2 = t_gm + t_l2
# 切 K 后流水时延缩 grid_k 倍
t_mmad = t_mmad_tile / grid_k
t_mte2 = t_mte2_tile / grid_k
# ---- MMAD: 芯片总量 ----
flops = 2.0 * b * m * n * k
t_mmad = flops / (s.aic_num * s.q16)
# 归约: 部分和 4B 驻留 L2, AIV 归约 (含部分和写/读回/求和/最终按 C dtype 写回)
# 口径 (issue#11/#17): 归约整体为串行追加 (t_drain=t_reduce, reduce_serial=True),
# 部分和写出已计入 eval_streamk_reduce 的 t_write_partial —— 稳态 Fixpipe 不再
# 重复计账. 若再按 grid_k*tile*4B/单核带宽份额另计一次, 既重复计账又把整组
# 部分和串行压到单核写口, 高估 grid_k 倍 (streamk_demo 曾虚高到 55us/FIXPIPE;
# 第三轮曾回退该修复, issue#17 恢复).
t_reduce = eval_streamk_reduce(tile_elems, grid_k, out_b, s)
# ---- 归约: 每 tile-组一次, 组间分波串行追加 ----
waves = ceil_div(b * m_cnt * n_cnt * grid_k, s.aic_num)
t_reduce_group = eval_streamk_reduce(tile, grid_k, out_b, s)
t_reduce = waves * t_reduce_group
fix_bytes = 0.0
t_fix = 0.0
flops_pc = 2.0 * tile_elems * k / grid_k
gm_bytes = k * (2 * math.sqrt(tile_elems)) * dt / grid_k
return assemble_timing(
t_mte2_gm=t_mte2, t_mte2_l2=0.0, t_dma_cmd=0.0,
t_mte2_gm=t_gm, t_mte2_l2=t_l2, t_dma_cmd=0.0,
t_mmad=t_mmad, t_fixpipe=t_fix, t_reduce=t_reduce,
t_drain=t_reduce, # 归约串行追加 (reduce_serial 默认 True, 不进稳态 max)
gm_read_bytes=gm_bytes, l2_read_bytes=0.0, dma_cmd_count=0.0,
cube_flops=flops_pc, fixpipe_bytes=fix_bytes,
reduce_serial=True,
gm_read_bytes=gm_read, l2_read_bytes=l2_read, dma_cmd_count=0.0,
cube_flops=flops, fixpipe_bytes=fix_bytes,
reduce_serial=True, fixpipe_to_gm=False,
)