Tensor Core编程PTX

可视化演示

双模式展示PTX编程：MMA指令执行流程（Shm→ldm加载→frag寄存器→mma.sync→输出）和代际对比（V100 FP16/A100 BF16/H100 FP8算力演进）。

PTX指令概览

码头调度室里有两套文档：一套是物流公司给的“运输总单”（CUDA C++ 源码），只写从哪个仓库取货、送到哪个泊位；另一套是龙门吊司机实际执行的“吊装作业单”（PTX 指令），精确到吊具型号、起吊角度、放箱位置。总单经过一层翻译，变成作业单，再经过一层翻译，变成司机手上的操作按钮——这就是 GPU 的两级编译：CUDA C++ 先编译成 PTX，再由驱动编译成 SASS 机器码。

PTX 的全称是 Parallel Thread Execution（并行线程执行），它是 NVIDIA 定义的虚拟指令集。它的巧妙之处在于：作业单不绑定具体某台吊机型号，老吊机和新吊机都能读懂同一份作业单，只是实际动作快慢不同。这保证了 PTX 编写的代码能跨越多代 GPU 运行，同时也让程序员能够触及硬件能力的“最后一公里”。

在 GPU 港口里，最小的“作业班组”是 warp，一个 warp 恰好有 32 名工人（线程）。龙门吊的每一次起吊，都需要这 32 名工人同时握住吊具、齐声喊号子一起发力。因此 mma 指令是 warp 级指令——它要求一个 warp 的 32 个线程全部参与、协同完成。在 CUDA 内核中，你可以用 asm volatile 嵌入 PTX 指令，这就是“内联 PTX”，相当于在总单中直接粘贴一份吊装作业单。

回到码头：如果你想让龙门吊高效工作，就先要学会读懂作业单的格式。PTX 是一个庞大的指令集，而我们最关心的那一页，写着 mma——矩阵乘加指令。它就是龙门吊最核心的“组合动作”手册。

mma指令语法

龙门吊最厉害的一招叫“组合吊装”：吊具同时抓起一整片集装箱（矩阵 A 的某个分块），在空中与另一片集装箱（矩阵 B 的分块）对位交错，落到已有的货堆（矩阵 C）上，瞬间完成“乘加叠放”。整个过程一条指令搞定，这就是 mma（matrix multiply-accumulate）。普通 CUDA 核心要一条条指令地做“乘一次、加一次”，而 Tensor Core 把整个分块的乘加打包成了一个硬件动作。

在 PTX 中，一次组合吊装写成这样的作业单：

逐个词拆开看：mma 是操作码，sync 表示班组内需要同步，aligned 表示所有地址与数据类型必须对齐；m16n8k8 是本次吊装的货物规格——M 维度 16、N 维度 8、K 维度 8，意味着 A 是 16×8 的货片，B 是 8×8 的货片，C 和 D 都是 16×8 的货堆。随后的 .row、.col 分别是 A、B 矩阵的“码放朝向”（行主序或列主序），最后的 .f32.f16.f16.f32 依次声明 D、A、B、C 四种货物的材质（数据类型）。

看到这里你可能会问：作业单上的 A、B、C、D 究竟存放在哪里？龙门吊的吊具上挂着四个作业区，每个区由 32 名工人分别掌管一部分货物——这就是矩阵片段寄存器，下一节我们走进这四个作业区。

mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {d}, {a}, {b}, {c};

矩阵片段寄存器

龙门吊的吊具不是一个大铁板，而是由 32 名工人手持的托盘拼成的。每名工人手里只攥着几件货物（几个数），这些货物就存放在工人胸前的工具口袋（寄存器）里。整个吊具的 A 区、B 区、C 区、D 区，分别对应 mma 指令的四个矩阵片段（fragment）。没有集中的大仓库，所有数据都分散在每名工人的口袋中。

