hgemm_naive_kernel详解

hgemm_naive_kernel详解

仓库链接：xlite-dev/HGEMM at eee72be829545bd6bd115a4b252b5068c9f61597

代码链接：HGEMM/kernels/hgemm/naive/hgemm.cu at eee72be829545bd6bd115a4b252b5068c9f61597 · xlite-dev/HGEMM

CUDA HGEMM Kernel 技术报告

Half-Precision General Matrix Multiplication (HGEMM) 优化实现详解

本报告详细分析了从朴素实现到高度优化版本的 HGEMM CUDA kernel，涵盖 Shared Memory 优化、Bank Conflict 分析、向量化访存、Double Buffering 等核心技术。

---

📑 目录

1. 概述
2. Kernel 版本总览
3. 朴素实现分析
4. Block Tiling 优化
5. Thread Tiling 优化
6. Bank Conflict 详解
7. Bank Conflict Free 优化
8. Double Buffering 优化
9. 性能对比与分析
10. 总结与最佳实践

---

1. 概述

1.1 问题定义

HGEMM (Half-precision GEMM) 计算矩阵乘法：C = A × B

• A: M × N 矩阵 (FP16)
• B: K × N 矩阵 (FP16)
• C: M×N 矩阵 (FP16)

1.2 核心优化技术

优化技术	目的	收益
Block Tiling	数据局部性	减少全局内存访问
Thread Tiling	计算密度	提高 FMA/Load 比率
Vectorization	访存带宽	减少指令数量
Bank Conflict Free	共享内存效率	消除访存冲突
Double Buffering	延迟隐藏	重叠计算与访存

1.3 硬件背景

CUDA Shared Memory Bank 结构：
├── 32 个 Bank
├── 每个 Bank 宽度: 4 bytes
├── 一个 Bank Cycle: 32 banks × 4 bytes = 128 bytes = 64 half
└── Bank 计算: Bank = (element_index × 2 / 4) % 32 = (element_index / 2) % 32

---

2. Kernel 版本总览

┌─────────────────────────────────────────────────────────────────────────┐
│                        HGEMM Kernel 演进路径                             │
├─────────────────────────────────────────────────────────────────────────┤
│                                                                         │
│  hgemm_naive_f16_kernel                                                 │
│      │                                                                  │
│      │ + Shared Memory + K Tiling                                       │
│      ▼                                                                  │
│  hgemm_sliced_k_f16_kernel                                              │
│      │                                                                  │
│      │ + Thread Tile 8×8 + Vectorization (f16x4)                        │
│      ▼                                                                  │
│  hgemm_t_8x8_sliced_k_f16x4_kernel                                      │
│      │                                                                  │
│      │ + LDST64BITS Pack                                                │
│      ▼                                                                  │
│  hgemm_t_8x8_sliced_k_f16x4_pack_kernel                                 │
│      │                                                                  │
│      │ + Transpose s_a (Bank Conflict Reduction)                        │
│      ▼                                                                  │
│  hgemm_t_8x8_sliced_k_f16x4_bcf_kernel                                  │
│      │                                                                  │
│      │ + Padding (OFFSET) + f16x8                                       │
│      ▼                                                                  │
│  hgemm_t_8x8_sliced_k_f16x8_pack_bcf_kernel                             │
│      │                                                                  │
│      │ + Double Buffering                                               │
│      ▼                                                                  │
│  hgemm_t_8x8_sliced_k_f16x8_pack_bcf_dbuf_kernel                        │
│                                                                         │
└─────────────────────────────────────────────────────────────────────────┘

2.1 各版本配置参数

Kernel	BM	BN	BK	TM	TN	Threads	Elements/Thread
naive	32	32	-	1	1	1024	1
sliced_k	32	32	32	1	1	1024	1
t_8x8_f16x4	128	128	8	8	8	256	64
bcf/dbuf	128	128	8	8	8	256	64

---

3. 朴素实现分析

3.1 代码实现

__global__ void hgemm_naive_f16_kernel(half* a, half* b, half* c, int M, int N, int K) {
  int n = blockIdx.x * blockDim.x + threadIdx.x;
  int m = blockIdx.y * blockDim.y + threadIdx.y;

  if (m < M && n < N) {
    half psum = 0.0;
    #pragma unroll
    for (int k = 0; k < K; k++) {
      psum += a[m * K + k] * b[k * N + n];  // 每次迭代 2 次全局内存访问
    }
    c[m * N + n] = psum;
  }
}

3.2 执行配置

dim3 block(32, 32);  // 1024 threads/block
dim3 grid((N + 31) / 32, (M + 31) / 32);

3.3 问题分析

🔴 问题1：无数据复用

对于计算 C[m][n]，需要读取：
- A[m][0:K] 的整行 (K 个元素)
- B[0:K][n] 的整列 (K 个元素)

整个 C 矩阵的计算：
- A 的每一行被读取 N 次 → 总共 M × N × K 次访问
- B 的每一列被读取 M 次 → 总共 M × N × K 次访问

数据复用率 = 0（每个数据仅服务于一次计算）

🔴 问题2：极低的计算密度

每次循环迭代：
├── 2 次 Global Memory Load (a[m*K+k], b[k*N+n])
├── 1 次 FMA 计算 (psum += a * b)
└── 计算密度 = 1 FMA / 2 Load = 0.5 FMA/Load

理想计算密度 > 10 FMA/Load

🔴 问题3：内存访问模式不佳

矩阵 B 的访问模式：
Thread 0:  b[k*N + 0]   ← 连续 4 bytes
Thread 1:  b[k*N + 1]   ← 连续 4 bytes (但与 Thread 0 不连续)
Thread 31: b[k*N + 31]  ← 连续 4 bytes

