CUDA cute库深度解析：基于Ampere架构的TN HGEMM与三级流水线实现

1 任务背景

本文将以cute库（Cutlass 3.9.0版本）中基于Ampere架构的GEMM示例代码为分析对象，具体路径为 cutlass/examples/cute/tutorial/sgemm_sm80.cu。该文件提供了TN GEMM、NT GEMM和TN HGEMM等多种高性能计算方案，本文将重点剖析TN HGEMM的实现。

TN表示矩阵乘法的模式：矩阵A转置（T），矩阵B不转置（N）。HGEMM表示数据类型为half。在典型的Attention计算场景中，矩阵A和B的形状常为[5120, 4096]。对应的核心PTX指令可能为mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16。

源代码的主要调用链为：main -> gemm -> gemm_tn -> gemm_device(Kernel)。核心计算逻辑集中在gemm_tn和gemm_device函数中。下文将对其关键代码进行逐行解析。

2 gemm_tn 函数

2.1 Layout 定义

2.1.1 Layout 参数设置

首先定义问题的形状及矩阵在全局内存中的步长。在TN模式下，矩阵A、B、C的形状分别为[m, k]、[n, k]、[m, n]，内存序分别为行主序、行主序、列主序。

auto M = int(m);
auto N = int(n);
auto K = int(k);
auto prob_shape = make_shape(M, N, K);              // (M, N, K)
auto dA = make_stride(ldA, Int<1>{});               // (dM, dK)
auto dB = make_stride(ldB, Int<1>{});               // (dN, dK)
auto dC = make_stride(Int<1>{}, ldC);               // (dM, dN)

接着定义线程块（CTA）级别的分块大小，并设置流水线缓冲区数量bP为3。

auto bM = Int<128>{};
auto bN = Int<128>{};
auto bK = Int<64>{};
auto cta_tiler = make_shape(bM, bN, bK);            // (BLK_M, BLK_N, BLK_K)
auto bP = Int<3>{};  // Pipeline

2.1.2 swizzle 机制

为避免从共享内存加载数据时的bank conflict，代码采用了swizzle策略。针对half类型和bK=64，选择了B=M=S=3的参数组合。

auto swizzle_atom = composition(Swizzle<3, 3, 3>{},
                                 Layout<Shape<_8, Shape<_8, _8>>,
                                        Stride<_8, Stride<_1, _64>>>{});
auto sA_layout = tile_to_shape(swizzle_atom, make_shape(bM, bK, bP));
auto sB_layout = tile_to_shape(swizzle_atom, make_shape(bN, bK, bP));
auto sC_layout = make_layout(make_shape(bM, bN));

swizzle_atom将一个(8, (8, 8)):(8, 1, 64)的布局与Swizzle变换组合。其效果是将原本连续加载的、可能导致bank conflict的一行数据（属于前4个bank），通过地址变换分散到32个不同的bank中，从而消除冲突。此变换需要在数据写入共享内存和从共享内存读出时各应用一次，以保证最终存入寄存器的数据布局与原始全局内存中的一致。

2.2 copy 方式

2.2.1 从全局内存移动到共享内存

此过程使用了Ampere架构引入的cp.async.ca.shared.global.L2::128B异步拷贝指令，能够将数据直接从全局内存搬移至共享内存，不经过寄存器。

TiledCopy copyA = make_tiled_copy(Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<uint128_t>, cute::half_t>{},
                                   Layout<Shape<_16,_8>,Stride<_8,_1>>{},
                                   Layout<Shape< _1,_8>>{});
TiledCopy copyB = make_tiled_copy(Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<uint128_t>, cute::half_t>{},
                                  Layout<Shape<_16,_8>,Stride<_8,_1>>{},
                                  Layout<Shape< _1,_8>>{});

Copy_Atom定义了一次拷贝128 bit（8个half元素）。ThrLayout定义了线程块内16x8（行主序）的线程排布。ValLayout设置为(1, 8)，与矩阵的行主序内存布局匹配。综合来看，一个线程块一次可完成形状为(16*1, 8*8)即(16, 64)的分片拷贝。

