发表更新16 分钟读完 (大约2452个字)

NVIDIA PTX 简单入门

先看代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64"
target triple = "nvptx64-nvidia-cuda"

declare i32 @llvm.nvvm.read.ptx.sreg.tid.x() nounwind readnone

define void @kernel(ptr addrspace(1) %A,
                    ptr addrspace(1) %B,
                    ptr addrspace(1) %C) {
entry:
  %id = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
  %ptrA = getelementptr inbounds float, ptr addrspace(1) %A, i32 %id
  %ptrB = getelementptr inbounds float, ptr addrspace(1) %B, i32 %id
  %ptrC = getelementptr inbounds float, ptr addrspace(1) %C, i32 %id

  %valA = load float, ptr addrspace(1) %ptrA, align 4
  %valB = load float, ptr addrspace(1) %ptrB, align 4
  %valC = fadd float %valA, %valB

  store float %valC, ptr addrspace(1) %ptrC, align 4
  ret void
}

!nvvm.annotations = !{!0}
!0 = !{ptr @kernel, !"kernel", i32 1}

这个kernel.ll文件是LLVM IR（Intermediate Representation）表示的CUDA核函数，专为NVIDIA PTX（Parallel Thread Execution）架构生成。以下是逐行解释：

逐行解释

  1. 目标架构定义
1
target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64"

作用：定义数据的内存布局规则。

关键参数：

  • e: 小端字节序
  • p:64:64:64: 指针占64位，对齐为64位
  • i1:8:8: 1位整数对齐为8位
  • f32:32:32: 单精度浮点数对齐为32位
  • v16:16:16: 16位向量对齐为16位
  • n16:32:64: 本地指针大小（用于GPU架构）
1
target triple = "nvptx64-nvidia-cuda"

作用：指定目标架构为NVIDIA PTX 64位架构，用于CUDA设备。

  1. 内置函数声明
1
declare i32 @llvm.nvvm.read.ptx.sreg.tid.x() nounwind readnone

作用：声明PTX内置函数，用于读取线程在X维度的ID（类似CUDA的threadIdx.x）。

特性：

  • nounwind: 保证不会抛出异常;
  • readnone: 函数无副作用，输出仅依赖输入;
  1. 核函数定义
1
2
3
define void @kernel(ptr addrspace(1) %A,
                    ptr addrspace(1) %B,
                    ptr addrspace(1) %C)

作用：定义名为kernel的核函数，接受三个全局内存指针参数。

关键点：

  • addrspace(1): 表示指针指向全局内存（GPU显存）
  • 对比CUDA中的参数类型：float* A → ptr addrspace(1) %A
  1. 入口基本块（Entry Block）
1
2
entry:
  %id = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()

作用：获取当前线程在X维度的ID，存入%id寄存器。

1
2
3
%ptrA = getelementptr inbounds float, ptr addrspace(1) %A, i32 %id
%ptrB = getelementptr inbounds float, ptr addrspace(1) %B, i32 %id
%ptrC = getelementptr inbounds float, ptr addrspace(1) %C, i32 %id

作用：计算每个线程访问的全局内存地址。

语法：

  • getelementptr inbounds: 计算数组元素地址（类似&A[threadIdx.x]）
  • float: 元素类型
  • ptr addrspace(1) %A: 基地址
  • i32 %id: 偏移量
1
2
%valA = load float, ptr addrspace(1) %ptrA, align 4
%valB = load float, ptr addrspace(1) %ptrB, align 4

作用：从全局内存加载数据到寄存器。

参数：

  • align 4: 确保4字节对齐访问（优化内存访问）
1
%valC = fadd float %valA, %valB

作用：执行浮点数加法运算。

1
2
store float %valC, ptr addrspace(1) %ptrC, align 4
ret void

作用：将结果写回全局内存，然后返回。

关键点：

  • store指令的内存地址必须指定正确的地址空间（addrspace(1)）
  1. 元数据注解
1
2
!nvvm.annotations = !{!0}
!0 = !{ptr @kernel, !"kernel", i32 1}

作用：标记@kernel函数为CUDA核函数。