问题：每个 warp 访问的是 B 矩阵同一行的连续元素
     → 可合并 (coalesced)，这点是 OK 的

矩阵 A 的访问模式：
同一个 warp 内，ty 相同的线程访问相同的 a[m*K+k]
     → 存在广播机会，但也意味着重复访问

🔴 问题4：无向量化

每次只加载 1 个 half (2 bytes)
理论上可以使用：
- half2:   4 bytes  (2x throughput)
- float4:  16 bytes (8x throughput)

3.4 性能特征

指标	值	评价
计算密度	0.5 FMA/Load	❌ 极低
数据复用	0	❌ 无
向量化	1 element	❌ 无
Shared Memory	未使用	❌ 未优化

---

4. Block Tiling 优化

4.1 核心思想

┌─────────────────────────────────────────────────────────────────────┐
│                      Block Tiling 原理图                            │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│     矩阵 A (M×K)           矩阵 B (K×N)           矩阵 C (M×N)      │
│   ┌────────────┐         ┌────────────┐         ┌────────────┐     │
│   │            │         │            │         │            │     │
│   │   ┌────┐   │   BK    │   ┌────────┤   BN    │   ┌────────┤     │
│   │   │ A' │◄──┼────────►│   │   B'   │◄───────►│   │   C'   │     │
│   │   │    │   │         │   │        │         │   │        │     │
│   │   └────┘   │   ▲     │   └────────┤         │   └────────┤     │
│   │     ▲      │   │     │            │         │            │     │
│   │     │ BM   │   │     │            │         │            │     │
│   └─────┼──────┘   │     └────────────┘         └────────────┘     │
│         │          │                                                │
│         └──────────┘                                                │
│              K 维度分块迭代                                          │
│                                                                     │
│   每个 Block 负责计算 C 中 BM×BN 大小的子块                          │
│   K 维度按 BK 大小分块迭代，每次迭代：                                │
│   1. 加载 A' (BM×BK) 到 Shared Memory                               │
│   2. 加载 B' (BK×BN) 到 Shared Memory                               │
│   3. 计算 C' += A' × B'                                             │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘

4.2 代码实现

template<const int BM=32, const int BN=32, const int BK=32>
__global__ void hgemm_sliced_k_f16_kernel(half* a, half* b, half* c, int M, int N, int K) {
  __shared__ half s_a[BM][BK], s_b[BK][BN];  // 共享内存缓存

  int bx = blockIdx.x, by = blockIdx.y;
  int tx = threadIdx.x, ty = threadIdx.y;
  int tid = ty * blockDim.x + tx;
  
  // 计算加载索引
  int load_smem_a_m = tid / 32;  // 行索引
  int load_smem_a_k = tid % 32;  // 列索引
  int load_smem_b_k = tid / 32;
  int load_smem_b_n = tid % 32;
  
  int load_gmem_a_m = by * BM + load_smem_a_m;
  int load_gmem_b_n = bx * BN + load_smem_b_n;
  
  half sum = 0;
  for (int bk = 0; bk < (K + BK - 1) / BK; ++bk) {
    // 协作加载数据到共享内存
    int load_gmem_a_k = bk * BK + load_smem_a_k;
    s_a[load_smem_a_m][load_smem_a_k] = a[load_gmem_a_m * K + load_gmem_a_k];
    
    int load_gmem_b_k = bk * BK + load_smem_b_k;
    s_b[load_smem_b_k][load_smem_b_n] = b[load_gmem_b_k * N + load_gmem_b_n];
    __syncthreads();
    
    // 从共享内存计算
    #pragma unroll
    for (int k = 0; k < BK; ++k) {
      sum += s_a[load_smem_a_m][k] * s_b[k][load_smem_b_n];
    }
    __syncthreads();
  }
  c[load_gmem_a_m * N + load_gmem_b_n] = sum;
}

4.3 数据复用分析

原始实现（无复用）：
├── 矩阵 A: 每个元素被读取 N 次
├── 矩阵 B: 每个元素被读取 M 次
└── 总全局内存访问: 2 × M × N × K

Block Tiling 实现：
├── 矩阵 A: 每个元素被读取 N/BN 次 (复用 BN 次)
├── 矩阵 B: 每个元素被读取 M/BM 次 (复用 BM 次)
└── 总全局内存访问: (M × N × K / BN) + (M × N × K / BM)

复用提升: BM 或 BN 倍 (取决于矩阵)

4.4 仍存在的问题

问题	说明
计算密度仍低	每线程仅计算 1 个输出元素
无向量化	单元素加载
同步开销	每次 K 迭代需要 2 次 __syncthreads()

---

5. Thread Tiling 优化

5.1 核心思想

┌─────────────────────────────────────────────────────────────────────┐
│                     Thread Tiling 原理图                            │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│   Block Tile (BM=128, BN=128)         Thread Tile (TM=8, TN=8)     │
│                                                                     │
│   ┌────────────────────────────┐      每个线程计算 8×8 = 64 个元素  │
│   │ ┌──┬──┬──┬──┬──┬──┬──┬──┐ │                                    │
│   │ ├──┼──┼──┼──┼──┼──┼──┼──┤ │      ┌──────────────┐              │
│   │ ├──┼──┼──┼──┼──┼──┼──┼──┤ │      │ r_c[8][8]    │              │
│   │ ├──┼──┼──┼──┼──┼──┼──┼──┤ │      │              │              │
│   │ ├──┼──┼──┼──┼──┼──┼──┼──┤ │  =   │  64 FP16     │              │
│   │ ├──┼──┼──┼──┼──┼──┼──┼──┤ │      │  寄存器      │              │
│   │ ├──┼──┼──┼──┼──┼──┼──┼──┤ │      │              │              │
│   │ ├──┼──┼──┼──┼──┼──┼──┼──┤ │      └──────────────┘              │
│   │ └──┴──┴──┴──┴──┴──┴──┴──┘ │                                    │
│   └────────────────────────────┘      Block: 16×16 = 256 线程      │
│                                       256 × 64 = 16384 = 128×128 ✓ │
│   线程布局: blockDim(16, 16)                                        │
│   tx: 0~15, ty: 0~15                                                │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘

