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

Fix review issues #11-#16: StreamK fixpipe 单次计账 / K=1 AIV单缓冲方案 / MergeBatch b0 L0A/L0B 上限+路由可行回退 / advice-StreamK / 输入校验 / .gitignore+死代码清理

Browse Source
This commit is contained in:
admin
2026-09-03 20:05:14 +08:00
parent 99a25a6b42
commit 4bedf0a109
14 changed files with 250 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
BMM/BMM_Theory/bmm_theory/branches/merge_batch.py
View File

@@ -94,10 +94,14 @@ class MergeBatchBranch(Branch):
b = case.batch_c
b_core = b // s.aic_num
# Step 1: 合并数 b0 (L0C + 算存比双上限, 尽量取 b_core 的因子)
# Step 1: 合并数 b0 (L0C + L0A/L0B + 算存比 + b_core 四类上限, 尽量取 b_core 的因子)
# L0A/L0B 上限 (issue#13): 合并 tile = (b0*M)x(b0*N), base_k 有 16 (fractal) 硬底,
# 须满足 b0*M*16*dt*2 <= L0A 且 b0*N*16*dt*2 <= L0B, 否则合并后 L0 tile 无法驻留.
b0_l0c = math.sqrt(s.l0c_bytes / (2 * m * n * 4))
b0_ai = s.r16 * (m + n) / (2 * m * n)
b0_max = int(min(b0_l0c, b0_ai, b_core))
b0_l0a = s.l0a_bytes / (2 * m * s.fractal * dt)
b0_l0b = s.l0b_bytes / (2 * n * s.fractal * dt)
b0_max = int(min(b0_l0c, b0_ai, b0_l0a, b0_l0b, b_core))
b0 = max(MIN_B0, self._factor_floor(b_core, b0_max))
# Step 2: L0 级 K 粒度 k_L0

22
BMM/BMM_Theory/bmm_theory/branches/special.py
View File

@@ -24,11 +24,14 @@ class SpecialBranch(Branch):
c1 = case.k <= 1
checks = [ConditionCheck("1_K<=1 (Cube 无用)", c1, f"K={case.k}")]
if case.k == 1:
# K=1 触发 AIV 通路需 B >= 2*AIV核数 且单 batch 输入输出能驻留 UB
c2 = case.batch_c >= 2 * self.spec.aiv_num
# K=1 的 AIV 通路恒可用 (issue#12): B>=2*AIV 开 UB 乒乓; B<128 退化为
# AIV 单缓冲 (无乒乓, 逐 batch 串行搬入), 不再是无方案空洞.
b = case.batch_c
pingpong = b >= 2 * self.spec.aiv_num
mode = "UB乒乓" if pingpong else "AIV单缓冲(逐batch串行, B<2*AIV)"
checks.append(ConditionCheck(
"2_K=1的AIV触发: B >= 2*AIV核数 (开UB乒乓)",
c2, f"B={case.batch_c} vs {2*self.spec.aiv_num}"))
"2_K=1的AIV通路: 恒可用 (B>=128 开UB乒乓, 否则单缓冲)",
True, f"B={b}, 模式={mode}"))
return checks
# ------------------------------------------------------------------
@@ -36,11 +39,15 @@ class SpecialBranch(Branch):
s = self.spec
if case.k == 0:
sub = "K=0纯写值"
mode = ""
note = "无任何计算, C=bias 或 0, 纯 AIV 写值; 按行均分到 AIV 核"
else:
sub = "K=1逐元素乘"
note = ("退化为 C=A⊙B 无累加深度, Cube 16x16x16 粒度浪费 15/16; "
"走 AIV 通路 GM->UB->Mul->GM, UB 乒乓")
pingpong = case.batch_c >= 2 * s.aiv_num
mode = "UB乒乓" if pingpong else "AIV单缓冲"
note = (f"退化为 C=A⊙B 无累加深度, Cube 16x16x16 粒度浪费 15/16; "
f"走 AIV 通路 GM->UB->Mul->GM, {mode} "
f"({'B>=2*AIV 双batch乒乓流水' if pingpong else 'B<2*AIV 逐batch单缓冲串行'})")
return ImplPlan(
case_id=case.case_id, branch=self.name,
used_core_num=s.aiv_num, # 用 AIV 核
@@ -48,7 +55,8 @@ class SpecialBranch(Branch):
core_map="AIV 核间按行均分 (无 Cube tile 概念)",
b_core=0, merge_b0=1,
single_core_m=0, single_core_n=0, single_core_k=case.k,
k_l1=0, b_l1=1, l1_form="UB驻留(AIV)",
k_l1=0, b_l1=1,
l1_form="UB驻留(AIV)" if case.k == 0 else "UB驻留(AIV) " + mode,
base_m=0, base_n=0, base_k=0,
l2_policy_in="allocate", l2_policy_out="direct_gm",
swizzle_w=0, workspace_bytes=0,

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

@@ -145,13 +145,14 @@ class StreamKBranch(Branch):
t_mmad = t_mmad_tile / grid_k
t_mte2 = t_mte2_tile / grid_k
# 归约: 部分和 4B 驻留 L2, AIV 归约 (含最终按 C dtype 写回)
# 归约: 部分和 4B 驻留 L2, AIV 归约 (含部分和写/读回/求和/最终按 C dtype 写回)
# 口径 (issue#11): 归约整体为串行追加 (t_drain=t_reduce, reduce_serial=True),
# 部分和写出已计入 eval_streamk_reduce 的 t_write_partial —— 稳态 Fixpipe 不再
# 重复计账. 此前按 grid_k*tile*4B/单核带宽份额另计一次, 既重复计账又把整组
# 部分和串行压到单核写口, 高估 grid_k 倍 (streamk_demo 曾虚高到 55us/FIXPIPE).
t_reduce = eval_streamk_reduce(tile_elems, grid_k, out_b, s)
# Fixpipe: 部分和写出按 4B (L0C dtype, 防精度丢失), 驻留 L2.
# 最终归约结果的 C dtype 写回已在 t_reduce 内计, 此处不重复 (issue#9 口径对齐).
fix_bytes = grid_k * tile_elems * 4
t_fix = fix_bytes / s.bw_l2_pc
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