2.2.2 从共享内存移动到寄存器

此操作为后续的MMA计算准备数据，需要根据MMA的布局来定义拷贝方式。首先定义TiledMMA对象。

TiledMMA mmaC = make_tiled_mma(SM80_16x8x8_F16F16F16F16_TN{},
    Layout<Shape<_2,_2>>{},    // 2x2x1 MMA Atoms
    Tile<_32,_32,_16>{});      // 32x32x16 Tiled MMA for LDSM

MMA_Atom SM80_16x8x8_F16F16F16F16_TN封装了mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16指令。ThrLayout (2,2)表示一个线程块内4个warp的排布。Tile<_32,_32,_16>参数将每个warp的MMA计算形状从m16n8k8拓展到线程块级别的m32n32k16。

接着定义从共享内存到寄存器的拷贝原子操作，它基于Turing架构的ldmatrix.sync.aligned.x4.m8n8.shared.b16指令，一个warp一次加载4个m8n8矩阵。

Copy_Atom<SM75_U32x4_LDSM_N, half_t> s2r_atom_A;
Copy_Atom<SM75_U32x4_LDSM_N, half_t> s2r_atom_B;

通过make_tiled_copy_A/B函数，将上述Copy_Atom与TiledMMA对象结合，生成最终的TiledCopy对象。该对象描述了数据如何从共享内存布局映射到满足MMA计算要求的寄存器布局。

2.3 启动 Kernel

首先计算并设置所需的动态共享内存大小。SharedStorage结构体根据ASmemLayout和BSmemLayout的cosize申请内存，cosize反映了布局实际占用的空间（可能包含填充）。

template <class ElementA,
          class ElementB,
          class SmemLayoutA,
          class SmemLayoutB>
struct SharedStorage
{
   cute::ArrayEngine<ElementA, cute::cosize_v<SmemLayoutA>> A;
   cute::ArrayEngine<ElementB, cute::cosize_v<SmemLayoutB>> B;
};
int smem_size = int(sizeof(SharedStorage<cute::half_t, cute::half_t, decltype(sA), decltype(sB)>));

设置核函数属性，将最大动态共享内存容量设为smem_size，并将Unified Cache中shared memory的比例设置为100%（即L1 Cache占比为0），以最大化共享内存容量。

cudaFuncSetAttribute(
    kernel_fptr,
    cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
cudaFuncSetAttribute(
    kernel_fptr,
    cudaFuncAttributePreferredSharedMemoryCarveout, 100);

最后配置线程块与网格维度，启动核函数。

3 gemm_device 函数

3.1 形状参数静态断言

核函数开头包含一系列编译期静态断言，用于校验输入Layout、形状、步长等参数是否匹配且有效，确保后续计算正确性。

3.2 定义全局内存和共享内存中的矩阵 Tensor

Tensor mA = make_tensor(make_gmem_ptr(A), select<0,2>(shape_MNK), dA); // (M,K)
Tensor mB = make_tensor(make_gmem_ptr(B), select<1,2>(shape_MNK), dB); // (N,K)
Tensor mC = make_tensor(make_gmem_ptr(C), select<0,1>(shape_MNK), dC); // (M,N)

定义代表完整矩阵的Tensor。

auto cta_coord = make_coord(blockIdx.x, blockIdx.y, _);              // (m,n,k)
Tensor gA = local_tile(mA, cta_tiler, cta_coord, Step<_1, X,_1>{});  // (BLK_M,BLK_K,k)
Tensor gB = local_tile(mB, cta_tiler, cta_coord, Step< X,_1,_1>{});  // (BLK_N,BLK_K,k)
Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{});  // (BLK_M,BLK_N)

通过local_tile获取当前线程块负责处理的全局内存分块Tensor。其中k = K/BLK_K。

extern __shared__ char shared_memory[];
SharedStorage& smem = *reinterpret_cast<SharedStorage*>(shared_memory);
Tensor sA = make_tensor(make_smem_ptr(smem.A.begin()), sA_layout);   // (BLK_M,BLK_K,PIPE)
Tensor sB = make_tensor(make_smem_ptr(smem.B.begin()), sB_layout);   // (BLK_N,BLK_K,PIPE)