5.2 计算密度提升

每次 K 迭代，每个线程：
├── 加载: r_comp_a[8] 从 s_a (8 个 half)
├── 加载: r_comp_b[8] 从 s_b (8 个 half)  
├── 计算: 8 × 8 = 64 次 FMA
└── 计算密度 = 64 FMA / 16 Load = 4 FMA/Load

对比 naive: 0.5 FMA/Load → 提升 8 倍！

5.3 代码实现

template<const int BM=128, const int BN=128, const int BK=8, const int TM=8, const int TN=8>
__global__ void hgemm_t_8x8_sliced_k_f16x4_kernel(half* a, half* b, half* c, int M, int N, int K) {
  int bx = blockIdx.x, by = blockIdx.y;
  int tx = threadIdx.x, ty = threadIdx.y;
  int tid = ty * blockDim.x + tx;  // 0~255
  
  __shared__ half s_a[BM][BK], s_b[BK][BN];  // 128×8 + 8×128 = 4KB
  
  // 加载索引计算
  // s_a[128][8]: 每行 8 元素，每线程加载 4 个，需要 2 线程/行，128 行需要 256 线程
  int load_smem_a_m = tid / 2;                    // 0~127
  int load_smem_a_k = (tid % 2 == 0) ? 0 : 4;     // 0 或 4
  
  // s_b[8][128]: 每行 128 元素，每线程加载 4 个，需要 32 线程/行，8 行需要 256 线程
  int load_smem_b_k = tid / 32;                   // 0~7
  int load_smem_b_n = (tid % 32) * 4;             // 0,4,8,...,124

  half r_c[TM][TN] = {0};  // 寄存器存储 8×8 结果
  
  for (int bk = 0; bk < (K + BK - 1) / BK; ++bk) {
    // 向量化加载到共享内存 (使用 half2)
    int load_gmem_a_k = bk * BK + load_smem_a_k;
    int load_gmem_a_addr = load_gmem_a_m * K + load_gmem_a_k;
    HALF2(s_a[load_smem_a_m][load_smem_a_k + 0]) = HALF2(a[load_gmem_a_addr + 0]);
    HALF2(s_a[load_smem_a_m][load_smem_a_k + 2]) = HALF2(a[load_gmem_a_addr + 2]);
    
    int load_gmem_b_k = bk * BK + load_smem_b_k;
    int load_gmem_b_addr = load_gmem_b_k * N + load_gmem_b_n;
    HALF2(s_b[load_smem_b_k][load_smem_b_n + 0]) = HALF2(b[load_gmem_b_addr + 0]);
    HALF2(s_b[load_smem_b_k][load_smem_b_n + 2]) = HALF2(b[load_gmem_b_addr + 2]);
    __syncthreads();
    
    #pragma unroll
    for (int k = 0; k < BK; k++) {
      #pragma unroll
      for (int m = 0; m < TM; m++) {
        #pragma unroll
        for (int n = 0; n < TN; n++) {
          int comp_smem_a_m = ty * TM + m;  // 线程 ty 负责的 M 范围
          int comp_smem_b_n = tx * TN + n;  // 线程 tx 负责的 N 范围
          r_c[m][n] += s_a[comp_smem_a_m][k] * s_b[k][comp_smem_b_n];
        }
      }
    }
    __syncthreads();
  }
  
  // 向量化写回结果
  #pragma unroll
  for (int m = 0; m < TM; ++m) {
    int store_gmem_c_m = by * BM + ty * TM + m;
    #pragma unroll
    for (int n = 0; n < TN; n += 2) {
      int store_gmem_c_n = bx * BN + tx * TN + n;
      HALF2(c[store_gmem_c_m * N + store_gmem_c_n]) = HALF2(r_c[m][n]);
    }
  }
}

5.4 线程到数据的映射

┌─────────────────────────────────────────────────────────────────────┐
│                  Thread (ty, tx) 到 C 子块的映射                     │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│   Block Tile C[128×128]                                             │
│                                                                     │
│        tx=0    tx=1    tx=2    ...    tx=15                         │
│       ┌──────┬──────┬──────┬───────┬──────┐                         │
│  ty=0 │ 8×8  │ 8×8  │ 8×8  │  ...  │ 8×8  │ ← 线程(0,0)~(0,15)     │
│       ├──────┼──────┼──────┼───────┼──────┤                         │
│  ty=1 │ 8×8  │ 8×8  │ 8×8  │  ...  │ 8×8  │ ← 线程(1,0)~(1,15)     │
│       ├──────┼──────┼──────┼───────┼──────┤                         │
│   ⋮   │  ⋮   │  ⋮   │  ⋮   │   ⋮   │  ⋮   │                         │
│       ├──────┼──────┼──────┼───────┼──────┤                         │
│ ty=15 │ 8×8  │ 8×8  │ 8×8  │  ...  │ 8×8  │ ← 线程(15,0)~(15,15)   │
│       └──────┴──────┴──────┴───────┴──────┘                         │
│                                                                     │
│   线程 (ty, tx) 计算 C 的子块:                                       │
│   ├── M 方向: [ty*TM, ty*TM+8) = [ty*8, ty*8+8)                     │
│   └── N 方向: [tx*TN, tx*TN+8) = [tx*8, tx*8+8)                     │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘

---

6. Bank Conflict 详解

6.1 Shared Memory Bank 结构

┌─────────────────────────────────────────────────────────────────────┐
│                    Shared Memory Bank 布局                          │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│   32 Banks, 每个 Bank 4 bytes, 总计 128 bytes = 64 half             │
│                                                                     │
│   Bank:    b0      b1      b2      b3     ...     b31               │
│          ┌─────┬─────┬─────┬─────┬─────┬─────┬─────┐               │
│   地址:  │ 0-3 │ 4-7 │8-11 │12-15│ ... │124-127│               │
│          │     │     │     │     │     │       │               │
│   half:  │h0,h1│h2,h3│h4,h5│h6,h7│ ... │h62,h63│               │
│          └─────┴─────┴─────┴─────┴─────┴───────┘               │
│                                                                     │
│   Bank 计算公式:                                                     │
│   Bank(element_index) = (element_index / 2) % 32                    │
│                                                                     │
│   例: half element[128]                                             │
│   element[0], element[1]   → Bank 0  (layer 0)                      │
│   element[2], element[3]   → Bank 1  (layer 0)                      │
│   ...                                                               │
│   element[62], element[63] → Bank 31 (layer 0)                      │
│   element[64], element[65] → Bank 0  (layer 1)  ← 回到 Bank 0!      │
│   element[66], element[67] → Bank 1  (layer 1)                      │
│   ...                                                               │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘

6.2 s_a[BM][BK] 的 Bank Conflict 分析

s_a[128][8] 布局 (行主序):

Row 0:   s_a[0][0~7]   → half[0~7]   → Bank 0~3 (layer 0)
Row 1:   s_a[1][0~7]   → half[8~15]  → Bank 4~7 (layer 0)
Row 2:   s_a[2][0~7]   → half[16~23] → Bank 8~11 (layer 0)
...
Row 7:   s_a[7][0~7]   → half[56~63] → Bank 28~31 (layer 0)
Row 8:   s_a[8][0~7]   → half[64~71] → Bank 0~3 (layer 1)  ← Bank 0~3 重复!
Row 9:   s_a[9][0~7]   → half[72~79] → Bank 4~7 (layer 1)
...
Row 15:  s_a[15][0~7]  → half[120~127] → Bank 28~31 (layer 1)
Row 16:  s_a[16][0~7]  → half[128~135] → Bank 0~3 (layer 2)  ← Bank 0~3 重复!
...

关键发现: 每 8 行，Bank 模式重复！
s_a[m][k] 与 s_a[m+8][k] 访问相同 Bank

6.3 读取 s_a 时的 Bank Conflict

计算阶段，线程访问模式:

comp_smem_a_m = ty * TM + m;  // TM=8
comp_smem_b_n = tx * TN + n;  // TN=8

对于 Warp 0 (tid 0~31, 即 ty=0~1, tx=0~15):

┌──────────────────────────────────────────────────────────────────┐
│   Warp 0 访问 s_a[comp_smem_a_m][k] 的 Bank 分布                  │
├──────────────────────────────────────────────────────────────────┤
│                                                                  │
│   ty=0 (tid 0~15):  comp_smem_a_m = 0×8 + m = m (0~7)           │
│   ty=1 (tid 16~31): comp_smem_a_m = 1×8 + m = 8+m (8~15)        │
│                                                                  │
│   当 m=0, k=0:                                                   │
│   ├── ty=0: s_a[0][0]  → Bank 0                                 │
│   └── ty=1: s_a[8][0]  → Bank 0  ← 同一 Bank!                   │
│                                                                  │
│   当 m=0, k=1:                                                   │
│   ├── ty=0: s_a[0][1]  → Bank 0                                 │
│   └── ty=1: s_a[8][1]  → Bank 0  ← 同一 Bank!                   │
│                                                                  │
│   结论: Warp 0 中 ty=0 和 ty=1 的线程总是访问相同 Bank            │
│         → 2-way Bank Conflict!                                   │
│                                                                  │
└──────────────────────────────────────────────────────────────────┘

6.4 完整 Warp 的 Bank Conflict 表

s_a[128][8] 读取 Bank Conflict (同一 k 值):

┌────────┬─────────────────┬───────────────────┬──────────────────┐
│ 线程   │ comp_smem_a_m   │ s_a[m][k] 地址    │ Bank             │
├────────┼─────────────────┼───────────────────┼──────────────────┤
│ ty=0   │ 0+m (0~7)       │ m*8+k             │ (m*8+k)/2 % 32   │
│ ty=1   │ 8+m (8~15)      │ (8+m)*8+k         │ ((8+m)*8+k)/2%32 │
│ ty=2   │ 16+m (16~23)    │ (16+m)*8+k        │ ...              │
│ ...    │ ...             │ ...               │ ...              │
│ ty=15  │ 120+m (120~127) │ (120+m)*8+k       │ ...              │
└────────┴─────────────────┴───────────────────┴──────────────────┘

当 m=0, k=0:
ty=0:  s_a[0][0]   = half[0]   → Bank 0
ty=1:  s_a[8][0]   = half[64]  → Bank 0   ← Conflict!
ty=2:  s_a[16][0]  = half[128] → Bank 0   ← Conflict!
...
ty=15: s_a[120][0] = half[960] → Bank 0   ← Conflict!

→ 16-way Bank Conflict! (ty=0~15 全部访问 Bank 0)

---

7. Bank Conflict Free 优化

7.1 转置 s_a: s_a[BK][BM]