参数：

  • ptr @kernel: 函数指针;
  • !”kernel”: 标记类型为核函数;
  • i32 1: 参数（通常表示核函数版本或特性）;

这个核函数的行为等价于以下CUDA代码：

1
2
3
4
__global__ void kernel(float*A, float* B, float* C) {
    int id = threadIdx.x;
    C[id] = A[id] + B[id];
}

每个线程负责将全局内存中A和B对应位置的值相加，结果写入C的相同位置。

关键概念补充

  1. 地址空间：

  • addrspace(1): 全局内存（显存）

  • addrspace(3): 共享内存

  • addrspace(5): 常量内存

  1. PTX寄存器访问：

  • tid.x对应threadIdx.x

  • 类似还有ctaid.x（blockIdx.x）、ntid.x（blockDim.x）等

  1. LLVM IR特性：

  • 静态单赋值（SSA）形式

  • 强类型系统

  • 显式内存地址空间管理

再看用法

将上文的LLVM IR代码编译为NVIDIA PTX 后端代码，命令为：

1
llc-15 -march=nvptx64 -mcpu=sm_80 kernel.ll -o kernel.ptx

  1. llc-15
    作用：LLVM 静态编译器工具,用于将 LLVM IR 代码编译为目标平台的汇编代码或二进制格式。此处用于生成 NVIDIA PTX 代码。

  2. -march=nvptx64
    作用：指定目标架构为 NVIDIA PTX 64 位，nvptx64 是 NVIDIA 的 Parallel Thread Execution (PTX) 虚拟指令集架构，专为 64 位 GPU 设计。PTX 代码可在支持该架构的 NVIDIA GPU 上运行（需进一步编译为实际 GPU 指令）。

  3. -mcpu=sm_80
    作用：指定目标 GPU 的计算能力版本。

  • sm_80 对应 NVIDIA Ampere 架构（如 A100、RTX 3090 等）。
  • sm 表示 “Streaming Multiprocessor”，数字 80 代表计算能力 8.0。
    此选项确保生成的 PTX 代码针对该架构优化。

  1. kernel.ll
    作用: 此文件通常由 Clang 或其他 LLVM 前端生成，包含高级语言（如 C/C++、CUDA 等）编译后的中间代码。

  2. -o kernel.ptx
    作用：输出的 PTX 文件是 NVIDIA GPU 可读的中间代码，后续可通过 NVIDIA 驱动或工具（如 nvcc）进一步编译为实际 GPU 指令（SASS，见下文）。

PTX 生成代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
//
// Generated by LLVM NVPTX Back-End
//

.version 7.0
.target sm_80
.address_size 64

 // .globl kernel                  // -- Begin function kernel
                                        // @kernel
.visible .entry kernel(
 .param .u64 kernel_param_0,
 .param .u64 kernel_param_1,
 .param .u64 kernel_param_2
)
{
 .reg .b32  %r<2>;
 .reg .f32  %f<4>;
 .reg .b64  %rd<8>;

// %bb.0:                               // %entry
 ld.param.u64  %rd1, [kernel_param_0];
 ld.param.u64  %rd2, [kernel_param_1];
 mov.u32  %r1, %tid.x;
 ld.param.u64  %rd3, [kernel_param_2];
 mul.wide.s32  %rd4, %r1, 4;
 add.s64  %rd5, %rd1, %rd4;
 add.s64  %rd6, %rd2, %rd4;
 add.s64  %rd7, %rd3, %rd4;
 ld.global.f32  %f1, [%rd5];
 ld.global.f32  %f2, [%rd6];
 add.rn.f32  %f3, %f1, %f2;
 st.global.f32  [%rd7], %f3;
 ret;
                                        // -- End function
}

SASS

SASS（Streaming ASSembly）是 NVIDIA GPU 的实际底层机器码（二进制指令集），直接由 GPU 硬件执行。它是 PTX 代码经过进一步编译后的最终产物，与 PTX 的关系类似于 CPU 汇编代码和中间语言（如 Java 字节码）的关系。