定义共享内存中的Tensor。

3.3 g2s 拷贝对象的线程级划分

将2.2.1中定义的TiledCopy对象划分到线程粒度。

ThrCopy thr_copy_a = copy_a.get_slice(threadIdx.x);
Tensor tAgA = thr_copy_a.partition_S(gA);          // (CPY,CPY_M,CPY_K,k)
Tensor tAsA = thr_copy_a.partition_D(sA);          // (CPY,CPY_M,CPY_K,PIPE)

tAgA和tAsA分别表示当前线程负责的源（全局内存）和目标（共享内存）数据视图。CPY=8（每个线程一次拷贝8个元素），CPY_M=8（在M方向重复8次），CPY_K=1（一次拷贝K方向的一整行）。

3.4 从全局内存预加载（prefetch）到共享内存

GEMM计算可被组织为三级流水线：1) 全局内存到共享内存（g2s），2) 共享内存到寄存器（s2r），3) MMA计算。要实现并行，需要在共享内存中设置多个缓冲区（本例中PIPE=3）。

auto K_PIPE_MAX = size<3>(tAsA);
int k_tile_count = size<3>(tAgA);
int k_tile_next = 0;
for (int k_pipe = 0; k_pipe < K_PIPE_MAX-1; ++k_pipe) {
    copy(copy_a, tAgA(_,_,_,k_tile_next), tAsA(_,_,_,k_pipe));
    copy(copy_b, tBgB(_,_,_,k_tile_next), tBsB(_,_,_,k_pipe));
    cp_async_fence();
    --k_tile_count;
    if (k_tile_count > 0) { ++k_tile_next; }
}

在主循环开始前，先异步提交前两个K分块（Tile）到共享内存的缓冲区0和1，实现预取。

3.5 mma 对象的线程级划分

将2.2.2中定义的TiledMMA对象划分到线程粒度，并分配寄存器。

ThrMMA thr_mma = mma.get_slice(threadIdx.x);
Tensor tCgC = thr_mma.partition_C(gC);                     // (MMA,MMA_M,MMA_N)
Tensor tCrA = thr_mma.partition_fragment_A(sA(_,_,0));     // (MMA,MMA_M,MMA_K)
Tensor tCrB = thr_mma.partition_fragment_B(sB(_,_,0));     // (MMA,MMA_N,MMA_K)
Tensor tCrC = thr_mma.make_fragment_C(tCgC);               // (MMA,MMA_M,MMA_N)
clear(tCrC);

MMA维度由具体的MMA_Atom指令决定（例如对于C矩阵，每个线程持有(2,2)个元素）。MMA_M、MMA_N、MMA_K表示在线程块维度，需要在M、N、K方向重复的次数。已知(BLK_M, BLK_N, BLK_K) = (128,128,64)，TiledMMA扩展后的单次计算形状为(32,16,8)，因此MMA_M = 128/32 = 4，MMA_N = 128/16 = 8，MMA_K = 64/8 = 8。tCrC被用作累加器并初始化为零。

3.6 s2r Copy Atom 重排

创建用于s2r拷贝的TiledCopy对象，并进行线程级划分。

TiledCopy s2r_copy_a = make_tiled_copy_A(s2r_atom_a, mma);
ThrCopy   s2r_thr_copy_a = s2r_copy_a.get_slice(threadIdx.x);
Tensor tXsA = s2r_thr_copy_a.partition_S(sA);              // (CPY,MMA_M,MMA_K,PIPE)
Tensor tXrA = s2r_thr_copy_a.retile_D(tCrA);               // (CPY,MMA_M,MMA_K)

tXsA是当前线程看到的共享内存源数据视图，tXrA则定义了如何将数据重排到寄存器tCrA的目标布局。CPY=8（ldmatrix指令每个线程加载8个元素），MMA_M=MMA_N=MMA_K=4（由Tile<_32,_32,_16>与分块大小计算得出）。