原始布局 s_a[BM][BK] = s_a[128][8]:
├── 问题: 同一 k 值时，不同 ty 访问相同 Bank
└── 原因: 行间距 8 half = 4 Banks，每 8 行回到相同 Bank

转置布局 s_a[BK][BM] = s_a[8][128]:
├── 每行 128 个 half = 64 Banks = 2 个完整 Bank Cycle
└── 效果: 不同 m 值的元素分布在不同 Bank

新布局下的 Bank 分布:
s_a[k][0~1]   → Bank 0
s_a[k][2~3]   → Bank 1
s_a[k][4~5]   → Bank 2
...
s_a[k][62~63] → Bank 31
s_a[k][64~65] → Bank 0 (layer 1)
...

7.2 转置后的访问分析

转置后访问: s_a[k][ty * TM + m]

对于 Warp 0 (ty=0~1):

当 k=0, m=0:
├── ty=0: s_a[0][0]  → Bank 0
└── ty=1: s_a[0][8]  → Bank 4   ← 不同 Bank!

当 k=0, m=0~7 (ty=0):
├── s_a[0][0]  → Bank 0
├── s_a[0][1]  → Bank 0
├── s_a[0][2]  → Bank 1
├── s_a[0][3]  → Bank 1
├── s_a[0][4]  → Bank 2
├── s_a[0][5]  → Bank 2
├── s_a[0][6]  → Bank 3
└── s_a[0][7]  → Bank 3

→ 如果使用 LDST128BITS 一次读 8 个 half，只涉及 4 个连续 Bank

7.3 bcf_kernel 实现

template<const int BM=128, const int BN=128, const int BK=8, const int TM=8, const int TN=8>
__global__ void hgemm_t_8x8_sliced_k_f16x4_bcf_kernel(...) {
  __shared__ half s_a[BK][BM];  // 转置: [8][128] 而非 [128][8]
  __shared__ half s_b[BK][BN];  // 保持: [8][128]
  
  half r_load_a[4], r_load_b[4];  // 加载缓冲
  half r_comp_a[TM], r_comp_b[TN];  // 计算缓冲
  half r_c[TM][TN] = {0};

  // 加载时转置存储
  int load_a_smem_m = tid / 2;      // 原始行索引
  int load_a_smem_k = (tid & 1) << 2;  // 原始列索引 (0 或 4)
  
  for (int bk = 0; bk < (K + BK - 1) / BK; bk++) {
    // 从全局内存按行读取 A
    HALF2(r_load_a[0]) = HALF2(a[load_a_gmem_addr + 0]);
    HALF2(r_load_a[2]) = HALF2(a[load_a_gmem_addr + 2]);
    
    // 转置存入共享内存: s_a[k][m] 而非 s_a[m][k]
    s_a[load_a_smem_k + 0][load_a_smem_m] = r_load_a[0];
    s_a[load_a_smem_k + 1][load_a_smem_m] = r_load_a[1];
    s_a[load_a_smem_k + 2][load_a_smem_m] = r_load_a[2];
    s_a[load_a_smem_k + 3][load_a_smem_m] = r_load_a[3];
    
    __syncthreads();
    
    for (int tk = 0; tk < BK; tk++) {
      // 计算时按列读取: s_a[tk][m_range]
      HALF2(r_comp_a[0]) = HALF2(s_a[tk][ty * TM / 2]);
      HALF2(r_comp_a[2]) = HALF2(s_a[tk][ty * TM / 2 + 2]);
      HALF2(r_comp_a[4]) = HALF2(s_a[tk][ty * TM / 2 + BM / 2]);
      HALF2(r_comp_a[6]) = HALF2(s_a[tk][ty * TM / 2 + BM / 2 + 2]);
      
      // FMA 计算
      for (int tm = 0; tm < TM; tm++)
        for (int tn = 0; tn < TN; tn++)
          r_c[tm][tn] += r_comp_a[tm] * r_comp_b[tn];
    }
  }
}

7.4 写入时的 Bank Conflict

转置存储: s_a[load_a_smem_k + i][load_a_smem_m]

tid 0: m=0, k=0~3 → s_a[0~3][0]
tid 1: m=0, k=4~7 → s_a[4~7][0]
tid 2: m=1, k=0~3 → s_a[0~3][1]
tid 3: m=1, k=4~7 → s_a[4~7][1]
...

写入 s_a[k][m]:
├── s_a[0][0] = half[0]   → Bank 0
├── s_a[1][0] = half[128] → Bank 0 (因为 128/2 % 32 = 0)
├── s_a[2][0] = half[256] → Bank 0
└── s_a[3][0] = half[384] → Bank 0

tid 0 写入 s_a[0~3][0]，全部访问 Bank 0
tid 2 写入 s_a[0~3][1]，全部访问 Bank 0

→ 写入时仍有 Bank Conflict (但写入次数少，影响较小)

7.5 Padding 优化 (OFFSET)

__shared__ half s_a[BK][BM + OFFSET];  // OFFSET=8

添加 Padding 后:
每行 128 + 8 = 136 half = 68 Banks

s_a[k][m] 新地址 = k * 136 + m

s_a[0][0] = half[0]     → Bank 0
s_a[1][0] = half[136]   → Bank (136/2) % 32 = 68 % 32 = 4
s_a[2][0] = half[272]   → Bank (272/2) % 32 = 136 % 32 = 8
s_a[3][0] = half[408]   → Bank (408/2) % 32 = 204 % 32 = 12

→ 不同 k 值访问不同 Bank，消除写入时的 Bank Conflict!

7.6 Padding 效果可视化

