阶梯 A 方案设计（ldmatrix vs ld.shared）—— 讨论稿 St1

我先读了 experiments/初步规划SM80.md 和仓库里相关脚本，下面是完整的设计说明（没有改动任何文件）。

⚠️ 先说环境：本机 nvidia-smi 不存在、/dev/nvidia* 不存在、torch 也没装 → 这台机器只能做设计和静态检查，编译和实测要在你的 SM80 机器上做。

1. 该看哪些脚本（阅读顺序）

2. 阶梯 A 的实验设计

2.1 固定量（其余全锁死）

GEMM: M=N=K=4096, bf16
BM=BN=128, BK=32, 256 线程 / 8 warp, warp tile 64×32, mma.m16n8k16
LOAD = cp.async.cg 16B，STAGES = 3，SWIZZLE = 1（BCF）

只有 smem→reg 这一步在变。

2.2 档位与 LSU 指令数（推导见 §3）

文档表 2 里只有 A1/A2/A3。我建议加 A2L，原因：A2 换掉了 smem 里 B 的布局（还换掉了 gmem 侧 B 的形态，见 §3.5），如果直接拿 A2 和 A3 比，就是两个变量同时变（读指令 + 搬运形态）。A2L 让 A2 和它布局完全相同、搬运完全相同，只差"ld.shared 还是 ldmatrix"这一个变量，这才是 §1.1 说的"必须给 ld.shared 最有利布局"的严格版本。

2.3 必须先承认的物理事实（决定实验结论怎么读）

• 搬运字节数三者完全相同：A 片段 16×16×2B=512B、B 片段 16×8×2B=256B。ldmatrix.x4 每 lane 拿 16B、ld.shared.b32×4 也是每 lane 16B → smem 带宽消耗一样。
• 所以 ldmatrix 省的是：① LSU issue slot（32→8 条指令）② 地址计算 ALU（swizzle XOR、地址加法）③ bank conflict。
• 带宽不变 ⇒ 在 compute-bound 的 4096³ 上，差异可能只有 1~5%，甚至被完全掩盖（规划文档「坑 2」）。A100 上 smem 带宽是 128B/clk/SM：按本 kernel 的数据复用率算，峰值 TC 时 smem 带宽占用约 75% → 属于"逼近瓶颈"，所以差异到底显不显现是本次实验最有价值的结论之一。
• ⇒ 必须配一个 smem→reg 微基准（只 load 不 mma，循环 N 次），把读路径单独测出来。这是阶梯 A 能否出结论的关键，不是可选项。

2.4 预期现象的解释框架（填数时用）

3. 关键技术推导（这部分决定你能不能自己写出来）

3.1 mma.m16n8k16 的 fragment 映射（必须背下来）

记 g = lane>>2（groupID，0..7）、t = lane&3（组内线程号，0..3）。

A 片段（16×16，M×K），4 个 32bit 寄存器：

reg_a[0] = {A[g    ][k0+2t  ], A[g    ][k0+2t+1]}   // 行 0-7,  列 0-7
reg_a[1] = {A[g+8  ][k0+2t  ], A[g+8  ][k0+2t+1]}   // 行 8-15, 列 0-7
reg_a[2] = {A[g    ][k0+2t+8], A[g    ][k0+2t+9]}   // 行 0-7,  列 8-15
reg_a[3] = {A[g+8  ][k0+2t+8], A[g+8  ][k0+2t+9]}   // 行 8-15, 列 8-15

B 片段（16×8，K×N），2 个 32bit 寄存器：

reg_b[0] = {B[k0+2t  ][n0+g], B[k0+2t+1][n0+g]}
reg_b[1] = {B[k0+2t+8][n0+g], B[k0+2t+9][n0+g]}

注意 B 的两个元素是 k 连续、n 固定 → 这正是 A1 痛苦的根源。

3.2 ldmatrix 与 fragment 的对应（现有代码为什么是对的）