3.7 缓冲区（管道）初始化

设置共享内存读写缓冲区的初始索引。

int smem_pipe_read  = 0;
int smem_pipe_write = K_PIPE_MAX-1;
Tensor tXsA_p = tXsA(_,_,_,smem_pipe_read);
Tensor tXsB_p = tXsB(_,_,_,smem_pipe_read);
auto K_BLOCK_MAX = size<2>(tCrA); // = MMA_K

3.8 初始化：从共享内存预加载数据到寄存器

在开始主循环计算前，确保第一个全局内存分块已加载到共享内存，并预取其第一个K_BLOCK（寄存器级分块）到寄存器。

if (K_BLOCK_MAX > 1) {
    cp_async_wait<K_PIPE_MAX-2>();
    __syncthreads();
    copy(s2r_atom_a, tXsA_p(_,_,Int<0>{}), tXrA(_,_,Int<0>{}));
    copy(s2r_atom_b, tXsB_p(_,_,Int<0>{}), tXrB(_,_,Int<0>{}));
}

3.9 主循环：流水线化的读写与计算

主循环实现了g2s、s2r、mma三级操作的重叠。

while (k_tile_count > -(K_PIPE_MAX-1)) {
    for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block) {
        // 1. 更新共享内存读取管道（当处理到最后一个寄存器分块时）
        if (k_block == K_BLOCK_MAX - 1) {
            tXsA_p = tXsA(_,_,_,smem_pipe_read);
            tXsB_p = tXsB(_,_,_,smem_pipe_read);
            cp_async_wait<K_PIPE_MAX-2>();
            __syncthreads();
        }
        // 2. smem→rmem复制（预取下一个寄存器分块）
        auto k_block_next = (k_block + Int<1>{}) % K_BLOCK_MAX;
        copy(s2r_atom_a, tXsA_p(_,_,k_block_next), tXrA(_,_,k_block_next));
        copy(s2r_atom_b, tXsB_p(_,_,k_block_next), tXrB(_,_,k_block_next));

        // 3. gmem→smem复制（加载下一个全局内存分块）
        if (k_block == 0) {
            copy(copy_a, tAgA(_,_,_,k_tile_next), tAsA(_,_,_,smem_pipe_write));
            copy(copy_b, tBgB(_,_,_,k_tile_next), tBsB(_,_,_,smem_pipe_write));
            cp_async_fence();
            --k_tile_count;
            if (k_tile_count > 0) { ++k_tile_next; }
            smem_pipe_write = smem_pipe_read;
            smem_pipe_read = (smem_pipe_read == K_PIPE_MAX-1) ? 0 : smem_pipe_read+1;
        }
        // 4. 寄存器级GEMM计算（当前分块）
        gemm(mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCrC);
    }
}

循环重叠逻辑精要：

• 内层循环 (k_block): 在K的寄存器分块维度循环。
• 步骤2 (s2r): 总是预取下一个(k_block_next)寄存器分块，与当前(k_block)的MMA计算重叠。
• 步骤3 (g2s): 仅在k_block == 0（处理一个新的全局分块的起始时）提交下一个全局分块的异步加载。更新缓冲区索引，循环利用共享内存。
• 步骤4 (mma): 对当前(k_block)分块进行计算。
• 同步: cp_async_wait确保待读取的数据已就绪；__syncthreads保证线程块内对共享内存的读写一致性。这种精细的并行计算控制是发挥GPU性能的关键。

3.10 收尾（Epilogue）

将寄存器中的累加结果与全局内存中的原始C矩阵进行线性组合，完成 C = alpha*A*B + beta*C 运算。

axpby(alpha, tCrC, beta, tCgC);

4 小结

本文以cute库中基于Ampere架构的TN HGEMM示例为蓝本，详细剖析了其利用Swizzle消除Bank Conflict、通过异步拷贝指令构建三级流水线以隐藏内存延迟等核心优化策略的具体实现。深入理解这些底层机制，对于在CUDA编程中实现极致性能的高性能计算内核至关重要。