┌─────────────────────────────────────────────────────────────────────┐
│                    OFFSET=8 Padding 效果                            │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│   无 Padding: s_a[8][128]                                           │
│                                                                     │
│   k=0: [Bank 0][Bank 1][Bank 2]...[Bank 31][Bank 0][Bank 1]...      │
│   k=1: [Bank 0][Bank 1][Bank 2]...[Bank 31][Bank 0][Bank 1]...      │
│   k=2: [Bank 0][Bank 1][Bank 2]...[Bank 31][Bank 0][Bank 1]...      │
│         ↑                                                           │
│         所有 k 的 m=0 都在 Bank 0                                    │
│                                                                     │
│   有 Padding: s_a[8][128 + 8]                                       │
│                                                                     │
│   k=0: [Bank 0][Bank 1][Bank 2]...[Bank 31][Bank 0]...[PAD]         │
│   k=1: [Bank 4][Bank 5][Bank 6]...[Bank 3][Bank 4]...[PAD]          │
│   k=2: [Bank 8][Bank 9][Bank 10]...[Bank 7][Bank 8]...[PAD]         │
│         ↑                                                           │
│         不同 k 的 m=0 在不同 Bank                                    │
│                                                                     │
│   Bank 偏移: Δ = (OFFSET / 2) % 32 = 4                              │
│   k=0 → Bank 0, k=1 → Bank 4, k=2 → Bank 8, ...                     │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘

---

8. Double Buffering 优化

8.1 核心思想

┌─────────────────────────────────────────────────────────────────────┐
│                    Double Buffering 原理                            │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│   传统实现 (无 Double Buffering):                                    │
│                                                                     │
│   ┌────────┐ ┌────────┐ ┌────────┐ ┌────────┐ ┌────────┐           │
│   │ Load 0 │→│Compute0│→│ Load 1 │→│Compute1│→│ Load 2 │→...       │
│   └────────┘ └────────┘ └────────┘ └────────┘ └────────┘           │
│   │←──────────串行执行，无重叠──────────→│                          │
│                                                                     │
│   Double Buffering 实现:                                            │
│                                                                     │
│   Buffer 0: │ Load 0 │         │ Load 2 │         │ Load 4 │       │
│   Buffer 1: │        │ Load 1 │         │ Load 3 │         │       │
│   Compute:  │        │Compute0│Compute1│Compute2│Compute3│         │
│                                                                     │
│   ┌────────┬────────┬────────┬────────┬────────┐                   │
│   │ Load 0 │ Load 1 │ Load 2 │ Load 3 │  ...   │ ← 加载流水线       │
│   ├────────┼────────┼────────┼────────┼────────┤                   │
│   │        │Compute0│Compute1│Compute2│  ...   │ ← 计算流水线       │
│   └────────┴────────┴────────┴────────┴────────┘                   │
│                                                                     │
│   关键: 加载 Buffer[i+1] 与计算 Buffer[i] 并行                       │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘

8.2 代码实现

template<const int BM=128, const int BN=128, const int BK=8, 
         const int TM=8, const int TN=8, const int OFFSET=0>
__global__ void hgemm_t_8x8_sliced_k_f16x8_pack_bcf_dbuf_kernel(...) {
  
  // 双缓冲: 两套 Shared Memory
  __shared__ half s_a[2][BK][BM + OFFSET];
  __shared__ half s_b[2][BK][BN + OFFSET];

  half r_load_a[4], r_load_b[4];  // 加载寄存器
  half r_comp_a[TM], r_comp_b[TN];  // 计算寄存器
  half r_c[TM][TN] = {0};

  // ========== 阶段 0: 预加载第一块到 Buffer 0 ==========
  {
    int load_a_gmem_addr = load_a_gmem_m * K + load_a_smem_k;  // bk=0
    int load_b_gmem_addr = load_b_smem_k * N + load_b_gmem_n;  // bk=0
    
    LDST64BITS(r_load_a[0]) = LDST64BITS(a[load_a_gmem_addr]);
    LDST64BITS(r_load_b[0]) = LDST64BITS(b[load_b_gmem_addr]);

    // 存入 Buffer 0
    s_a[0][load_a_smem_k + 0][load_a_smem_m] = r_load_a[0];
    s_a[0][load_a_smem_k + 1][load_a_smem_m] = r_load_a[1];
    s_a[0][load_a_smem_k + 2][load_a_smem_m] = r_load_a[2];
    s_a[0][load_a_smem_k + 3][load_a_smem_m] = r_load_a[3];
    LDST64BITS(s_b[0][load_b_smem_k][load_b_smem_n]) = LDST64BITS(r_load_b[0]);
  }
  __syncthreads();

  // ========== 主循环: bk = 1 开始 ==========
  for (int bk = 1; bk < (K + BK - 1) / BK; bk++) {
    
    int smem_sel = (bk - 1) & 1;      // 当前计算使用的 Buffer
    int smem_sel_next = bk & 1;       // 下一次加载使用的 Buffer

    // Step 1: 开始加载下一块数据到寄存器 (异步)
    int load_a_gmem_k = bk * BK + load_a_smem_k;
    int load_a_gmem_addr = load_a_gmem_m * K + load_a_gmem_k;
    int load_b_gmem_k = bk * BK + load_b_smem_k;
    int load_b_gmem_addr = load_b_gmem_k * N + load_b_gmem_n;
    
    LDST64BITS(r_load_a[0]) = LDST64BITS(a[load_a_gmem_addr]);  // 异步加载
    LDST64BITS(r_load_b[0]) = LDST64BITS(b[load_b_gmem_addr]);  // 异步加载
    
    // Step 2: 使用当前 Buffer 进行计算 (与加载并行)
    #pragma unroll
    for (int tk = 0; tk < BK; tk++) {
      LDST128BITS(r_comp_a[0]) = LDST128BITS(s_a[smem_sel][tk][ty * TM]);
      LDST128BITS(r_comp_b[0]) = LDST128BITS(s_b[smem_sel][tk][tx * TN]);

      #pragma unroll
      for (int tm = 0; tm < TM; tm++) {
        #pragma unroll
        for (int tn = 0; tn < TN; tn++) {
          r_c[tm][tn] = __hfma(r_comp_a[tm], r_comp_b[tn], r_c[tm][tn]);
        }
      }
    }

    // Step 3: 将加载的数据写入下一个 Buffer (计算完成后)
    s_a[smem_sel_next][load_a_smem_k + 0][load_a_smem_m] = r_load_a[0];
    s_a[smem_sel_next][load_a_smem_k + 1][load_a_smem_m] = r_load_a[1];
    s_a[smem_sel_next][load_a_smem_k + 2][load_a_smem_m] = r_load_a[2];
    s_a[smem_sel_next][load_a_smem_k + 3][load_a_smem_m] = r_load_a[3];
    LDST64BITS(s_b[smem_sel_next][load_b_smem_k][load_b_smem_n]) = LDST64BITS(r_load_b[0]);

    __syncthreads();  // 仅需一次同步!
  }
  
  // ========== 阶段 Final: 计算最后一块 (Buffer 1) ==========
  #pragma unroll
  for (int tk = 0; tk < BK; tk++) {
    LDST128BITS(r_comp_a[0]) = LDST128BITS(s_a[1][tk][ty * TM]);
    LDST128BITS(r_comp_b[0]) = LDST128BITS(s_b[1][tk][tx * TN]);

    #pragma unroll
    for (int tm = 0; tm < TM; tm++) {
      #pragma unroll
      for (int tn = 0; tn < TN; tn++) {
        r_c[tm][tn] = __hfma(r_comp_a[tm], r_comp_b[tn], r_c[tm][tn]);
      }
    }
  }

  // 写回结果
  #pragma unroll
  for (int i = 0; i < TM; i++) {
    int store_c_gmem_m = by * BM + ty * TM + i;
    int store_c_gmem_n = bx * BN + tx * TN;
    LDST128BITS(c[store_c_gmem_m * N + store_c_gmem_n]) = LDST128BITS(r_c[i][0]);
  }
}