• ldmatrix.m8n8.x4：lane 0-7 给 tile0 的 8 个行地址，8-15 给 tile1，16-23 给 tile2，24-31 给 tile3；每个地址指向 8 个连续 b16（16B）。
• 现有代码 a_row = warp_m*64 + m_idx*16 + (lane%16)、a_col = k_offset + (lane/16)*8 ⇒ tile0=行0-7/列0-7，tile1=行8-15/列0-7，tile2=行0-7/列8-15，tile3=行8-15/列8-15 ⇒ r0..r3 正好 = reg_a[0..3] ✔
• 无 .trans 时，thread 拿 M[lane>>2][2t] 和 M[lane>>2][2t+1]；有 .trans 时拿 M[2t][lane>>2] 和 M[2t+1][lane>>2]（M = smem 里那个 8×8）。
  • 现有 B 用 LDMATRIX_X2_TRANS：b_row = k_offset + (lane%16)、b_col = warp_n*32 + n_idx*8 ⇒ tile0 = B 行 k0..k0+7 / 列 n0..n0+7，tile1 = 行 k0+8.. ⇒ 转置后 thread 拿 B[k0+2t][n0+g]、B[k0+2t+1][n0+g] ✔ 正好是 reg_b[0]；tile1 同理给 reg_b[1] ✔

3.3 swizzle 的"透明性"（为什么 ldmatrix + swizzle 是对的）

SWIZZLE_B(row,col) = col ^ ((row&7)<<3)：store 时元素 (k,n) 落到列 n ^ ((k&7)*8)。 ldmatrix 读时给的地址也是 &Bs[r][SWIZZLE_B(r, c)]，第 r 行读到的是物理列块 (c/8)^(r&7)，再 XOR 回去逻辑列恰好是 c..c+7（因为 (c/8)^(r&7)^(r&7) = c/8）⇒ 每个 8×8 tile 逻辑上仍是自然矩阵，转置语义不受影响 ✔

关键推论：ld.shared 变体必须用同一套 swizzle 地址去读，否则读到的是错位数据。

3.4 三个变体的具体寻址（照抄这个写）

As[BM][BK]（行 = m，列 = k），Bs[BK][BN]（行 = k，列 = n）
SWZ_A(r,c) = c ^ (((r>>1)&3)<<3)      SWZ_B(r,c) = c ^ ((r&7)<<3)

A1 / A2 的 A 片段（ld.shared.b32，一次拿 2 个连续 k）：

int g = lane>>2, t = lane&3;
int m0 = m_base + g;          // m_base = warp_m*64 + m_idx*16
int k  = k0 + 2*t;            // k0 = k_step*16
reg_a[0] = lds32(&As[m0    ][SWZ_A(m0,    k  )]);
reg_a[1] = lds32(&As[m0 + 8][SWZ_A(m0 + 8,k  )]);
reg_a[2] = lds32(&As[m0    ][SWZ_A(m0,    k+8)]);
reg_a[3] = lds32(&As[m0 + 8][SWZ_A(m0 + 8,k+8)]);

（lds32 = asm volatile("ld.shared.b32 %0,[%1]";)；建议用 inline PTX 而不是 C++ 表达式，防止编译器把两条相邻 b32 合并成一条 v4，否则指令数不可控。）

A1 的 B 片段（自然布局，k 不连续 → 只能 b16）：

int n = n_base + g;           // n_base = warp_n*32 + n_idx*8
reg_b[0] = lds16(&Bs[k    ][SWZ_B(k,    n)]) | (lds16(&Bs[k+1  ][SWZ_B(k+1,  n)]) << 16);
reg_b[1] = lds16(&Bs[k + 8][SWZ_B(k + 8,n)]) | (lds16(&Bs[k+9  ][SWZ_B(k+9,  n)]) << 16);

→ 4 条 LDS.U16 + 打包 ALU（这条 ALU 是 ld.shared 方案的固有成本，ldmatrix 把它变成了硬件免费的转置）。

A2 的 B 片段（预转置 Bs[BN][BK]，k 连续 → b32）：

reg_b[0] = lds32(&BsT[n][SWZ_BT(n, k    )]);   // BsT[n][k], BsT[n][k+1]
reg_b[1] = lds32(&BsT[n][SWZ_BT(n, k + 8)]);

→ 2 条 LDS.32，无打包 ALU。

