关于Nvidia PTX中mma指令的技术疑问及代码解析

对PTX中MMA指令的初步理解 #

• 它是按warp执行的指令
• 需要将特定元素加载到目标warp内每个线程的寄存器中
• 乘法的尺寸是固定的（有有限的可选集合）

示例PTX指令片段 #

//
// Generated by LLVM NVPTX Back-End
//

.version 8.4
.target sm_89
.address_size 64

    // .globl   matmul_kernel
.extern .shared .align 16 .b8 global_smem[];

.visible .entry matmul_kernel(
    .param .u64 matmul_kernel_param_0,
    .param .u64 matmul_kernel_param_1,
    .param .u64 matmul_kernel_param_2,
    .param .u32 matmul_kernel_param_3,
    .param .u32 matmul_kernel_param_4,
    .param .u32 matmul_kernel_param_5,
    .param .u32 matmul_kernel_param_6,
    .param .u32 matmul_kernel_param_7,
    .param .u32 matmul_kernel_param_8
)
.maxntid 128, 1, 1
{
    ...
    ldmatrix.sync.aligned.m8n8.x4.shared.b16 { %r3100, %r3101, %r3102, %r3103 }, [ %r561 + 0 ];
    ...
    ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 { %r3084, %r3085, %r3086, %r3087 }, [ %r581 + 0 ];
    ...
    mov.f32     %f2306, 0f00000000;
    mov.b32     %r3107, 2;
    mov.b32     %r3106, 0;
    shl.b32     %r2885, %r100, 1;
    shl.b32     %r2894, %r101, 1;
    shl.b32     %r2895, %r102, 1;
    shl.b32     %r2896, %r103, 1;
    shl.b32     %r2897, %r104, 1;
    shl.b32     %r2898, %r105, 1;
    shl.b32     %r2899, %r106, 1;
    mov.u32     %r3104, %r765;
    mov.u32     %r3105, %r758;
    mov.f32     %f2307, %f2306;
    mov.f32     %f2308, %f2306;
    mov.f32     %f2309, %f2306;
    ...
    mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 { %f2306, %f2307, %f2308, %f2309 }, { %r3100, %r3101, %r3102, %r3103 }, { %r3084, %r3085 }, { %f2306, %f2307, %f2308, %f2309 };
}

技术疑问解答 #

1. 矩阵元素的寄存器分配逻辑 #

你的理解完全正确。

• A矩阵（16×16的f16）：每个32位寄存器可存储2个f16元素，4个寄存器对应8个元素；warp共32线程，总元素数8×32=256，刚好匹配16×16的规模。
• B矩阵（16×8的f16）：每个线程用2个32位寄存器存储4个f16元素，32线程总元素数4×32=128，对应16×8的尺寸。
• 结果矩阵D（16×8的f32）：每个线程持有4个f32元素，32线程总元素数4×32=128，正好匹配16×8的f32矩阵规模。

2. sync标记的作用 #

sync不是等待Tensor Core完成运算，而是保证warp内所有线程都完成了前置的ldmatrix加载操作，同步后再执行MMA指令。因为MMA指令需要warp内所有线程的寄存器数据全部准备到位才能启动，sync用于避免数据不一致的问题。

3. MMA尺寸的兼容性 #

PTX 8.4中，大尺寸的MMA类型支持拆分使用更小的合法子尺寸。比如m16n8k64的Tensor Core可以被配置为执行m16n8k16的运算，本质是只利用大尺寸Tensor Core中对应k=16的部分。官方文档列出的是基础支持的最大尺寸，实际可向下兼容符合m≤16、n≤8、k≤64的子尺寸。

4. 密集型与稀疏型矩阵的定义 #

• 密集型（dense）：矩阵中绝大多数元素为非零值，运算时会处理每一个元素，无稀疏性优化。
• 稀疏型（sparse）：矩阵中存在大量零元素，对应的MMA指令会跳过零元素的乘法运算以提升效率。这类指令要求输入矩阵按特定稀疏格式存储，Tensor Core可直接识别并处理稀疏数据。

内容的提问来源于stack exchange，提问作者Shore