8.3 同步次数优化

传统实现 (无 Double Buffering):
每次 K 迭代需要 2 次 __syncthreads()
├── 第 1 次: 确保数据加载完成后再计算
└── 第 2 次: 确保计算完成后再加载下一批

总同步次数 = 2 × K/BK

Double Buffering 实现:
每次 K 迭代仅需 1 次 __syncthreads()
├── 计算使用 Buffer[i]，与加载 Buffer[i+1] 无依赖
└── 只需在写入 Buffer[i+1] 完成后同步

总同步次数 = K/BK

节省同步次数 = K/BK 次

8.4 流水线时序

┌─────────────────────────────────────────────────────────────────────┐
│                 Double Buffering 流水线时序图                        │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│   bk=0    │  Load[0]→Buf0  │                                        │
│           │  __sync__      │                                        │
│                                                                     │
│   bk=1    │  Load[1]→Reg   │  Compute Buf0  │  Reg→Buf1  │ __sync__ │
│                                                                     │
│   bk=2    │  Load[2]→Reg   │  Compute Buf1  │  Reg→Buf0  │ __sync__ │
│                                                                     │
│   bk=3    │  Load[3]→Reg   │  Compute Buf0  │  Reg→Buf1  │ __sync__ │
│                                                                     │
│   ...                                                               │
│                                                                     │
│   Final   │                │  Compute Buf1  │                       │
│                                                                     │
│   时间轴  │←────────────────────────────────────────────────────→│  │
│                                                                     │
│   关键: Load[bk]→Reg 与 Compute Buf[bk-1] 并行执行                  │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘

---

9. 性能对比与分析

9.1 各版本性能对比 (相对性能)

┌──────────────────────────────────────────────────────────────────────┐
│                        性能对比 (估算)                               │
├──────────────────────────────────────────────────────────────────────┤
│                                                                      │
│   Kernel                              │ 相对性能 │ 主要优化           │
│   ────────────────────────────────────┼──────────┼──────────────────  │
│   hgemm_naive_f16_kernel              │   1.0×   │ 基准               │
│   hgemm_sliced_k_f16_kernel           │   5.0×   │ +Shared Memory     │
│   hgemm_t_8x8_sliced_k_f16x4_kernel   │  20.0×   │ +Thread Tile       │
│   hgemm_t_8x8_sliced_k_f16x4_pack     │  25.0×   │ +Pack Load/Store   │
│   hgemm_t_8x8_sliced_k_f16x4_bcf      │  40.0×   │ +Transpose BCF     │
│   hgemm_t_8x8_sliced_k_f16x8_pack_bcf │  50.0×   │ +Padding OFFSET    │
│   hgemm_t_8x8_sliced_k_f16x8_dbuf     │  60.0×   │ +Double Buffering  │
│   cuBLAS hgemm                        │ ~80.0×   │ 高度优化库         │
│                                                                      │
└──────────────────────────────────────────────────────────────────────┘

9.2 各版本计算密度

Kernel	每次 K 迭代 FMA	每次 K 迭代 Load	计算密度
naive	1	2	0.5
sliced_k	1	0 (from smem)	~1
t_8x8	64	16	4.0
t_8x8_f16x4	64	16	4.0
bcf/dbuf	64	16	4.0

9.3 Shared Memory 使用量