A2L 的 B 片段（预转置 + ldmatrix，不加 trans）： 因为 BsT[n][k] 的连续方向是 k，而 ldmatrix 把"地址指向的 8 个连续元素"当作矩阵的行 ⇒ 载入的矩阵 M[r=n][c=k]，M[g][2t] = BsT[n0+g][k0+2t] = B[k0+2t][n0+g] ✔ 正好是 reg_b[0]，所以要 plain ldmatrix.x2：

int n = n_base + lane;                       // lane 0-7 → n0..n0+7 (tile0), lane 8-15 → n0..n0+7 (tile1)
int kk = (lane < 16) ? k0 : k0 + 8;
addr = &BsT[n_base + (lane%8)][SWZ_BT(n_base + (lane%8), kk)];

（lane 16-31 的地址被忽略，写什么都行，但别越界访问 smem。）

3.5 ⚠️ 最大的实现难点：B 的"预转置"从哪来

ld.shared 想要"k 连续"，就必须在 smem 里存成 Bs[BN][BK]。但 cp.async 不能转置（它只能整 16B 直拷）。三种做法：

→ 我建议 T1，并在 readme 里写明 caveat。Python 侧只需把 b_t 提前算好并复用（不计入计时）。

3.6 bank conflict 预分析（用来解释结果，不是设计目标）

• A 片段（ld.shared.b32）：行距 64B。地址 = 64r + 4t + 16·S(r)，bank = (16r + t + 4·S(r)) mod 32，S(r)=(r>>1)&3 ⇒ 8 个 group（r=base..base+7）恰好铺满 32 个 bank 零冲突 ✔（现有 SWIZZLE_A 恰好是为 ldmatrix 设计的，但顺带也让 ld.shared 无冲突——这点要在报告里点出来，说明 A1 的劣势不来自 A 侧）。
• B 片段自然布局（b16）：行距 256B，bank 与 k 无关 ⇒ 单条指令 32 lane 只落在 8 个 bank ⇒ ~4 路冲突；加 SWIZZLE_B 后 XOR 分散到 16 个 bank ⇒ ~2 路冲突。
• B 片段预转置（b32）：行距 64B ⇒ 与 A 侧同构 ⇒ 零冲突（用 SWZ_BT(r,c)=c^(((r>>1)&3)<<3)）。

⇒ 结论预告：A1 的 B 侧 conflict 是它落后的主因之一，A2 把它消掉了。所以 A1 必须配 A2 才有信息量。

4. 代码骨架（新增文件里长什么样）

4.1 kernels/ablation/ablation.cu

// 复用 hgemm.cu 的宏（直接 copy 一份或 #include 拆出的 common 头）
template <int LOAD_MODE,   // 0 = ldg+st.shared, 1 = cp.async
          int READ_MODE,   // 0 = ld.shared,    1 = ldmatrix
          int B_LAYOUT,    // 0 = Bs[BK][BN],   1 = Bs[BN][BK]（gmem 侧 B 需为 [N][K]）
          int STAGES,      // 1..4
          int SWIZZLE,     // 0/1
          typename T>
__global__ void hgemm_ablation_kernel(T* a, T* b, T* c, int m, int n, int k)

内部用 if constexpr 分三块：

1. smem 声明：用 dynamic smem（extern __shared__ char smem_raw[];）+ 启动时 cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes)。

   必须 dynamic：STAGES=3 → 3×(8KB+8KB) = 48KB，静态 __shared__ 的 48KB 上限很紧；STAGES=4 → 64KB 更是必须 opt-in（A100 每 block 上限 163KB）。 布局：As[STAGES][BM][BK] + Bs[STAGES][BK][BN]（或 [BN][BK]），stage-major，offset 用 constexpr 算。

2. gmem→smem 装载：LOAD_MODE==1 用 CP_ASYNC_CG（B_LAYOUT==1 时源指针换成 b_t[(n0+row)*K + k0]）；LOAD_MODE==0 用 ld.global.v4 + st.shared.v4 + __syncthreads()。
3. smem→reg + mma：READ_MODE 决定用 §3.4 的哪一套寻址；mma body 与 hgemm_bcf_dbf 完全一致（16 次 M16N8K16_BF16）。

多级流水骨架（STAGES 通用，比现有 bcf_dbf 的硬编码 2 更清楚）：