以 m16n8k8 的 A 片段为例：A 矩阵共有 16×8=128 个元素，分给 32 名工人，每人恰好拿 4 个元素。如果元素是 fp16 半精度（16 位），那么每 2 个元素可以塞进一个 32 位寄存器，所以每名工人口袋里是 2 个寄存器。同理，B 片段 8×8=64 个元素分给 32 人，每人 2 个元素（1 个寄存器）；C 和 D 片段是 fp32 累加结果，16×8=128 个 32 位数分下去，每人 4 个寄存器。

这些片段不是随意分配的，而是严格按 lane ID（工人在班组中的编号）映射。第 0 号工人拿 A 矩阵第一行的前几个元素，第 1 号工人拿紧接着的几个……硬件工程师在文档里画好了完整的“分货表”。正因如此，同一个寄存器在不同 mma 指令中可以被不同片段复用——上一轮还是 A 片段的寄存器，下一轮经过 ldmatrix 重新装载，就变成了新的 B 片段。这就是寄存器复用：工具口袋不变，装的东西可以换。理解了四个作业区，下一步要搞清楚货物从大仓库（全局内存）运到作业区途中的码放规则。

内存布局约束

集装箱在货场上码放是有规章的：箱子必须全部朝一个方向，吊具才能一次抓牢；如果朝向混乱，龙门吊一抓就散。同样，Tensor Core 要求矩阵以特定的“主序”排列——行主序（row-major）或列主序（col-major）。mma 指令中的 .row 和 .col 修饰符，就是告诉硬件“这批货是横着码还是竖着码”。行主序的相邻元素在内存中横向挨着，列主序则纵向挨着。

货物从港口外的大仓库（全局内存）运到龙门吊旁边的中转堆场（共享内存）之后，还需要一辆专用的“集装箱转运车”把它们整齐装到工人的口袋里，这辆车就是 ldmatrix。一条 ldmatrix.sync.aligned.m8n8.x4 指令，可以同时从共享内存抓取 4 个 8×8 的矩阵块，精准分配到 32 名工人的寄存器中——这个动作比普通 ld.shared 逐元素装载快得多，因为它完全摸清了龙门吊的分货表。

但转运车对货物摆放有苛刻要求：每个 8×8 矩阵块的起始地址必须 16 字节对齐。除此之外，32 名工人同时去共享内存取货时，还可能撞上同一个“装卸通道”（bank），导致取货排队。共享内存有 32 个 bank，每个 bank 每次服务一个 4 字节单元；如果多名工人同时访问同一个 bank，就会被串行处理，这就是 bank 冲突。解决方式是使用 swizzle 模式：码放货物时故意按异或规则打乱位置，让工人的访问均匀分散到不同 bank。

回到码头：货物码放规则解决了，接下来要考虑选什么材质的集装箱。不同的货物需要不同的箱子，而 Tensor Core 对不同“材质”的运算速度与精度也大不相同。

数据类型精度

码头上备有四种集装箱：标准钢箱（FP16）、重型加固箱（TF32）、轻量铝箱（BF16）和速装箱（INT8）。选用哪种箱子，直接影响龙门吊的单次吊装重量（吞吐量）和货物安全（精度）。FP16 半精度是 Tensor Core 最经典的搭档：只有 16 位，尾数 10 位，数值范围有限，但两个 FP16 能拼进一个 32 位寄存器，计算吞吐翻倍，是深度学习训练和推理的主力材质。

TF32 是专门为“既要精度又要速度”准备的折中方案：它保留了 FP32 的 8 位指数（数值范围与 FP32 相同），但把尾数砍到 10 位。相当于重型箱——外面尺寸和 FP32 一样大，但内壁加厚，装贵重货物不会溢出。BF16 则反过来，指数 8 位保范围、尾数仅 7 位，适合对数值范围敏感但对尾数精度不敏感的场合。INT8 则是把货物压缩打包，配合 FP32 累加器，常用于推理阶段的量化加速。

如何从数学上理解 TF32 的截断？FP32 有 23 位尾数，TF32 只保留高 10 位，等价于把尾数的小数部分“四舍五入”到 10 位精度——或者说，把原数放大 2 的 13 次方倍、截掉小数部分、再缩小回来。

