PTX 技术详解

1. GPU 架构基础与 PTX 概述

1.1 GPU 与 CPU 架构对比

现代异构计算系统中，GPU 和 CPU 承担着本质不同的计算任务。理解两者的架构差异是掌握 GPU 编程的基础。

核心洞察：GPU 的性能优势来源于 用面积换并行度。相同芯片面积下，GPU 选择堆叠大量简单核心而非少量复杂核心，这种设计在数据并行场景下能获得数量级的性能提升。

1.2 SIMT 架构原理

SIMT（Single Instruction Multiple Threads）是 NVIDIA GPU 的核心执行模型，它在 SIMD 基础上引入了线程抽象，提供更灵活的编程模型。

flowchart TB
    subgraph SM["Streaming Multiprocessor (SM)"]
        subgraph WS["Warp Scheduler"]
            W1[Warp 0]
            W2[Warp 1]
            W3[Warp N]
        end
        
        subgraph EXEC["Execution Units"]
            ALU1[ALU x32]
            ALU2[ALU x32]
            SFU[SFU]
            LSU[Load/Store]
        end
        
        subgraph MEM["Memory"]
            RF[Register File<br/>64K x 32-bit]
            SMEM[Shared Memory<br/>64-164 KB]
            L1[L1 Cache]
        end
    end
    
    WS --> EXEC
    EXEC <--> MEM

1.2.1 SIMT vs SIMD 关键区别

1.2.2 Warp 执行机制

Warp 是 GPU 调度的基本单位，理解其执行机制对性能优化至关重要：

Warp (32 threads)
├── Thread 0-31 共享同一条指令
├── 每个线程有独立的：
│   ├── 寄存器状态（最多 255 个 32-bit 寄存器）
│   ├── 程序计数器（Volta+ 独立 PC）
│   └── 本地内存空间
└── 分支发散时：
    ├── 硬件序列化执行不同分支
    └── 通过 Predication 或 Re-convergence 处理

// 糟糕的写法：50% 线程执行 if，50% 执行 else
if (threadIdx.x % 2 == 0) {
    // 分支 A
} else {
    // 分支 B  
}
// 实际执行时间 = 分支A时间 + 分支B时间

// 优化写法：按 Warp 边界划分
if (threadIdx.x < 32) {
    // Warp 0 全部执行此分支
} else {
    // Warp 1 全部执行此分支
}

1.3 PTX 在 CUDA 生态中的地位

PTX（Parallel Thread Execution）是 NVIDIA 设计的虚拟 ISA，作为 CUDA 编译工具链的中间表示，它在源码与硬件之间提供了关键的抽象层。

flowchart TD
    subgraph 源码层
        CUDA[CUDA C/C++]
        FORT[CUDA Fortran]
        PY[Python/Numba]
    end
    
    subgraph 编译层
        NVCC[nvcc 前端]
        LLVM[NVVM/LLVM]
    end
    
    subgraph 中间表示
        PTX[PTX<br/>虚拟 ISA]
    end
    
    subgraph 后端
        PTXAS[ptxas 汇编器]
        JIT[Driver JIT]
    end
    
    subgraph 目标码
        SASS_SM80[SASS sm_80<br/>Ampere]
        SASS_SM89[SASS sm_89<br/>Ada]
        SASS_SM90[SASS sm_90<br/>Hopper]
    end
    
    CUDA --> NVCC
    FORT --> NVCC
    PY --> LLVM
    NVCC --> LLVM
    LLVM --> PTX
    PTX --> PTXAS
    PTX --> JIT
    PTXAS --> SASS_SM80
    PTXAS --> SASS_SM89
    JIT --> SASS_SM90

PTX 设计目标

1. 跨代兼容：PTX 代码可在未来 GPU 上通过 JIT 编译运行
2. 性能透明：提供足够底层的控制，不牺牲性能
3. 机器无关：屏蔽不同 GPU 架构的硬件差异
4. 可读性：类汇编语法，便于调试和手工优化