for (int s = 0; s < STAGES-1; ++s) { load_tile(s, /*stage*/s); commit_group(); }
for (int kt = 0; kt < num_tiles; ++kt) {
    cp_async_wait_group<STAGES-2>();      // 保证 stage kt 的那一组已落地
    __syncthreads();                       // 也保证所有人不再读被覆盖的 stage
    int rs = kt % STAGES;
    if (kt + STAGES-1 < num_tiles) load_tile(kt+STAGES-1, (kt+STAGES-1)%STAGES);
    commit_group();
    mma_on_stage(rs);
}

（现有 bcf_dbf 写死了 2 buffer、且把最后一块 MMA 放在 loop 外做 epilogue；改成模板后要重写这段，注意 wait_group 是按组计数，STAGES=3 时 wait_group<1> 才对，用 wait_group<0> 会退化成单 buffer —— 规划文档 §1.4 第 4 条。）

4.2 变体注册（照抄 binding_tiled_func_gen 的套路）

template <int LM, int RM, int BL, int ST, int SW>
void abl_launcher(torch::Tensor a, torch::Tensor b, torch::Tensor c) { ... }

#define ABL_BIND(NAME, LM, RM, BL, ST, SW) \
  void NAME(torch::Tensor a, torch::Tensor b, torch::Tensor c) { abl_launcher<LM,RM,BL,ST,SW>(a,b,c); }

ABL_BIND(abl_A1, 1, 0, 0, 3, 1)   // ld.shared, 自然 B
ABL_BIND(abl_A2, 1, 0, 1, 3, 1)   // ld.shared, 预转置 B (b 传 [N][K])
ABL_BIND(abl_A2L,1, 1, 1, 3, 1)   // ldmatrix,  预转置 B
ABL_BIND(abl_A3, 1, 1, 0, 3, 1)   // ldmatrix, 自然 B  (== 现有 hgemm_bcf_dbf 的读路径)

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    torch_pybinding_func(hgemm_cublas);   // extern 声明，链接 cublas.cu
    torch_pybinding_func(abl_A1); ... torch_pybinding_func(abl_A3);
}

4.3 kernels/ablation/micro.cu（smem→reg 微基准，强烈建议）

同一个 kernel 模板加一个 DO_MMA 开关：

• DO_MMA=0：只循环做 fragment 载入 + acc ^= reg（防 DCE），最后写一个标量 → 纯读路径吞吐。
• DO_MMA=1：加上 16 次 mma → "读 + 算"。 两者一减，差值就是"读路径在真实 kernel 里被掩盖了多少"。这是让阶梯 A 出结论的关键。

5. 会新增哪些文件

kernels/ablation/
├── ablation.cu       # 模板化 kernel（LOAD/READ/B_LAYOUT/STAGES/SWIZZLE）+ A1/A2/A2L/A3 binding
├── test.py           # 编译 + 正确性 + 计时 + 打印 Markdown 表（表 2）
├── test_for_ncu.py   # 单次运行版，供 ncu 采样（照抄 hgemm/test_for_ncu.py）
├── run_ncu.sh        # ncu -k regex:"ablation" --set full -o prof -f python test_for_ncu.py
├── micro.cu          # smem→reg 微基准（只 load / load+mma 两种）
├── micro_test.py     # 微基准的计时脚本
└── readme.md         # 结果表（照 sgemm/readme.md 的格式：每档一行 ms + TFLOPS）

