HGEMM SMEM Swizzle 代码学习
本人学习笔记，AI总结
09 篇讲 bank conflict 时给了两把刀：padding（行尾塞废位错开 bank）和 CuTe swizzle（位交错）。当时 swizzle 只在 10 篇的 CuTe 转置里露过脸，被标记为"不深入"。本篇把它拆透：17 篇的 MMA kernel 去掉 A_PAD/B_PAD，换上smem swizzle——同样消 bank conflict，但零浪费、零对齐副作用，是 Tensor Core kernel 的工业标配（CUTLASS 内部就这么干）。
实测对照（RTX 4090，方阵）：
| 形状 | 无 swizzle（PAD=0） | + smem swizzle | 提升 |
|---|---|---|---|
| 2048³ | 122.7 | 130.8 | +6.6% |
| 4096³ | 145.8 | 148.5 | +1.8% |
| 8192³ | 153.1 | 151.9 | -0.8% |
| 12288³ | 148.6 | 147.3 | -0.9% |
| 16384³ | 155.9 | 154.3 | -1.0% |
有意思：只有小形状赚，大形状反而微亏。这个反直觉的结果本身就是本篇最好的教材——先把 swizzle 机理讲透，再回来解释为什么。
问题：128bit 访存的 bank conflict 有多隐蔽
复习 bank 布局（09 篇）：smem 32 个 bank，每 bank 4 字节宽。128bit（16 字节）访存是本 kernel 的常态——cp.async 每次搬 16B，ldmatrix 每次读 16B。一条 16B 访存横跨 4 个 bank。
17 篇的 smem 布局（A 矩阵，s_a[BM][BK] = 128 行 × 16 列 half，half 2 字节）：
行 i 的起始字节 = i * 16 * 2 = i * 32 字节 → 行起点 = bank (i*32/4) % 32 = (i*8) % 32
行号: 0 1 2 3 4 ...
起点bank: 0 8 16 24 0(32%32) ...
每行的 16B(4 bank): 行0 占 bank 0-3, 行1 占 bank 8-11, 行2 占 16-19, 行3 占 24-27
行4 又回到 bank 0-3!周期是 4 行：行 0/4/8/12 起点全是 bank 0。17 篇 kernel 的装载映射 load_a_smem_m = tid / 2——warp 内相邻 tid 的 m 差 1，但 16B 写落在 4 行一循环的 bank 上，warp 32 线程 → 8 路冲突（32 线程的写请求挤 4 组重复 bank）。ldmatrix 读端同理。
09 篇的 padding 刀法在这里失灵一半：给 BK=16 加 8 个 half 的 PAD，行距变 48 字节 = 12 bank，周期变 8——冲突降但不消，而且 128×24×2B×stage 数的 smem 膨胀会挤掉一级流水线（smem 总量 96KB 上限，stage2 就要 16KB×2）。大 tile 的 Tensor Core kernel 付不起这个浪费。
swizzle 核心：一次异或，行内块重排
swizzle 的思路和 padding 完全不同：不动布局尺寸，改元素位置。存储时对列坐标做位异或重排，让"物理相邻的行"在 16B 粒度上互相错开。
核心函数（hgemm_mma_stage_tn_swizzle.cu:124，一行就是全部魔法）：
c++
// i: 行号; j: 列号; 返回: 重排后的列号
// kColStride=16 (BK), kStep=8 (128bit = 8 half)
swizzle_permuted_j = (((j >> 3) ^ (i >> 2)) % (kColStride >> 3)) << 3;
// 块号 j/8 XOR 行组 i/4 块总数 2 回到字节粒度拆开读：每行 16 个 half 分成 2 个 8-half 块（每块 16B=128bit）；重排 = 块号 XOR (行号/4)。
具体算（i=行, 块号 0/1，i>>2 的值每 4 行变一次）：
行 0~3 (i>>2=0): 块号 XOR 0 → 不变 [块0|块1]
行 4~7 (i>>2=1): 块号 XOR 1 → 交换 [块1|块0]
行 8~11 (i>>2=2): 块号 XOR 0 → 不变 [块0|块1] ← 2 xor 回 0
行 12~15(i>>2=3): 块号 XOR 1 → 交换 [块1|块0]物理 smem 里相邻 4 行一组交替"正常/交换"。效果：行 0 的 16B 从 bank 0-3 出发，行 4 的 16B 现在取的是原块 1 的数据——bank 8-11。原来 4 行一循环的 bank 重复模式被搅成 8 行才循环，装载/读取时相邻行不再撞 bank。
为什么是 i>>2（4 行一组）而不是 i>>1？必须和访存形态匹配：本 kernel 装载时 2 个 tid 合作搬一行（tid/2=m），warp 内 16 行同时访问——组距 4 保证同 warp 的 32 个 16B 请求均匀落在 8×4=32 个 bank 上。kStep=8（128bit）对齐 cp.async 的 16B 粒度，kColStride=16 对应 BK——参数全部由 kernel 的访存形态决定，不是玄学调参。
写端与读端：一处重排，处处配合
swizzle 是侵入式的：存储布局变了，所有访问 smem 的地方都要用同一个函数换算。
写端（cp.async 装载，源码 263 行）：
c++
// 存的时候: 列坐标经过 swizzle 换算
int permuted_a_k = swizzle_permuted_A_j<MMA_K>(load_smem_a_m, load_smem_a_k);
uint32_t load_smem_a_ptr = smem_a_base_ptr +
(smem_sel * s_a_stage_offset + load_smem_a_m * (BK) + permuted_a_k) * sizeof(half);
CP_ASYNC_CG(load_smem_a_ptr, &A[load_gmem_a_addr], 16);读端（ldmatrix，297 行）：同一个 swizzle_permuted_A_j 换算——写进去用重排列号，读出来也用同一个重排列号，逻辑上数据没动（s_a[i][swizzle(j)] 两端一致），物理上 bank 分布被打散。
这引出 swizzle 的本质：它是"纯物理层"的优化——逻辑布局（tile 语义）不变，计算代码一行不用改，只有 smem 地址换算插入。对照 padding 的"物理塞废位 + 逻辑索引偏移"，swizzle 没有任何存储浪费（BK 保持 16 不膨胀），stage 数不用降。
对照 09/10 篇：三把刀的完整对比
到本篇，bank conflict 的武器库集齐了，横向对比（这是贯穿 09→10→18 的总结账）：
| padding | CuTe Swizzle（10 篇） | 手写 swizzle（本篇） | |
|---|---|---|---|
| 原理 | 行尾塞废位错开行距 | 位异或交错 | 位异或交错 |
| 存储 | 有浪费（PAD 列） | 零浪费 | 零浪费 |
| 对齐 | 破坏 16B 对齐（04 篇的坑） | 保持 | 保持 |
| 表达 | s[BK][BN+PAD] 直白 | composition(Swizzle, layout) 声明式 | 一个 inline 函数 |
| 灵活性 | 固定 | 任意 shape/stride 组合的代数推导 | 按访存形态定制 |
10 篇的 CuTe swizzle (c^r) 位交错和本篇的 ((j>>3)^(i>>2)) 是同一个东西的不同记法——CuTe 把它抽象成可组合的代数对象（Swizzle<B,M,S>：B=位数、M=异或掩码起点、S=位移），手写版直接针对 m16n8k16 + 128bit 访存形态定制。10 篇当时"不深入"的部分，现在从手写侧完全打通了——这也正是 19 篇（CuTe 版 HGEMM）要兑现的闭环：用 CuTe 把 17/18 篇的所有手排布局重新表达一遍。
回到数据：为什么小形状赚、大形状亏
现在能解释开头那个反直觉了。冲突消除的收益 = 冲突造成的延迟占总耗时的比例：
- 小形状（2048³）：K/BK = 128 轮迭代，每轮 ldmatrix/cp.async 的 bank 排队是循环体的主要开销之一——消掉它，+6.6%
- 大形状（8192³+）：计算量是 3 次方涨（8192³ vs 2048³ = 64 倍），mma 指令流占满 SM，smem 延迟被计算密度掩盖（15 篇"瓶颈翻转"的又一实例）——冲突收益趋于 0，反而 swizzle 的 XOR 换算指令是每循环白白多几条整数运算，-1%
这验证了 12 篇立的铁律：优化手段必须打在瓶颈上。swizzle 打的是"访存冲突"，但大形状 HGEMM 的瓶颈根本不在那。工业库（CUTLASS）仍然默认带 swizzle——因为它们要服务所有形状，小形状的 +6.6% 也是 +，且 XOR 两条指令在编译器调度好的大 kernel 里经常能藏进流水线空隙（我们 -O2 的手写版藏不住，cuBLAS 级的 SASS 调度可以）。
另一个实验证据：x4 版（A&B 都 swizzle + ldmatrix.x4，README 里 4090 的最优配置之一）在 4096³ 打到 149.7 vs swizzle 版 148.5——组合优化在大中形状又开始赚，说明 x4 把读端指令数砍半的收益盖过了 XOR 开销。swizzle 不是孤立开关，和 ldmatrix 宽度、tile 形状联动调优才有完整收益。
完整代码
hgemm_mma_stage_tn_swizzle.cu
hgemm_mma_stage_tn_swizzle.cu
c++
#include <algorithm>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <cuda_fp8.h>
#include <cuda_runtime.h>
#include <float.h>
#include <mma.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <vector>
using namespace nvcuda;
#define WARP_SIZE 32
#define DEVICE_INLINE __device__ inline
#define HOST_DEVICE_INLINE __device__ __host__ inline
#define INT4(value) (reinterpret_cast<int4 *>(&(value))[0])
#define FLOAT4(value) (reinterpret_cast<float4 *>(&(value))[0])
#define HALF2(value) (reinterpret_cast<half2 *>(&(value))[0])
#define BFLOAT2(value) (reinterpret_cast<__nv_bfloat162 *>(&(value))[0])
#define LDST32BITS(value) (reinterpret_cast<half2 *>(&(value))[0])
#define LDST64BITS(value) (reinterpret_cast<float2 *>(&(value))[0])
#define LDST128BITS(value) (reinterpret_cast<float4 *>(&(value))[0])
// gmem -> smem
#define CP_ASYNC_COMMIT_GROUP() asm volatile("cp.async.commit_group;\n" ::)
#define CP_ASYNC_WAIT_ALL() asm volatile("cp.async.wait_all;\n" ::)
#define CP_ASYNC_WAIT_GROUP(n) \
asm volatile("cp.async.wait_group %0;\n" ::"n"(n))
// ca(cache all, L1 + L2): support 4, 8, 16 bytes, cg(cache global, L2): only
// support 16 bytes.
#define CP_ASYNC_CA(dst, src, bytes) \
asm volatile( \
"cp.async.ca.shared.global.L2::128B [%0], [%1], %2;\n" ::"r"(dst), \
"l"(src), "n"(bytes))
#define CP_ASYNC_CG(dst, src, bytes) \
asm volatile( \
"cp.async.cg.shared.global.L2::128B [%0], [%1], %2;\n" ::"r"(dst), \
"l"(src), "n"(bytes))
// smem -> gmem: requires sm_90 or higher.
#define CP_ASYNC_BULK_COMMIT_GROUP() \
asm volatile("cp.async.bulk.commit_group;\n" ::)
#define CP_ASYNC_BULK_WAIT_ALL() asm volatile("cp.async.bulk.wait_all;\n" ::)
#define CP_ASYNC_BULK_WAIT_GROUP(n) \
asm volatile("cp.async.bulk.wait_group %0;\n" ::"n"(n))
#define CP_ASYNC_BULK(dst, src, bytes) \
asm volatile( \
"cp.async.bulk.global.shared::cta.bulk_group.L2::128B [%0], [%1], " \
"%2;\n" ::"r"(dst), \
"l"(src), "n"(bytes))
// ldmatrix
#define LDMATRIX_X1(R, addr) \
asm volatile("ldmatrix.sync.aligned.x1.m8n8.shared.b16 {%0}, [%1];\n" \
: "=r"(R) \
: "r"(addr))
#define LDMATRIX_X2(R0, R1, addr) \
asm volatile("ldmatrix.sync.aligned.x2.m8n8.shared.b16 {%0, %1}, [%2];\n" \
: "=r"(R0), "=r"(R1) \
: "r"(addr))
#define LDMATRIX_X4(R0, R1, R2, R3, addr) \
asm volatile( \
"ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n" \
: "=r"(R0), "=r"(R1), "=r"(R2), "=r"(R3) \
: "r"(addr))
#define LDMATRIX_X1_T(R, addr) \
asm volatile("ldmatrix.sync.aligned.x1.trans.m8n8.shared.b16 {%0}, [%1];\n" \
: "=r"(R) \
: "r"(addr))
#define LDMATRIX_X2_T(R0, R1, addr) \
asm volatile( \
"ldmatrix.sync.aligned.x2.trans.m8n8.shared.b16 {%0, %1}, [%2];\n" \
: "=r"(R0), "=r"(R1) \
: "r"(addr))
#define LDMATRIX_X4_T(R0, R1, R2, R3, addr) \
asm volatile( \
"ldmatrix.sync.aligned.x4.trans.m8n8.shared.b16 {%0, %1, %2, %3}, " \
"[%4];\n" \
: "=r"(R0), "=r"(R1), "=r"(R2), "=r"(R3) \
: "r"(addr))
// stmatrix: requires sm_90 or higher.
#define STMATRIX_X1(addr, R) \
asm volatile( \
"stmatrix.sync.aligned.x1.m8n8.shared.b16 [%0], {%1};\n" ::"r"(addr), \
"r"(R))
#define STMATRIX_X2(addr, R0, R1) \
asm volatile( \
"stmatrix.sync.aligned.x2.m8n8.shared.b16 [%0], {%1, %2};\n" ::"r"( \
addr), \
"r"(R0), "r"(R1))
#define STMATRIX_X4(addr, R0, R1, R2, R3) \
asm volatile( \
"stmatrix.sync.aligned.x4.m8n8.shared.b16 [%0], {%1, %2, %3, %4};\n" :: \
"r"(addr), \
"r"(R0), "r"(R1), "r"(R2), "r"(R3))
#define STMATRIX_X1_T(addr, R) \
asm volatile( \
"stmatrix.sync.aligned.x1.trans.m8n8.shared.b16 [%0], {%1};\n" ::"r"( \
addr), \
"r"(R))
#define STMATRIX_X2_T(addr, R0, R1) \
asm volatile( \
"stmatrix.sync.aligned.x2.trans.m8n8.shared.b16 [%0], {%1, %2};\n" :: \
"r"(addr), \
"r"(R0), "r"(R1))
#define STMATRIX_X4_T(addr, R0, R1, R2, R3) \
asm volatile( \
"stmatrix.sync.aligned.x4.trans.m8n8.shared.b16 [%0], {%1, %2, %3, " \
"%4};\n" ::"r"(addr), \
"r"(R0), "r"(R1), "r"(R2), "r"(R3))
// mma m16n8k16
#define HMMA16816(RD0, RD1, RA0, RA1, RA2, RA3, RB0, RB1, RC0, RC1) \
asm volatile( \
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3, " \
"%4, %5}, {%6, %7}, {%8, %9};\n" \
: "=r"(RD0), "=r"(RD1) \
: "r"(RA0), "r"(RA1), "r"(RA2), "r"(RA3), "r"(RB0), "r"(RB1), "r"(RC0), \
"r"(RC1))
HOST_DEVICE_INLINE
int div_ceil(int a, int b) { return (a % b != 0) ? (a / b + 1) : (a / b); }
// i: row index; j: col index.
// e.g kColStride = 16, kStep = 8 -> load 8 half as 128 bits memory issue.
template <const int kColStride = 16, const int kStep = 8>
static __device__ __forceinline__ int swizzle_permuted_j(int i, int j) {
// for col_stride > 16, we have to permute it using col major ZigZag order.
// e.g, A smem logical layout [Br,d]=[Br,64] -> store layout [4][Br][16].
static_assert(kColStride <= 16, "kColStride must <= 16");
// swizzle: ((int(j / kStep) ^ int(i / 4)) % int(kColStride / kStep)) * kStep;
static_assert(kStep == 4 || kStep == 8, "kStep must be 8 or 4.");
static_assert(kColStride % kStep == 0,
"kColStride must be multiple of kStep.");
if constexpr (kStep == 8) {
return (((j >> 3) ^ (i >> 2)) % (kColStride >> 3)) << 3;
} else {
static_assert(kStep == 4);
return (((j >> 2) ^ (i >> 2)) % (kColStride >> 2)) << 2;
}
}
// i: row index; j: col index
template <const int kMmaAtomK = 16>
static __device__ __forceinline__ int swizzle_permuted_A_j(int i, int j) {
// -------------------
// -col 0~16, step 8--
// -------------------
// | row 0 | (0, 8) |
// | row 1 | (0, 8) |
// | row 2 | (0, 8) |
// | row 3 | (0, 8) |
// -------------------
// | row 4 | (8, 0) |
// | row 5 | (8, 0) |
// | row 6 | (8, 0) |
// | row 7 | (8, 0) |
// -------------------
// | row 8 | (0, 8) |
// | row 9 | (0, 8) |
// | row 10 | (0, 8) |
// | row 11 | (0, 8) |
// -------------------
// | row 12 | (8, 0) |
// | row 13 | (8, 0) |
// | row 14 | (8, 0) |
// | row 15 | (8, 0) |
// -------------------
return swizzle_permuted_j<kMmaAtomK, 8>(i, j);
}
// i: row index; j: col index
template <const int kMmaAtomK = 16>
static __device__ __forceinline__ int swizzle_permuted_B_j(int i, int j) {
// -------------------
// -col 0~16, step 8--
// -------------------
// | row 0 | (0, 8) |
// | row 1 | (0, 8) |
// | row 2 | (0, 8) |
// | row 3 | (0, 8) |
// -------------------
// | row 4 | (8, 0) |
// | row 5 | (8, 0) |
// | row 6 | (8, 0) |
// | row 7 | (8, 0) |
// -------------------
// | row 8 | (0, 8) |
// | row 9 | (0, 8) |
// | row 10 | (0, 8) |
// | row 11 | (0, 8) |
// -------------------
// | row 12 | (8, 0) |
// | row 13 | (8, 0) |
// | row 14 | (8, 0) |
// | row 15 | (8, 0) |
// -------------------
return swizzle_permuted_j<kMmaAtomK, 8>(i, j);
}
// NN: A/B/C All row major
// TN: A row major MxK, B col major NxK, C row major MxN
// 128x128, mma2x4, warp4x4(64,32,16), stages, block swizzle, dsmem
template <const int MMA_M = 16, const int MMA_N = 8, const int MMA_K = 16,
const int MMA_TILE_M = 2, const int MMA_TILE_N = 4,
const int WARP_TILE_M = 4, const int WARP_TILE_N = 4,
const int A_PAD = 0, const int B_PAD = 0, const int K_STAGE = 2,
const bool BLOCK_SWIZZLE = false>
__global__ void __launch_bounds__(256)
hgemm_mma_m16n8k16_mma2x4_warp4x4_stages_dsmem_tn_swizzle_kernel(
half *A, half *B, half *C, int M, int N, int K) {
// BLOCK_SWIZZLE 0/1 control use block swizzle or not.
const int bx = ((int)BLOCK_SWIZZLE) * blockIdx.z * gridDim.x + blockIdx.x;
const int by = blockIdx.y;
const int NUM_K_TILES = div_ceil(K, MMA_K);
constexpr int BM = MMA_M * MMA_TILE_M * WARP_TILE_M; // 16*2*4=128
constexpr int BN = MMA_N * MMA_TILE_N * WARP_TILE_N; // 8*4*4=128
constexpr int BK = MMA_K; // 16
extern __shared__ half smem[];
half *s_a = smem;
half *s_b = smem + K_STAGE * BM * (BK + A_PAD);
constexpr int s_a_stage_offset = BM * (BK + A_PAD); // BMxBK 128*16
constexpr int s_b_stage_offset = BN * (BK + B_PAD); // BNxBK 128*16
const int tid = threadIdx.y * blockDim.x + threadIdx.x; // within block
const int warp_id = tid / WARP_SIZE; // 0~7 warp_id within block
const int lane_id = tid % WARP_SIZE; // 0~31
const int warp_m = warp_id % 2; // 0,1
const int warp_n = warp_id / 2; // 0,1,2,3
int load_smem_a_m = tid / 2; // row 0~127
int load_smem_a_k = (tid % 2 == 0) ? 0 : 8; // col 0,8
int load_smem_b_n = tid / 2; // row 0~127
int load_smem_b_k = (tid % 2 == 0) ? 0 : 8; // col 0,8
int load_gmem_a_m = by * BM + load_smem_a_m; // global row of c
int load_gmem_b_n = bx * BN + load_smem_b_n; // global col of c
if (load_gmem_a_m >= M || load_gmem_b_n >= N)
return;
uint32_t RA[WARP_TILE_M][4];
uint32_t RB[WARP_TILE_N][2];
uint32_t RC[WARP_TILE_M][WARP_TILE_N][2];
#pragma unroll
for (int i = 0; i < WARP_TILE_M; ++i) {
#pragma unroll
for (int j = 0; j < WARP_TILE_N; ++j) {
RC[i][j][0] = 0;
RC[i][j][1] = 0;
}
}
uint32_t smem_a_base_ptr = __cvta_generic_to_shared(s_a);
uint32_t smem_b_base_ptr = __cvta_generic_to_shared(s_b);
#pragma unroll
for (int k = 0; k < (K_STAGE - 1); ++k) { // 0, 1
int load_gmem_a_k = k * BK + load_smem_a_k; // global col of a
int load_gmem_a_addr = load_gmem_a_m * K + load_gmem_a_k;
int load_gmem_b_k = k * BK + load_smem_b_k; // global col of b
int load_gmem_b_addr = load_gmem_b_n * K + load_gmem_b_k;
uint32_t load_smem_a_ptr =
(smem_a_base_ptr +
(k * s_a_stage_offset + load_smem_a_m * (BK + A_PAD) +
swizzle_permuted_A_j<MMA_K>(load_smem_a_m, load_smem_a_k)) *
sizeof(half));
CP_ASYNC_CG(load_smem_a_ptr, &A[load_gmem_a_addr], 16);
uint32_t load_smem_b_ptr =
(smem_b_base_ptr +
(k * s_b_stage_offset + load_smem_b_n * (BK + B_PAD) +
swizzle_permuted_B_j<MMA_K>(load_smem_b_n, load_smem_b_k)) *
sizeof(half));
CP_ASYNC_CG(load_smem_b_ptr, &B[load_gmem_b_addr], 16);
CP_ASYNC_COMMIT_GROUP();
}
CP_ASYNC_WAIT_GROUP(K_STAGE - 2); // s2->0, s3->1, s4->2
__syncthreads();
#pragma unroll
for (int k = (K_STAGE - 1); k < NUM_K_TILES; ++k) {
// gmem -> smem
// s2/4 can use bitwise ops but s3 can not, so, we use mod
// ops for all stages kernel. s2: (k + 1)&1, s4: (k + 1)&3
// s3: (k + 1) % 3
int smem_sel = (k + 1) % K_STAGE; // s3 k 2->0, k 3->1, k 4->2...
int smem_sel_next = k % K_STAGE; // s3 k 2->2, k 3->0, k 4->1...
int load_gmem_a_k = k * BK + load_smem_a_k; // global col of a
int load_gmem_a_addr = load_gmem_a_m * K + load_gmem_a_k;
int load_gmem_b_k = k * BK + load_smem_b_k; // global col of b
int load_gmem_b_addr = load_gmem_b_n * K + load_gmem_b_k;
uint32_t load_smem_a_ptr =
(smem_a_base_ptr +
(smem_sel_next * s_a_stage_offset + load_smem_a_m * (BK + A_PAD) +
swizzle_permuted_A_j<MMA_K>(load_smem_a_m, load_smem_a_k)) *
sizeof(half));
CP_ASYNC_CG(load_smem_a_ptr, &A[load_gmem_a_addr], 16);
uint32_t load_smem_b_ptr =
(smem_b_base_ptr +
(smem_sel_next * s_b_stage_offset + load_smem_b_n * (BK + B_PAD) +
swizzle_permuted_B_j<MMA_K>(load_smem_b_n, load_smem_b_k)) *
sizeof(half));
CP_ASYNC_CG(load_smem_b_ptr, &B[load_gmem_b_addr], 16);
CP_ASYNC_COMMIT_GROUP();
// smem -> reg
#pragma unroll
for (int i = 0; i < WARP_TILE_M; ++i) {
int warp_smem_a_m = warp_m * (MMA_M * WARP_TILE_M) + i * MMA_M;
int lane_smem_a_m = warp_smem_a_m + lane_id % 16; // 0~15
int lane_smem_a_k = (lane_id / 16) * 8; // 0,8
uint32_t lane_smem_a_ptr =
(smem_a_base_ptr +
(smem_sel * s_a_stage_offset + lane_smem_a_m * (BK + A_PAD) +
swizzle_permuted_A_j<MMA_K>(lane_smem_a_m, lane_smem_a_k)) *
sizeof(half));
LDMATRIX_X4(RA[i][0], RA[i][1], RA[i][2], RA[i][3], lane_smem_a_ptr);
}
#pragma unroll
for (int j = 0; j < WARP_TILE_N; ++j) {
int warp_smem_b_n = warp_n * (MMA_N * WARP_TILE_N) + j * MMA_N;
int lane_smem_b_n = warp_smem_b_n + lane_id % 8; // 0~7, MMA_N=8
int lane_smem_b_k = ((lane_id / 8) % 2) * 8; // 0,8
uint32_t lane_smem_b_ptr =
(smem_b_base_ptr +
(smem_sel * s_b_stage_offset + lane_smem_b_n * (BK + B_PAD) +
swizzle_permuted_B_j<MMA_K>(lane_smem_b_n, lane_smem_b_k)) *
sizeof(half));
LDMATRIX_X2(RB[j][0], RB[j][1], lane_smem_b_ptr);
}
// MMA compute
#pragma unroll
for (int i = 0; i < WARP_TILE_M; ++i) {
#pragma unroll
for (int j = 0; j < WARP_TILE_N; ++j) {
HMMA16816(RC[i][j][0], RC[i][j][1], RA[i][0], RA[i][1], RA[i][2],
RA[i][3], RB[j][0], RB[j][1], RC[i][j][0], RC[i][j][1]);
}
}
CP_ASYNC_WAIT_GROUP(K_STAGE - 2);
__syncthreads();
}
// make sure all memory issues ready.
if ((K_STAGE - 2) > 0) {
CP_ASYNC_WAIT_GROUP(0);
__syncthreads();
}
// processing last (K_STAGE-1) k iters.
{
#pragma unroll
for (int k = 0; k < (K_STAGE - 1); k++) {
int stage_sel = ((NUM_K_TILES - (K_STAGE - 1) + k) % K_STAGE);
// ldmatrix for s_a, ldmatrix.trans for s_b.
// smem -> reg
#pragma unroll
for (int i = 0; i < WARP_TILE_M; ++i) {
int warp_smem_a_m = warp_m * (MMA_M * WARP_TILE_M) + i * MMA_M;
int lane_smem_a_m = warp_smem_a_m + lane_id % 16; // 0~15
int lane_smem_a_k = (lane_id / 16) * 8; // 0,8
uint32_t lane_smem_a_ptr =
(smem_a_base_ptr +
(stage_sel * s_a_stage_offset + lane_smem_a_m * (BK + A_PAD) +
swizzle_permuted_A_j<MMA_K>(lane_smem_a_m, lane_smem_a_k)) *
sizeof(half));
LDMATRIX_X4(RA[i][0], RA[i][1], RA[i][2], RA[i][3], lane_smem_a_ptr);
}
#pragma unroll
for (int j = 0; j < WARP_TILE_N; ++j) {
int warp_smem_b_n = warp_n * (MMA_N * WARP_TILE_N) + j * MMA_N;
int lane_smem_b_n = warp_smem_b_n + lane_id % 8; // 0~7, MMA_N=8
int lane_smem_b_k = ((lane_id / 8) % 2) * 8; // 0,8
uint32_t lane_smem_b_ptr =
(smem_b_base_ptr +
(stage_sel * s_b_stage_offset + lane_smem_b_n * (BK + B_PAD) +
swizzle_permuted_B_j<MMA_K>(lane_smem_b_n, lane_smem_b_k)) *
sizeof(half));
LDMATRIX_X2(RB[j][0], RB[j][1], lane_smem_b_ptr);
}
// MMA compute
#pragma unroll
for (int i = 0; i < WARP_TILE_M; ++i) {
#pragma unroll
for (int j = 0; j < WARP_TILE_N; ++j) {
HMMA16816(RC[i][j][0], RC[i][j][1], RA[i][0], RA[i][1], RA[i][2],
RA[i][3], RB[j][0], RB[j][1], RC[i][j][0], RC[i][j][1]);
}
}
}
}
{
for (int i = 0; i < WARP_TILE_M; ++i) {
// How to use LDST128BITS here? __shfl_sync -> lane 0 -> store 8 half.
// thus, we only need 8 memory issues with 128 bits after shfl_sync.
#pragma unroll
for (int j = 0; j < WARP_TILE_N; ++j) {
// reuse RA[4][4] for collective store, reduce register pressure.
RA[0][0] = RC[i][j][0];
RA[1][0] = RC[i][j][1];
RA[0][1] = __shfl_sync((0xffffffff), RC[i][j][0], lane_id + 1);
RA[0][2] = __shfl_sync((0xffffffff), RC[i][j][0], lane_id + 2);
RA[0][3] = __shfl_sync((0xffffffff), RC[i][j][0], lane_id + 3);
RA[1][1] = __shfl_sync((0xffffffff), RC[i][j][1], lane_id + 1);
RA[1][2] = __shfl_sync((0xffffffff), RC[i][j][1], lane_id + 2);
RA[1][3] = __shfl_sync((0xffffffff), RC[i][j][1], lane_id + 3);
if (lane_id % 4 == 0) {
int store_warp_smem_c_m = warp_m * (MMA_M * WARP_TILE_M) + i * MMA_M;
int store_lane_gmem_c_m = by * BM + store_warp_smem_c_m + lane_id / 4;
int store_warp_smem_c_n = warp_n * (MMA_N * WARP_TILE_N) + j * MMA_N;
int store_lane_gmem_c_n = bx * BN + store_warp_smem_c_n;
int store_gmem_c_addr_0 =
store_lane_gmem_c_m * N + store_lane_gmem_c_n;
int store_gmem_c_addr_1 =
(store_lane_gmem_c_m + 8) * N + store_lane_gmem_c_n;
LDST128BITS(C[store_gmem_c_addr_0]) = LDST128BITS(RA[0][0]);
LDST128BITS(C[store_gmem_c_addr_1]) = LDST128BITS(RA[1][0]);
}
} // end for WARP_TILE_N
} // end for WARP_TILE_M
}
}
// build cpp binary
#ifndef NO_MMA_HGEMM_BIN
#include "utils.h"
// 128x128, mma2x4, warp4x4(64,32,16), stages, block&smem swizzle, dsmem, TN
#define LAUNCH_16816_STAGE_SWIZZLE_MMA2x4_WARP4x4_DSMEM_TN_X2_KERNEL(stages, \
stride) \
{ \
const int smem_max_size = ((stages) * BM * (BK + A_PAD) * sizeof(half) + \
(stages) * BN * (BK + B_PAD) * sizeof(half)); \
cudaFuncSetAttribute( \
hgemm_mma_m16n8k16_mma2x4_warp4x4_stages_dsmem_tn_swizzle_kernel< \
MMA_M, MMA_N, MMA_K, MMA_TILE_M, MMA_TILE_N, WARP_TILE_M, \
WARP_TILE_N, A_PAD, B_PAD, (stages), true>, \
cudaFuncAttributeMaxDynamicSharedMemorySize, 98304); \
const int N_SWIZZLE = (N + (stride) - 1) / (stride); \
dim3 block(NUM_THREADS); \
dim3 grid((div_ceil(N, BN) + N_SWIZZLE - 1) / N_SWIZZLE, div_ceil(M, BM), \
N_SWIZZLE); \
hgemm_mma_m16n8k16_mma2x4_warp4x4_stages_dsmem_tn_swizzle_kernel< \
MMA_M, MMA_N, MMA_K, MMA_TILE_M, MMA_TILE_N, WARP_TILE_M, WARP_TILE_N, \
A_PAD, B_PAD, (stages), true> \
<<<grid, block, smem_max_size>>>(a, b, c, M, N, K); \
}
template <const int K_STAGE = 2, const int BLOCK_SWIZZLE_STRIDE = 2048>
void lanunch_hgemm_mma_m16n8k16_swizzle_tn(half *a, half *b, half *c, int M,
int N, int K) {
constexpr int MMA_M = 16;
constexpr int MMA_N = 8;
constexpr int MMA_K = 16;
constexpr int MMA_TILE_M = 2;
constexpr int MMA_TILE_N = 4;
constexpr int WARP_TILE_M = 4;
constexpr int WARP_TILE_N = 4;
constexpr int A_PAD = 0; // apply smem swizzle
constexpr int B_PAD = 0; // apply smem swizzle
constexpr int NUM_THREADS =
(MMA_TILE_M * MMA_TILE_N * WARP_SIZE); // 2 * 4 * 32 = 256
constexpr int BM = MMA_M * MMA_TILE_M * WARP_TILE_M;
constexpr int BN = MMA_N * MMA_TILE_N * WARP_TILE_N;
constexpr int BK = MMA_K;
// s2: 2*128*(32)*2=16KB, 2*32*(128+16)*2=18KB, ~35KB
// s3: 3*128*(32)*2=24KB, 3*32*(128+16)*2=27KB, ~51KB
// s4: 4*128*(32)*2=32KB, 4*32*(128+16)*2=36KB, ~68KB
// s5: 5*128*(32)*2=40KB, 5*32*(128+16)*2=45KB, ~85KB
LAUNCH_16816_STAGE_SWIZZLE_MMA2x4_WARP4x4_DSMEM_TN_X2_KERNEL(
K_STAGE, BLOCK_SWIZZLE_STRIDE);
}
#ifdef HGEMM_MMA_DEBUG
#include <iostream>
#endif
int main(int argc, char *argv[]) {
// ncu --metrics
// sm__sass_l1tex_data_bank_conflicts_pipe_lsu_mem_shared_op_ldsm
// ./hgemm_mma_stage_tn_swizzle.80.debug.bin 1024 1024 1024 0 1
#ifdef HGEMM_MMA_DEBUG
const int test_num = 1;
#else
const int test_num = 64;
#endif
int M_list[test_num];
int N_list[test_num];
int K_list[test_num];
for (int i = 0; i < test_num; i++) {
M_list[i] = (i + 1) * 256;
N_list[i] = (i + 1) * 256;
K_list[i] = (i + 1) * 256;
}
#ifdef HGEMM_MMA_DEBUG
if (argc > 1)
M_list[0] = std::stoi(argv[1]);
if (argc > 2)
N_list[0] = std::stoi(argv[2]);
if (argc > 3)
K_list[0] = std::stoi(argv[3]);
#endif
#ifdef HGEMM_MMA_DEBUG
int outer_repeat = 1, inner_repeat = 1, warmup = 1;
if (argc > 4)
warmup = std::stoi(argv[4]);
if (argc > 5)
inner_repeat = std::stoi(argv[5]);
#else
int outer_repeat = 10, inner_repeat = 1, warmup = 1;
#endif
printf(
"ALGO = MMA16816 HGEMM TN MMA=2x4 WARP=4x4 STAGES=2 BLOCK SWIZZLE=2048, "
"SMEM SWIZZLE, B LDMATRIX.X2\n");
#ifndef HGEMM_MMA_DEBUG
for (int j = 0; j < 5; j++) {
int M = M_list[j], N = N_list[j], K = K_list[j];
float max_error = gemm_error_check_tn<half>(
lanunch_hgemm_mma_m16n8k16_swizzle_tn<2, 2048>, M, N, K);
printf("M N K = %6d %6d %6d, ", M, N, K);
printf("Max Error = %f\n", max_error);
}
#endif
for (int j = 0; j < test_num; j++) {
int M = M_list[j], N = N_list[j], K = K_list[j];
double max_sec = 0.0;
double min_sec = DBL_MAX;
double total_sec = 0.0;
for (int k = 0; k < outer_repeat; k++) {
double this_sec =
perf_gemm<half>(lanunch_hgemm_mma_m16n8k16_swizzle_tn<2, 2048>, M, N,
K, inner_repeat, warmup);
max_sec = max(max_sec, this_sec);
min_sec = min(min_sec, this_sec);
total_sec += this_sec;
}
// 1 TFLOPS = 10^12 FLOPS
// ref: https://imgtec.eetrend.com/blog/2021/100062210.html.
double avg_sec = total_sec / outer_repeat;
double avg_Tflops = ((double)M) * N * K * 2 * 1e-12 / avg_sec;
printf("M N K = %6d %6d %6d, W = %1d, R = %2d ", M, N, K, warmup,
inner_repeat);
printf("Time = %12.8lf %12.8lf %12.8lf s, ", min_sec, avg_sec, max_sec);
printf("AVG Performance = %10.4lf Tflops\n", avg_Tflops);
}
return 0;
}
#else
#include <torch/extension.h>
#include <torch/types.h>
#define STRINGFY(str) #str
#define TORCH_BINDING_COMMON_EXTENSION(func) \
m.def(STRINGFY(func), &func, STRINGFY(func));
#define CHECK_TORCH_TENSOR_DTYPE(T, th_type) \
if (((T).options().dtype() != (th_type))) { \
std::cout << "Tensor Info:" << (T).options() << std::endl; \
throw std::runtime_error("values must be " #th_type); \
}
#define CHECK_TORCH_TENSOR_SHAPE(T, S0, S1) \
if (((T).size(0) != (S0)) || ((T).size(1) != (S1))) { \
throw std::runtime_error("Tensor size mismatch!"); \
}
// 128x128, mma2x4, warp4x4(64,32,16), stages, block&smem swizzle, dsmem, TN
#define LAUNCH_16816_STAGE_SWIZZLE_MMA2x4_WARP4x4_DSMEM_TN_X2_KERNEL(stages, \
stride) \
{ \
const int smem_max_size = ((stages) * BM * (BK + A_PAD) * sizeof(half) + \
(stages) * BN * (BK + B_PAD) * sizeof(half)); \
cudaFuncSetAttribute( \
hgemm_mma_m16n8k16_mma2x4_warp4x4_stages_dsmem_tn_swizzle_kernel< \
MMA_M, MMA_N, MMA_K, MMA_TILE_M, MMA_TILE_N, WARP_TILE_M, \
WARP_TILE_N, A_PAD, B_PAD, (stages), true>, \
cudaFuncAttributeMaxDynamicSharedMemorySize, 98304); \
const int N_SWIZZLE = (N + (stride) - 1) / (stride); \
dim3 block(NUM_THREADS); \
dim3 grid((div_ceil(N, BN) + N_SWIZZLE - 1) / N_SWIZZLE, div_ceil(M, BM), \
N_SWIZZLE); \
hgemm_mma_m16n8k16_mma2x4_warp4x4_stages_dsmem_tn_swizzle_kernel< \
MMA_M, MMA_N, MMA_K, MMA_TILE_M, MMA_TILE_N, WARP_TILE_M, WARP_TILE_N, \
A_PAD, B_PAD, (stages), true><<<grid, block, smem_max_size>>>( \
reinterpret_cast<half *>(a.data_ptr()), \
reinterpret_cast<half *>(b.data_ptr()), \
reinterpret_cast<half *>(c.data_ptr()), M, N, K); \
}
#define LAUNCH_16816_STAGE_NO_SWIZZLE_MMA2x4_WARP4x4_DSMEM_TN_X2_KERNEL( \
stages) \
{ \
const int smem_max_size = ((stages) * BM * (BK + A_PAD) * sizeof(half) + \
(stages) * BN * (BK + B_PAD) * sizeof(half)); \
cudaFuncSetAttribute( \
hgemm_mma_m16n8k16_mma2x4_warp4x4_stages_dsmem_tn_swizzle_kernel< \
MMA_M, MMA_N, MMA_K, MMA_TILE_M, MMA_TILE_N, WARP_TILE_M, \
WARP_TILE_N, A_PAD, B_PAD, (stages), false>, \
cudaFuncAttributeMaxDynamicSharedMemorySize, 98304); \
dim3 block(NUM_THREADS); \
dim3 grid(div_ceil(N, BN), div_ceil(M, BM)); \
hgemm_mma_m16n8k16_mma2x4_warp4x4_stages_dsmem_tn_swizzle_kernel< \
MMA_M, MMA_N, MMA_K, MMA_TILE_M, MMA_TILE_N, WARP_TILE_M, WARP_TILE_N, \
A_PAD, B_PAD, (stages), false><<<grid, block, smem_max_size>>>( \
reinterpret_cast<half *>(a.data_ptr()), \
reinterpret_cast<half *>(b.data_ptr()), \
reinterpret_cast<half *>(c.data_ptr()), M, N, K); \
}
// 128x128, mma2x4, warp4x4(64,32,16), stages, block&smem swizzle, dsmem
void hgemm_mma_m16n8k16_mma2x4_warp4x4_stages_dsmem_tn_swizzle(
torch::Tensor a, torch::Tensor b, torch::Tensor c, int stages, bool swizzle,
int swizzle_stride) {
CHECK_TORCH_TENSOR_DTYPE(a, torch::kHalf)
CHECK_TORCH_TENSOR_DTYPE(b, torch::kHalf)
CHECK_TORCH_TENSOR_DTYPE(c, torch::kHalf)
const int M = a.size(0);
const int K = a.size(1);
const int N = b.size(1);
CHECK_TORCH_TENSOR_SHAPE(a, M, K)
CHECK_TORCH_TENSOR_SHAPE(b, K, N)
CHECK_TORCH_TENSOR_SHAPE(c, M, N)
constexpr int MMA_M = 16;
constexpr int MMA_N = 8;
constexpr int MMA_K = 16;
constexpr int MMA_TILE_M = 2;
constexpr int MMA_TILE_N = 4;
constexpr int WARP_TILE_M = 4;
constexpr int WARP_TILE_N = 4;
constexpr int A_PAD = 0; // 0,8,16
constexpr int B_PAD = 0; // 0,8,16
constexpr int NUM_THREADS =
(MMA_TILE_M * MMA_TILE_N * WARP_SIZE); // 2 * 4 * 32 = 256
constexpr int BM = MMA_M * MMA_TILE_M * WARP_TILE_M;
constexpr int BN = MMA_N * MMA_TILE_N * WARP_TILE_N;
constexpr int BK = MMA_K;
if (swizzle) {
// assert(swizzle_stride % 256 == 0);
switch (stages) {
case 2:
LAUNCH_16816_STAGE_SWIZZLE_MMA2x4_WARP4x4_DSMEM_TN_X2_KERNEL(
2, swizzle_stride);
break;
case 3:
LAUNCH_16816_STAGE_SWIZZLE_MMA2x4_WARP4x4_DSMEM_TN_X2_KERNEL(
3, swizzle_stride);
break;
case 4:
LAUNCH_16816_STAGE_SWIZZLE_MMA2x4_WARP4x4_DSMEM_TN_X2_KERNEL(
4, swizzle_stride);
break;
case 5:
LAUNCH_16816_STAGE_SWIZZLE_MMA2x4_WARP4x4_DSMEM_TN_X2_KERNEL(
5, swizzle_stride);
break;
default:
LAUNCH_16816_STAGE_SWIZZLE_MMA2x4_WARP4x4_DSMEM_TN_X2_KERNEL(
2, swizzle_stride);
break;
}
} else {
switch (stages) {
case 2:
LAUNCH_16816_STAGE_NO_SWIZZLE_MMA2x4_WARP4x4_DSMEM_TN_X2_KERNEL(2);
break;
case 3:
LAUNCH_16816_STAGE_NO_SWIZZLE_MMA2x4_WARP4x4_DSMEM_TN_X2_KERNEL(3);
break;
case 4:
LAUNCH_16816_STAGE_NO_SWIZZLE_MMA2x4_WARP4x4_DSMEM_TN_X2_KERNEL(4);
break;
case 5:
LAUNCH_16816_STAGE_NO_SWIZZLE_MMA2x4_WARP4x4_DSMEM_TN_X2_KERNEL(5);
break;
default:
LAUNCH_16816_STAGE_NO_SWIZZLE_MMA2x4_WARP4x4_DSMEM_TN_X2_KERNEL(2);
break;
}
}
}
#endif本篇小结
- 128bit 访存的 bank 冲突：16B 横跨 4 bank，
s_a[128][16]的行距 32B 造成 4 行一循环的 bank 重复，8 路冲突——大 tile + 128bit 访存下 padding 的浪费付不起（smem 预算挤掉流水级数） - swizzle 核心：
块号 XOR (行号/组距)，一行位运算让相邻行组在 16B 块粒度互相错开——零存储浪费、零对齐破坏、参数（kStep=8/kColStride=16/组距4）全部由访存形态决定 - 侵入式但局部：写端 cp.async 和读端 ldmatrix 插同一个换算函数，逻辑布局不变——"纯物理层优化"
- 与 10 篇打通：CuTe 的
Swizzle<B,M,S>和手写((j>>3)^(i>>2))同源异记——10 篇欠的"不深入"在这里还清，19 篇 CuTe 版是完整闭环 - 瓶颈原则第四次出场：小形状（循环体访存主导）+6.6%，大形状（计算主导，冲突被掩盖）-1%——swizzle 是对的工具，但要打在对的形状上