SASS 的关键特性

  • 二进制格式：

    • SASS 是 GPU 硬件直接执行的二进制指令，不可读
    • PTX 是文本格式的中间代码，人类可读（但一般由编译器生成）

  • 硬件绑定：

    • SASS 直接对应具体 GPU 架构（如 Ampere、Ada Lovelace、Hopper）。
    • 不同架构的 SASS 不兼容（例如 sm_80 的 SASS 无法在 sm_70 的 GPU 上运行）。

  • 性能优化

    • SASS 经过 NVIDIA 驱动或工具（如 nvcc）的优化，包含特定 GPU 的指令调度、寄存器分配等。
    • PTX 是通用中间表示，需进一步编译为 SASS 才能高效执行。

  • 隐蔽性：

    • NVIDIA 未公开 SASS 的完整指令集和编码规范，普通开发者通常无需直接操作 SASS。

PTX 与 SASS 的编译流程

1
2
3
4
5
6
7
高级语言（如 CUDA C++）
  ↓ 编译（nvcc/clang）
LLVM IR（.ll 文件）
  ↓ 编译（llc）
PTX 代码（.ptx 文件）
  ↓ 运行时编译（NVIDIA 驱动）或离线编译（nvcc）
SASS 机器码（二进制，GPU 直接执行）
  • PTX 是跨 GPU 架构的中间层代码（类似虚拟指令集）;
  • SASS 是最终在 GPU 上运行的机器码（绑定具体架构）;

如何查看 SASS 代码？

  1. 使用 cuobjdump：
1
cuobjdump -sass compiled_gpu_binary.cubin

输出 SASS 指令的文本表示（例如指令操作码、寄存器分配）。

  1. 使用 nvcc 生成：
1
nvcc --keep --gpu-architecture=sm_80 -c code.cu

保留中间文件（如 .sass 或 .cubin）。

  1. Nsight Compute：

NVIDIA 官方工具，可分析 SASS 指令的执行效率和资源使用。

假设一段简单的加法操作，PTX 代码：

1
2
3
4
5
6
7
8
9
10
11
12
.version 7.8
.target sm_80
.global .func(.param.b64 %out), add(
  .param.b64 %a,
  .param.b64 %b
) {
  ld.param.u64 %r1, [%a];
  ld.param.u64 %r2, [%b];
  add.u64 %r3, %r1, %r2;
  st.param.b64 [%out], %r3;
  ret;
}

对应的SASS 代码为：

1
2
IADD R1, R2, R3;
MOV R4, R1;

PTX 怎么用

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
#include <iostream>
#include <fstream>
#include <cassert>
#include <cuda.h>


void checkCudaErrors(CUresult err) {
  assert(err == CUDA_SUCCESS);
}