PTX vs SASS

# 使用 PTX 6.0 编译，可在任何 sm_70+ GPU 运行
nvcc -arch=sm_70 -ptx kernel.cu -o kernel.ptx

# 运行时 Driver 自动 JIT 编译为目标 GPU 的 SASS
# 无需重新编译源码即可支持新硬件

2. PTX 指令集架构

2.1 PTX 程序结构

一个典型的 PTX 程序包含以下组成部分：

// PTX 版本和目标架构声明
.version 7.8
.target sm_80
.address_size 64

// 全局变量声明
.global .align 4 .b8 data[1024];

// Kernel 函数定义
.visible .entry vector_add(
    .param .u64 param_a,
    .param .u64 param_b,
    .param .u64 param_c,
    .param .u32 param_n
)
{
    // 寄存器声明
    .reg .pred p0;
    .reg .b32 r<10>;
    .reg .b64 rd<5>;
    
    // 指令体
    ld.param.u64 rd0, [param_a];
    ld.param.u64 rd1, [param_b];
    ld.param.u64 rd2, [param_c];
    ld.param.u32 r0, [param_n];
    
    // 计算线程索引
    mov.u32 r1, %tid.x;
    mov.u32 r2, %ctaid.x;
    mov.u32 r3, %ntid.x;
    mad.lo.u32 r4, r2, r3, r1;  // idx = blockIdx.x * blockDim.x + threadIdx.x
    
    // 边界检查
    setp.ge.u32 p0, r4, r0;
    @p0 bra EXIT;
    
    // 地址计算和数据加载
    mul.wide.u32 rd3, r4, 4;
    add.u64 rd0, rd0, rd3;
    add.u64 rd1, rd1, rd3;
    add.u64 rd2, rd2, rd3;
    
    ld.global.f32 r5, [rd0];
    ld.global.f32 r6, [rd1];
    
    // 计算和存储
    add.f32 r7, r5, r6;
    st.global.f32 [rd2], r7;
    
EXIT:
    ret;
}

2.2 核心指令类别

2.3 内存层次与访问模式

flowchart TB
    subgraph 线程私有
        REG[寄存器<br/>~1 cycle]
        LOCAL[Local Memory<br/>~200 cycles]
    end
    
    subgraph Block 共享
        SMEM[Shared Memory<br/>~20 cycles]
    end
    
    subgraph 全局
        L1[L1 Cache]
        L2[L2 Cache<br/>~100 cycles]
        GMEM[Global Memory<br/>~400 cycles]
    end
    
    REG <--> SMEM
    SMEM <--> L1
    L1 <--> L2
    L2 <--> GMEM
    LOCAL --> L2

PTX 内存访问指令示例：

// 全局内存访问
ld.global.f32 r0, [addr];           // 普通加载
ld.global.ca.f32 r0, [addr];        // Cache all levels
ld.global.cg.f32 r0, [addr];        // Cache at L2 only
ld.global.nc.f32 r0, [addr];        // Non-coherent (texture path)

// 共享内存访问
ld.shared.f32 r0, [smem_addr];
st.shared.f32 [smem_addr], r0;

// 原子操作
atom.global.add.f32 r0, [addr], r1;
atom.shared.cas.b32 r0, [addr], r1, r2;

3. 编译与优化

3.1 编译流程详解

# 生成 PTX（保留中间文件）
nvcc -ptx -arch=sm_80 kernel.cu -o kernel.ptx

# 生成 SASS（cubin）
nvcc -cubin -arch=sm_80 kernel.cu -o kernel.cubin

# 同时嵌入 PTX 和 SASS（推荐发布方式）
nvcc -arch=sm_80 -code=sm_80,compute_80 kernel.cu

# 查看 SASS 反汇编
cuobjdump -sass kernel.cubin

3.2 性能优化要点

4. 参考资料

• PTX ISA Reference
• CUDA C++ Programming Guide
• NVIDIA GPU Architecture Whitepapers