Kernel	s_a	s_b	总计	备注
sliced_k	32×32×2 = 2KB	32×32×2 = 2KB	4KB	基础
t_8x8	128×8×2 = 2KB	8×128×2 = 2KB	4KB	Tiling
bcf (OFFSET=4)	8×132×2 = 2.1KB	8×132×2 = 2.1KB	4.2KB	Padding
dbuf (OFFSET=8)	2×8×136×2 = 4.3KB	2×8×136×2 = 4.3KB	8.6KB	Double

9.4 Bank Conflict 分析

Kernel	s_a 写入	s_a 读取	s_b 写入	s_b 读取
t_8x8_f16x4	无冲突	16-way	2-way	2-way
bcf (转置)	4-way	2-way	2-way	2-way
bcf (OFFSET)	无冲突	无冲突	无冲突	4-way
dbuf	无冲突	无冲突	无冲突	4-way

---

10. 总结与最佳实践

10.1 优化技术总结

┌─────────────────────────────────────────────────────────────────────┐
│                        HGEMM 优化技术栈                             │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│   Level 1: 内存层次优化                                              │
│   ├── Shared Memory Tiling (Block Tile)                             │
│   ├── Register Tiling (Thread Tile)                                 │
│   └── K 维度分块 (K Tile)                                           │
│                                                                     │
│   Level 2: 访存优化                                                  │
│   ├── 向量化加载/存储 (half2, float4, LDST128BITS)                  │
│   ├── Bank Conflict 消除 (转置 + Padding)                           │
│   └── 内存合并访问 (Coalesced Access)                               │
│                                                                     │
│   Level 3: 延迟隐藏                                                  │
│   ├── Double Buffering (流水线)                                     │
│   ├── 寄存器预取                                                     │
│   └── 循环展开 (#pragma unroll)                                     │
│                                                                     │
│   Level 4: 指令级优化                                                │
│   ├── FMA 指令 (__hfma)                                             │
│   ├── 减少分支                                                       │
│   └── 指令级并行 (ILP)                                              │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘

10.2 最佳实践清单

✅ Shared Memory 使用

// 1. 选择合适的 Tile 大小
constexpr int BM = 128, BN = 128, BK = 8;  // 平衡 smem 使用和复用

// 2. 添加 Padding 消除 Bank Conflict
__shared__ half s_a[BK][BM + OFFSET];  // OFFSET = 8

// 3. 考虑转置以优化访问模式
// s_a[BK][BM] 而非 s_a[BM][BK]

✅ 向量化访存

// 1. 使用向量类型
#define LDST64BITS(value)  (reinterpret_cast<float2*>(&(value))[0])
#define LDST128BITS(value) (reinterpret_cast<float4*>(&(value))[0])

// 2. 合并小的加载/存储
LDST128BITS(r_comp_a[0]) = LDST128BITS(s_a[tk][ty * TM]);  // 8 half 一次

✅ Double Buffering

// 1. 分配双缓冲
__shared__ half s_a[2][BK][BM + OFFSET];

// 2. 预加载第一块
// 3. 主循环中交替使用
int smem_sel = (bk - 1) & 1;
int smem_sel_next = bk & 1;

// 4. 先加载到寄存器，再计算，最后写 smem

✅ 循环展开

// 1. 展开内层循环
#pragma unroll
for (int tm = 0; tm < TM; tm++) {
  #pragma unroll
  for (int tn = 0; tn < TN; tn++) {
    r_c[tm][tn] = __hfma(r_comp_a[tm], r_comp_b[tn], r_c[tm][tn]);
  }
}

10.3 进一步优化方向

方向	技术	预期收益
Tensor Core	WMMA/MMA 指令	2-4×
Swizzle	手动地址重映射	10-20%
Async Copy	cp.async 指令	15-25%
Multi-Stage	多级流水线	10-15%
Warp Specialization	分离 Load/Compute	10-20%

10.4 参考资源

• CUTLASS: CUDA Templates for Linear Algebra Subroutines
• NVIDIA CUDA Programming Guide
• Bank Conflicts in Shared Memory
• Optimizing Matrix Multiplication

---

附录 A: 宏定义说明

// 向量类型转换宏
#define HALF2(value)      (reinterpret_cast<half2*>(&(value))[0])
#define LDST64BITS(value) (reinterpret_cast<float2*>(&(value))[0])   // 8 bytes = 4 half
#define LDST128BITS(value)(reinterpret_cast<float4*>(&(value))[0])  // 16 bytes = 8 half

// 使用示例
half r_load[4];
LDST64BITS(r_load[0]) = LDST64BITS(a[addr]);  // 一次加载 4 个 half

half s_a[8][128];
LDST128BITS(r_comp[0]) = LDST128BITS(s_a[k][m]);  // 一次加载 8 个 half

---

附录 B: Kernel 配置参数

// 优化版 Kernel 的标准配置
constexpr int BM = 128;       // Block Tile M 维度
constexpr int BN = 128;       // Block Tile N 维度
constexpr int BK = 8;         // K 维度分块大小
constexpr int TM = 8;         // Thread Tile M 维度
constexpr int TN = 8;         // Thread Tile N 维度
constexpr int OFFSET = 8;     // Bank Conflict Padding

// 执行配置
dim3 block(BN/TN, BM/TM);     // (16, 16) = 256 threads
dim3 grid((N+BN-1)/BN, (M+BM-1)/BM);

// Shared Memory 大小
// 单缓冲: 2 × BK × (BM + OFFSET) × sizeof(half) = 2 × 8 × 136 × 2 = 4.25 KB
// 双缓冲: 2 × 2 × BK × (BM + OFFSET) × sizeof(half) = 8.5 KB

---

文档版本: v1.0 最后更新: 2026-02-28 作者: CUDA HGEMM 优化研究