/// main - Program entry point
int main(int argc, char **argv) {
  CUdevice    device;
  CUmodule    cudaModule;
  CUcontext   context;
  CUfunction  function;
  CUlinkState linker;
  int         devCount;

  // CUDA initialization
  checkCudaErrors(cuInit(0));
  checkCudaErrors(cuDeviceGetCount(&devCount));
  checkCudaErrors(cuDeviceGet(&device, 0));

  char name[128];
  checkCudaErrors(cuDeviceGetName(name, 128, device));
  std::cout << "Using CUDA Device [0]: " << name << "\n";

  int devMajor, devMinor;
  checkCudaErrors(cuDeviceComputeCapability(&devMajor, &devMinor, device));
  std::cout << "Device Compute Capability: "
            << devMajor << "." << devMinor << "\n";
  if (devMajor < 2) {
    std::cerr << "ERROR: Device 0 is not SM 2.0 or greater\n";
    return 1;
  }

  std::ifstream t("kernel.ptx");
  if (!t.is_open()) {
    std::cerr << "kernel.ptx not found\n";
    return 1;
  }
  std::string str((std::istreambuf_iterator<char>(t)),
                    std::istreambuf_iterator<char>());

  // Create driver context
  checkCudaErrors(cuCtxCreate(&context, 0, device));

  // Create module for object
  checkCudaErrors(cuModuleLoadDataEx(&cudaModule, str.c_str(), 0, 0, 0));

  // Get kernel function
  checkCudaErrors(cuModuleGetFunction(&function, cudaModule, "kernel"));

  // Device data
  CUdeviceptr devBufferA;
  CUdeviceptr devBufferB;
  CUdeviceptr devBufferC;

  checkCudaErrors(cuMemAlloc(&devBufferA, sizeof(float)*16));
  checkCudaErrors(cuMemAlloc(&devBufferB, sizeof(float)*16));
  checkCudaErrors(cuMemAlloc(&devBufferC, sizeof(float)*16));

  float* hostA = new float[16];
  float* hostB = new float[16];
  float* hostC = new float[16];

  // Populate input
  for (unsigned i = 0; i != 16; ++i) {
    hostA[i] = (float)i;
    hostB[i] = (float)(2*i);
    hostC[i] = 0.0f;
  }

  checkCudaErrors(cuMemcpyHtoD(devBufferA, &hostA[0], sizeof(float)*16));
  checkCudaErrors(cuMemcpyHtoD(devBufferB, &hostB[0], sizeof(float)*16));


  unsigned blockSizeX = 16;
  unsigned blockSizeY = 1;
  unsigned blockSizeZ = 1;
  unsigned gridSizeX  = 1;
  unsigned gridSizeY  = 1;
  unsigned gridSizeZ  = 1;

  // Kernel parameters
  void *KernelParams[] = { &devBufferA, &devBufferB, &devBufferC };

  std::cout << "Launching kernel\n";

  // Kernel launch
  checkCudaErrors(cuLaunchKernel(function, gridSizeX, gridSizeY, gridSizeZ,
                                 blockSizeX, blockSizeY, blockSizeZ,
                                 0, NULL, KernelParams, NULL));

  // Retrieve device data
  checkCudaErrors(cuMemcpyDtoH(&hostC[0], devBufferC, sizeof(float)*16));


  std::cout << "Results:\n";
  for (unsigned i = 0; i != 16; ++i) {
    std::cout << hostA[i] << " + " << hostB[i] << " = " << hostC[i] << "\n";
  }


  // Clean up after ourselves
  delete [] hostA;
  delete [] hostB;
  delete [] hostC;

  // Clean-up
  checkCudaErrors(cuMemFree(devBufferA));
  checkCudaErrors(cuMemFree(devBufferB));
  checkCudaErrors(cuMemFree(devBufferC));
  checkCudaErrors(cuModuleUnload(cudaModule));
  checkCudaErrors(cuCtxDestroy(context));

  return 0;
}

这里是代码的简单行为：

  • 初始化 CUDA 环境：

    • 检测 GPU 设备，验证计算能力；
    • 加载 kernel.ptx 文件。

  • 准备数据：

    • 在 GPU 上分配内存（devBufferA, devBufferB, devBufferC）；
    • 在 CPU 上初始化输入数据 hostA 和 hostB（值分别为 0, 1, 2, …, 15 和 0, 2, 4, …, 30）。
    • 将数据从 CPU 拷贝到 GPU。

  • 启动 GPU 核函数：

    • 调用 kernel 函数，执行 A + B = C 的加法操作；
    • 核函数使用 1 个线程块，包含 16 个线程（每个线程处理一个元素）。

  • 验证结果：

    • 将 GPU 计算结果 devBufferC 拷贝回 CPU 的 hostC；
    • 打印 hostA[i] + hostB[i] = hostC[i]，验证结果是否正确。

  • 清理资源

相关输出为：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
Using CUDA Device [0]: NVIDIA H100 80GB HBM3
Device Compute Capability: 9.0
Launching kernel
Results:
0 + 0 = 0
1 + 2 = 3
2 + 4 = 6
3 + 6 = 9
4 + 8 = 12
5 + 10 = 15
6 + 12 = 18
7 + 14 = 21
8 + 16 = 24
9 + 18 = 27
10 + 20 = 30
11 + 22 = 33
12 + 24 = 36
13 + 26 = 39
14 + 28 = 42
15 + 30 = 45

NVIDIA PTX 简单入门

https://devillove084.github.io/2025/02/01/PTX-1/

作者

devillove084

发布于

2025-02-01

更新于

2025-02-16

许可协议

#

喜欢这篇文章？打赏一下作者吧

评论

Your browser is out-of-date!

Update your browser to view this website correctly.&npsb;Update my browser now

×

×