不修改 kernels/hgemm/*（保持历史数据可复现）；若想复用 cublas.cu，在 test.py 的 sources 里加 ../hgemm/cublas.cu 即可。

6. .py 是怎么触发 kernel 的（完整链路）

python test.py
   │
   ├─(1) torch.utils.cpp_extension.load(name="ablation_lib",
   │        sources=["ablation.cu", "../hgemm/cublas.cu"],
   │        extra_cuda_cflags=["-O3","-std=c++17", "-U__CUDA_NO_HALF_OPERATORS__", ...,
   │                           "--use_fast_math", "-Xptxas -v"],
   │        extra_cflags=["-O3","-std=c++17"], verbose=True)
   │        └─→ ninja 现场编译成 ablation_lib.so（缓存在 ~/.cache/torch_extensions），
   │           作为一个普通 python module 被 import
   │
   ├─(2) C++ 侧：ablation.cu 里的 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) 执行，
   │        里面 torch_pybinding_func(abl_A1) 展开成 m.def("abl_A1", &abl_A1, "abl_A1")
   │        → 于是 python 侧出现属性 lib.abl_A1
   │
   ├─(3) python: lib.abl_A1(a, b, c)  → 调进 C++ 函数：
   │        · CHECK_T(a/b/c)：必须是 cuda + contiguous
   │        · 读 M/K/N，判断 dtype（bf16 / fp16）
   │        · 计算 grid（(N+BN-1)/BN, (M+BM-1)/BM）、256 线程
   │        · cudaFuncSetAttribute(..., MaxDynamicSharedMemorySize, 48KB+)  ← SM80 上 >48KB 必须
   │        · kernel<<<grid, 256, 0, at::cuda::getCurrentCUDAStream()>>>(...)
   │        · kernel 内部自己完成 grid swizzle（linear_id → bx,by）、load、mma、写回 C
   │
   └─(4) python: benchmark() 计时
            warmup 10 次 → torch.cuda.synchronize() → perf_counter() 计时 rep 次 → sync
            TFLOPS = M*N*K*2*rep / duration / 1e12
            diff_check(c, torch.matmul(...)) 校验

两个必做的环境准备：

• export TORCH_CUDA_ARCH_LIST=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | head -n 1) → 否则 sm_80 可能不被编出来。
• 测性能前锁频（nvidia-smi -lgc/-lmc，测完 -rgc/-rmc），否则噪声 > 效应（规划文档「硬性要求」）。

7. 怎么验证"写对了"

1. 数值：A1/A2/A2L/A3 四档结果必须逐位相同（ldmatrix.trans 只是重排，不换算）→ torch.equal 级别 diff = 0，与 torch.matmul 的 diff 只需在 bf16 误差内。
2. 指令数：cuobjdump -sass ablation_lib.so | grep -c "LDSM\|LDS" 数每档的 LDS/LDSM 条数，应严格等于 32/24/8/8（不含 epilogue）。
3. ncu 指标：
   • smsp__inst_executed_pipe_lsu_mem_shared_op_ld.sum（LSU 条数）
   • l1tex__data_bank_conflicts_pipe_lsu_mem_shared_op_ld.sum（bank conflict）
   • sm__inst_executed_pipe_lsu.avg.pct_of_peak_sustained_active
   • smsp__warp_issue_stalled_long_scoreboard_per_warp_active.pct / ..._mio_throttle_...
4. smem 用量：-Xptxas -v 会打印每个模板实例的 smem/寄存器 → 确认 STAGES=3 是 48KB、占用率没被压死。

8. 需要你先拍板的 4 件事（我再动手写文件）

1. B 预转置用 T1（host 侧 b.t().contiguous()）可以吗？ 我推荐 T1。代价：A2/A2L 的 gmem 侧形态变了、计时不含转置开销（会写进 readme 的 caveat）。若你坚持 kernel 内转置，就要改成 T2，代价是搬运路径变了（实验就不干净了）。
2. 加不加 A2L 这一档？ 我建议加（否则 A2 vs A3 是双变量）。加的话 A2L 需要为预转置布局单独设计一个 SWZ_BT 和一套 ldmatrix 寻址。
3. STAGES 固定 3 还是同时测 2/4？ 阶梯 A 的主表建议固定 STAGES=3；但 STAGES=4 需要 64KB dynamic smem + cudaFuncSetAttribute，建议作为"表 4 形状敏感性"的一部分再展开。
4. epilogue 用哪个？ bcf_dbf 的简单写回（c 直接按 lane 映射写，coalescing 一般）还是 bcf_dbf_rw 的 smem 暂存 + coalesced 写回？4096³ 下 epilogue 影响 <1%，我倾向先用简单的 bcf_dbf 版本（少一堆 union/smem 复杂度），把 rw 留到阶梯 B。

另外确认一下范围：这次只做阶梯 A（表 2），还是顺带把阶梯 B 的 LOAD_MODE 模板骨架先搭好（但不填 B 档）？我建议骨架一次搭好（LOAD_MODE 参数位留着），这样阶梯 B 后面只加变体、不动结构。