选好了集装箱材质，最后一步是优化整个码头的作业节奏。再快的龙门吊，如果工人和转运车总是互相等待，吞吐量也上不去——这就是性能调优的用武之地。

性能调优技巧

码头运营的核心目标是“吊机不停”——龙门吊每秒钟都应该在干活，而不是等货。最有效的办法是双缓冲：准备两块中转堆场，龙门吊正在计算堆场 A 里的货物时，转运车已经往堆场 B 卸下了下一批货；等 A 算完，立刻切换到 B，转运车再回头给 A 补货。计算与搬运像两条传送带一样重叠进行，这就是软件流水线。

在 PTX 层面，这需要精细控制每条 ldmatrix 和 mma 的发射顺序。更深一层的约束来自寄存器压力：每名工人的工具口袋数量是有限的（每个线程最多 255 个寄存器）。mma 的四个片段本就要占用不少寄存器，如果为了流水线多缓存几批货，寄存器就会超支，导致内核无法启动或占用率暴跌。占用率（occupancy）决定了一个 SM 上能同时驻留多少个 warp——驻留越多，当一个 warp 等待内存时，另一个 warp 就能顶上干活，延迟就被“藏”了起来。

至此，我们从读懂作业单、掌握组合动作、分配四区货物、规范码放规则、选择集装箱材质，到最终优化码头调度，完整走了一遍 Tensor Core 的 PTX 编程流程。龙门吊依然在轰鸣，但现在的你，已经能亲手写出那份让它发挥极限效率的吊装作业单了。

AI 应用场景

Python 演示

import numpy as np  def to_fp16(x):    """模拟FP16在Tensor Core中的输入：转半精度后转回fp32参与累加。"""    return x.astype(np.float16).astype(np.float32)  def to_tf32(x):    """模拟TF32截断：保留fp32的8位指数,尾数从23位截到10位。"""    bits = np.ascontiguousarray(x, dtype=np.float32).view(np.uint32)    bits = bits & np.uint32(0xFFFFE000)    return bits.view(np.float32)  def to_fp32(x):    """FP32基准:不做任何精度截断。"""    return x.astype(np.float32)  def mma_accumulate(A, B, C, cast_fn):    """模拟mma指令语义 D = A x B + C,输入先按指定精度舍入。"""    Aq = cast_fn(A)    Bq = cast_fn(B)    return Aq @ Bq + C  def run_demo():    np.random.seed(42)    M, K, N = 128, 64, 64    A = np.random.randn(M, K).astype(np.float32)    B = np.random.randn(K, N).astype(np.float32)    C = np.random.randn(M, N).astype(np.float32)     ref = A @ B + C    schemes = {        "fp16 (Tensor Core)": to_fp16,        "tf32  (Tensor Core)": to_tf32,        "fp32  (CUDA Core)": to_fp32,    }    for name, cast in schemes.items():        out = mma_accumulate(A, B, C, cast)        err = np.abs(out - ref).mean()        print(f"{name:22s} MAE = {err:.6e}")  if __name__ == "__main__":    run_demo()

1. def to_fp16(x)把 FP32 张量先降为 FP16 再升回 FP32，模拟 Tensor Core 以 FP16 读取输入、以 FP32 累加的过程。
2. def to_tf32(x)通过位运算把 FP32 的尾数低 13 位清零，等价于 TF32 的 10 位尾数截断。
3. bits & np.uint32(0xFFFFE000)0xFFFFE000 的二进制低 13 位全为 0，按位与即丢弃低 13 位尾数。
4. def mma_accumulate(A, B, C, cast_fn)用 Aq @ Bq + C 一句话复现 mma 指令的数学语义 D = A×B+C。
5. schemes = { ... }用字典把三种精度策略映射到对应的舍入函数，便于统一循环对比。
6. err = np.abs(out - ref).mean()计算平均绝对误差（MAE），量化每种精度相对 FP32 基准的偏差。

常见误区

面试题