Nvidia GPU架构与并行计算分享

学习笔记

Posted by donpromax on September 29, 2025

CPU & GPU

计算机的基本组成结构

计算机的基本组成单元(冯诺依曼结构):

  • 存储器
  • 控制单元
  • 算数逻辑单元
  • 输入输出

读写:由RAM(Random Access Memory)存储指令和数据;

控制单元:读取指令/数据,解码指令并串行地执行;

算数逻辑单元:基本的运算单元;

输入/输出:人机交互的输入输出接口。

What is CPU?

CPU(Central Processing Unit)是计算机系统的重要组成部分,经常被比喻为计算机的“大脑”。CPU的主要职责是执行位于存储器(Memory)中的指令(Instructions),进行算数、逻辑运算,并管理计算机中不同不同部分之间的数据交换。

CPU的性能瓶颈

CPU的性能评价指标经常是时钟频率(Clock speed),表示CPU每秒可以执行的指令数;以及核心数(number of cores),表示CPU包含几个独立的计算单元。

早期CPU的发展以提高时钟频率为主,但这种趋势在2004~2005年就得以减缓,主要原因是由功耗限制Power Wall和访存限制Memory wall两大阻碍。

CPU的时钟频率的提高也同时带来了更大的功耗,意味着需要更大的电源,以及可能面临的运行时高温度问题。另一方面,受限于存储技术的发展限制,在很长一段时间内计算机程序的主要性能瓶颈并不在计算上,而是在访存上。

因此,CPU从不断的提高时钟频率转而向提高核心数的方向发展,例如现在常见的都是4核、8核CPU。CPU核心更擅长于串行任务的计算。

What is GPU?

GPU (Graphics Processing Units,图形处理单元,也称为显卡) ,从命名可能看出来最早是为了加速图形处理(特别是电子游戏)而设计的,后来也被广泛运用于其他并行计算领域。

与CPU的设计理念不同,GPU主要被设计用于处理数以千计的重复性任务。以图形处理为例,可以把显示器中的每个像素点的计算看成是一个任务,每个像素点之间相对独立互不干扰。GPU可以高效的并行处理每个像素点运算。

GPU的特性以及优势

GPU(Throughput-Oriented Design)相较于CPU(Latency-Oriented Design)的最显著的优势就是拥有更高的峰值访存与计算性能。如果能利用这一特性,可以在GPU上开发运行某些应用时可以较CPU带来显著的性能提升。例如:生成式AI、数字信号处理、视频图像处理、生物信息学、网络信息搜索、气象预报、加密系统等领域。

造成GPU与CPU的峰值访存与计算性能差距的最主要原因是两者的设计理念区别,它们被设计用于处理不同类型的任务。

CPU被设计得更擅长处理复杂的任务序列,一个程序也可能包含数十个线程(Threads)的并行。其中每个线程可能运行于不同的核心(Cores)上,CPU的每个核心都很强大,因此每个线程的任务相对复杂。

而GPU是被设计用于并行处理成千上万个任务的,GPU处理每个任务的核心相较CPU更弱,但胜在数量。

从下图可以看到,GPU芯片的更多晶体管都用于数据处理核心;而CPU则有更多的晶体管用于缓存(Cache)和控制单元(Control unit)。

CUDA C++ Programming Guide

使用更多的晶体管用于数据处理核心(例如浮点数计算单元)也就带来了更高的浮点计算性能(FLOPs, floating-point operations per second)。

另一方面,GPU可以使用延迟隐藏(Latency hiding)等手段来降低访存带来的等待。

总的来说,通常的应用都是串行计算和并行计算的混合。因此计算机系统被设计为由CPU与GPU的组合来最大化性能表现。

Flynn矩阵

一种并行计算机的分类方法,根据指令流和数据流的不同状态来分类。指令流(Instruction stream)与数据流(Data stream)各自可以是单(Single)和多(Multiple)两种状态,因此可以形成一个2×2的矩阵。

SISD:非并行计算机,一种古老的单核计算机架构:

SIMD: 最适合用于处理任务之间具有高度相似规则的情形,例如图形处理器(GPUs):

MISD:实际应用中较少见,一种理论上用途可能是对单个信号应用不同滤波函数:

MIMD:现代多核CPU的计算机最常见结构:

Amdahl’s law

阿姆达定律指出了程序的加速比取决于程序中可以并行的部分(P)的占比:

\[\text{speedup} = \frac{1}{P/N+S}\]

P = parallel fraction, N = number of processors and S = serial fraction.

                            speedup
              -------------------------------------
        N     P = .50   P = .90   P = .95   P = .99
      -----   -------   -------   -------   -------
         10      1.82      5.26      6.89      9.17
        100      1.98      9.17     16.80     50.25    
      1,000      1.99      9.91     19.62     90.99
     10,000      1.99      9.91     19.96     99.02
    100,000      1.99      9.99     19.99     99.90

You can spend a lifetime getting 95% of your code to be parallel, and never achieve better than 20x speedup no matter how many processors you throw at it!

一个著名的quote:你可以用一辈子的时间来把代码的95%并行化,但是无论你加多少处理器,都不可能获得高于20x的加速比。

GPU计算单元结构

“学校模型”

在介绍GPU的物理结构和逻辑结构前,可以用一种直观的比喻来帮助理解,“学校模型”

例子:一个学校中有多个教室,每个教室中有多个学生。接下来,我们有若干个任务(具体数量可以根据情况有所不同)需要分配给学校的学生来处理。只不过需要遵循一定的规则:

  1. 每个教室最多只能处理1024个任务(每个班级最多只有1024个学生);
  2. 学生按照32个一组处理Task(线程按照Warps为单位进行调度)。

物理结构

http://users.umiacs.umd.edu/~ramani/cmsc828e_gpusci/lecture9.pdf

Let’s take Kepler or Fermi for example:

从上面的示意图可以看到,该GPU有十六个SMs,并且每个SM有32个Cores。

流处理器 Streaming Processors(SPs or cores):SPs是GPU中的主要处理单元,可以用于并行地在不同的数据上做计算。SP可以对应于“学校模型”中的课桌,每个SP都是用于处理一个个的任务。因此可以认为,有越多的SPs,就可以越多得并行处理任务。

流多处理器 Streaming Multiprocessor(SM or multiprocessor):SM是多个SP的组合。SM可以认为是“学校模型”中的教室。SM作为一个更高级别的SP集合,可以统一管理内部所有SP的任务处理调度。

具体到不同的GPU型号,SM和SP的数量各自不同。以Nvidia T4为例,T4总共有40个SMs,每个SM包含64个Streaming Processors。每个GPU的具体参数细节可以查阅Nvidia官方网站获得。

An illustration of a streaming multiprocessor of RTX4090

以Nvidia RTX4090为例,上图是一个SM的构成示例图,可以看到除了SPs(Cores or Processors),还包含几个其他的组成部分:

Caches: L1 Cache/Shared Memory、Texture cache等;

Schedulers for warps:以Warps为单独调度Threads。以“学校模型”为例,可以理解为将班级划分为若干个小组。每个小组由Scheduler for warps调度。

Registers:供Thread使用,提供最快的访存速度。

Total view of RTX4090

逻辑结构

https://en.wikipedia.org/wiki/Thread_block_(CUDA_programming>)

可以用下图来理解CUDA的编程结构:

从上面的示意图可以看到此逻辑结构有2×3的6块blocks组成的grid,其中每个block由3×4个Threads组成。

Thread:一个任务的执行单元,每个Thread负责处理一个Task,可以对应于“学校模型”中的学生(苦力)。

Block:若干个Threads组成的集合,可以是一维、二维或者三维的排布。可以对应于“学校模型”中的班级。单个Block中的Threads数量存在上限,通常是最多可以包含1024个线程,Max dimension size of a thread block (x,y,z): (1024, 1024, 64). For example, a block with dimensions of (32,32,1) will have 1024 threads。

Grid:若干个Blocks组成的集合,可以是一维、二维或者三维的排布。可以对应于“学校模型”中的年级。Max dimension size of a grid size (x,y,z): (2147483647, 65535, 65535)。

Warps:每个Block中的Thread将会被分为若干个Warps(对于CUDA编程者来说不可感知)。可以对应于“学校模型”中将班级划分为若干个小组。每个Warp包含连续的32个Threads,Warp内的32个线程执行相同的指令(Warp内的线程是同步的)。

https://slideplayer.com/slide/16690299/


关于Warp进阶阅读可以参考

https://developer.nvidia.com/blog/using-cuda-warp-level-primitives/

问:以Nvidia G80为例。如果某个SM被分配了3个Thread Blocks,每个Block包含256个Threads,该SM中有多少个Warps?

答:每个Block被划分为 256/32 = 8 Warps。共有 8 * 3 = 24 Warps。在任意时刻,这24个Warps中只有一个Warp被Scheduler调度运行。


问:一个程序包含,1次对global memory的读操作(200 cycles)和4次独立的mutiply/add操作(每个4 cycles)。warps的context switch是零开销的,那么需要多少个warps才可以完全隐藏访存延迟?

答:每个warp有4个mutiply/add操作,4×4=16 个周期。我们需要覆盖200个周期,200/16=12.5,ceil(12.5)=13,至少需要13个warps可以完全隐藏访存延迟。


GPU的物理结构与逻辑结果之间存在对应关系:

GPU的存储结构

GPU的逻辑存储结构

View of GPU memory

局部内存(Local Memory):每个线程可以使用其自己的局部内存,用于存储临时变量。具有最小的作用域,并且是专门为每个单独的线程分配的。“学校模型”中的课桌,每个学生可以访问自己的课桌。

共享内存(Shared Memory):同一线程块内的线程可以通过共享内存共享数据。这允许同一线程块内的线程之间更快地通信和访问数据,相比访问全局内存来说速度更快。“学校模型”中的讲台,每个班级都有一个讲台。

全局内存(Global Memory):这是GPU中最大的内存,可以被所有线程跨所有线程块访问。然而,访问全局内存通常比其他内存类型慢,因此需要进行优化以避免性能下降。“学校模型”中的操场,学校有一个操场。

纹理内存和常量内存(Texture Memory and Constant Memory):这些是GPU中的特殊内存类型,针对访问特定数据类型(如纹理或常量值)进行了优化。所有线程跨所有线程块都可以访问这些内存类型。

物理访存过程

不同的Memory之间的带宽有着显著的区别。CPU (host) and GPU (device)之间通过PCIe连接,速度最慢(Host2DeviceorDevice2Host,~32 GB/s)。

Device Memory(显存)其次。然后是L2 Cache、L1 Cache/Shared Memory,最快的是Registers。

Logical GPU memory attributes

Physical GPU memory attributes(outdated)

Global Memory

当某个CUDA kernel尝试访问Global Memory(Logical Memory Space)时,实际的的过程如下:

  1. 尝试从L1 Cache中查找,if hit then return;
  2. (If L1 Cache miss,存在latency penalty)尝试从L2 Cache中查找,if hit then return;
  3. (If L2 Cache miss,存在latency penalty)尝试从Device Memory(显存)查找。

Global Memory Access Process

Shared Memory

是一个较小的物理空间。可以被同一个Block中的所有Threads所共享。在物理空间上,Shared Memory与L1 Cache位于同一片上Memory。CUDA编程者甚至可以自由地分配L1 Cache与Shared Memory的比例。

cudaFuncSetAttribute(
    my_kernel,
    cudaFuncAttributePreferredSharedMemoryCarveout,
    20 // Use 20% of combined L1/Shared Memory for Shared Memory
);

为了利用Shared Memory来加速访存,在CUDA编程中,使用__shared__关键字修饰的变量为Shared Memory。如果某个地址空间在多个线程中需要反复多次使用,使用Shared Memory可以显著提升访存效率。因为可以避免每次从Global Memory中读取(可能存在 cache missing 带来 latency penalty)。

Local Memory与Registers

在CUDA kernel中声明的标量(scalar-type)局部变量(int、float),默认都会存储在寄存器(Registers)中。

以Nvidia A100为例,每个Thread最多可以有255个Registers,超出该部分的局部变量(寄存器耗尽)或者是数组等动态大小的类型会被存储到Local Memory中,Local Memory与Glocal Memory同样,是一个包含多级缓存(L1 Cache -> L2 Cache -> Device Memory)的逻辑存储空间。

使用Nvidia Nsight Compute可以直观的看到CUDA代码到PTX汇编中局部变量的内存空间分配的过程:

常量局部变量被定义为寄存器变量

数组局部变量被定义为local memory

Some Official Specs

Graphics Card NVIDIA GeForce RTX 2080 NVIDIA GeForce RTX 3090 NVIDIA GeForce RTX 4090
GPU Architecture NVIDIA Turing NVIDIA Ampere NVIDIA Ada Lovelace
SMs 46 82 128
CUDA Cores / SM 64 128 192
CUDA Cores / GPU 2944 10496 24,576
Tensor Cores / SM 8 (2nd Gen) 4 (3rd Gen) 6 (3rd Gen)
Tensor Cores / GPU 368 (2nd Gen) 328 (3rd Gen) 768 (3rd Gen)
RT Cores 46 (1st Gen) 82 (2nd Gen) 128 (3rd Gen)
Memory Bandwidth 448 GB/sec 936 GB/sec 1008 GB/sec
L1 Data Cache/Shared Memory 4,416 KB 10,496 KB 16,384 KB
L2 Cache Size 4096 KB 6144 KB 73,728 KB
Register File Size 11,776 KB 20,992 KB 32,768 KB
TGP (Total Graphics Power) 225W 350W 450W
Transistor Count 13.6 Billion 28.3 Billion 76.3 Billion

CUDA编程

Why CUDA?

CUDA In Large Language Models

CUDA对于Generative AI的训练推理加速起着至关重要的作用,例如对于Qwen2的一次前向传播过程,虽然经过了Torch、 vLLM等框架的包装,但底层还是调用了CUDA来进行并行计算加速:

AI模型推理过程包含的一个个CUDA kernel的调用

学习CUDA编程可以做什么?

  • 针对特定场景的并行计算的加速,例如数字信号处理、视频编解码等领域;
  • 特别的,For AI,可以手写CUDA kernel并集成到PyTorch、vLLM等框架中使用,以实现训练、推理加速等目的;

What is CUDA?

Nvidia GPU上的并行计算编程涉及到CPU与GPU之间的数据传输、处理、计算,支持C、C++等语言。

CUDA (Compute Unified Device Architecture)是Nvidia开发的一种并行计算开发框架,利用CUDA编程者可以利用GPU的硬件特性来进行并行计算的开发。

CUDA Compiler-NVCC

使用CUDA C/C++的文件的后缀为.cu,我们需要一个可以同时支持C/C++和CUDA的编译器。Nvidia开发了NVCC(NVIDIA CUDA Compiler),可以将C/C++和CUDA代码分别编译为在CPU以及GPU运行的二进制文件。NVCC的主要功能包括:

  • 代码划分:区分.cu代码中哪些部分是CPU的代码,哪些部分是GPU的代码;
  • 编译优化:Host代码(CPU)使用传统的C/C++编译器,而Device代码(GPU)使用CUDA编译器;
  • 生成PTX:生成PTX代码 (PTX - Parallel Thread Execution),是一种运行在GPU上代码的中间表示;
  • PTX翻译优化:将PTX代码翻译为GPU-specific machine code (SASS - Scalable Assembly) ,会根据硬件架构做专门的优化;
  • Linking:将编译完成的CPU、GPU代码链接到一起生成最终的可执行二进制文件。

Hello World!

Getting Started

  1. 根据官网的教程安装

Cuda Toolkit

  1. 如果你的电脑没有Nvidia GPU,可以使用Google Colab,教程

in here

确保CUDA被正确地添加到环境变量中,在终端中输入nvcc -V可以验证是否安装成功

Code Example

#include <stdio.h>


__global__ void kernel()
{

    printf("hello world");
}

int main()
{
    kernel<<<1,1>>>();
    cudaDeviceSynchronize();

    return 0;
}

__ host __: 调用运行于CPU上的代码块,如果一个函数不含任何CUDA修饰符,则默认为host代码块,例如上面的main()函数。<span style="color:rgb(31, 35, 40);">int main()</span><span style="color:rgb(31, 35, 40);">__host__ int main()</span>是等效的。

__device__调用运行于GPU上的代码块,

__ global __:被host代码块(CPU)调用,在device(GPU)上运行

«<1,1»>: 调用kernel时需要指定的维度信息,第一个1表示Blocks的数量,第二个1表示每个Block中的Threads数量。用“学校模型”来理解,你作为校长,对于这个任务指定了一个班级,每个班级使用一个学生来处理。

cudaDeviceSynchronize():cuda device线程同步语句,可以确保GPU上的kernel线程全部执行完毕后,再运行后面的CPU代码。另外,在kernel内部,还有__syncthreads()可以确保同一个Block内的Threads完成同步。


Run Hello World!

将上面的代码保存为 .cu 的格式,然后用nvcc来编译:

  • nvcc .cu (-o )

然后运行生成的二进制可执行文件:

  • ./a.out

Vector Add

Vector add

A CPU version of Vector Add

float *x, *y;

/* ==========================================
  Allocate arrays to store vector x and y
  ========================================== */
x = calloc( N, sizeof(float) );
y = calloc( N, sizeof(float) );

/* ===============================================
  CPU version of the vector addition algorithm
  =============================================== */
for (i = 0; i < N; i++)    // Add N elements...
  y[i] = x[i] + y[i];     // Add one element at a time...

i=0

y[0] = x[0] + y[0]

And so on….

Vector Add CUDA Kernel

// CUDA kernel to add two vectors
__global__ void vectorAdd(int *a, int *b, int *c) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid < N) {
        c[tid] = a[tid] + b[tid];
    }
}

__global__ void定义了该函数是一个由CPU调用,GPU执行的CUDA kernel。

其中blockIdx.x * blockDim.x + threadIdx.x用于计算当前线程处理Vector中第几个元素的索引。

if (tid < N)保证了Vector的长度N不能被BLOCK_SIZE整除时,最后一个Block的线程不会访问越界;


About 1-D indexing:Why blockIdx.x * blockDim.x + threadIdx.x

假设我们有4个Blocks,每个Block有8个Threads,来处理长度是32的Vector Add。那么Vector Index 21 位于第3个Block的第6个线程:


Vector Add Final Code

#include <stdio.h>
#include <stdlib.h>

// Size of the vector
#define N 32
#define BLOCK_SIZE 32

// CUDA kernel to add two vectors
__global__ void vectorAdd(int *a, int *b, int *c) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid < N) {
        c[tid] = a[tid] + b[tid];
    }
}

int main() {
    int *h_a, *h_b, *h_c; // Host vectors
    int *d_a, *d_b, *d_c; // Device vectors

    // Initialize host vectors
    h_a = (int *)malloc(N * sizeof(int));
    h_b = (int *)malloc(N * sizeof(int));
    h_c = (int *)malloc(N * sizeof(int));

    // Initialize host vectors with random values
    for (int i = 0; i < N; i++) {
        h_a[i] = rand() % 10;
        h_b[i] = rand() % 10;
    }

    // Allocate device memory for vectors
    cudaMalloc((void **)&d_a, N * sizeof(int));
    cudaMalloc((void **)&d_b, N * sizeof(int));
    cudaMalloc((void **)&d_c, N * sizeof(int));

    // Copy data from CPU to GPU
    cudaMemcpy(d_a, h_a, N * sizeof(int), cudaMemcpyHostToDevice);
    cudaMemcpy(d_b, h_b, N * sizeof(int), cudaMemcpyHostToDevice);

    // Call the CUDA kernel to perform vector addition
    int numBlocks = (N + BLOCK_SIZE - 1) / BLOCK_SIZE;
    vectorAdd<<<numBlocks, BLOCK_SIZE>>>(d_a, d_b, d_c);

    // Copy the result from GPU to CPU
    cudaMemcpy(h_c, d_c, N * sizeof(int), cudaMemcpyDeviceToHost);

    // Print the result
    for (int i = 0; i < N; i++) {
        printf("h_a[%d] %d + h_b[%d] %d = %d\n", i, h_a[i], i, h_b[i], h_c[i]);
    }

    // Free memory
    free(h_a);
    free(h_b);
    free(h_c);
    cudaFree(d_a);
    cudaFree(d_b);
    cudaFree(d_c);

    return 0;
}

Initialization of values on the GPU

cudaMalloc: 为device申请内存空间,与C中的malloc类似。你阅读其他文档或者代码时,可能看到cudaMalloc((void **)&d_a, N * sizeof(int))或者cudaMalloc(&d_a, N * sizeof(int))两种写法,这两种写法是等效的;

cudaFree:为申请好的device空间进行内存释放,与C中的free类似;

cudaMemcpy:在host和device之间进行数据传输,分为cudaMemcpyHostToDevice(CPU主存到显存)和cudaMemcpyDeviceToHost(显存到CPU主存)两种方向;


关于kernel调用的部分:

// Call the CUDA kernel to perform vector addition
int numBlocks = (N + BLOCK_SIZE - 1) / BLOCK_SIZE;
vectorAdd<<<numBlocks, BLOCK_SIZE>>>(d_a, d_b, d_c);

设BLOCK_SIZE为每个Block中的Threads数,那么可以用$ \lceil N / BLOCK_SIZE \rceil $来计算需要的Blocks数。(N + BLOCK_SIZE - 1) / BLOCK_SIZE实现了向上取整的除法。

Run Vector Add!

Run kernel with PyTorch

Install PyBind11

pip install "pybind11[global]"

我们使用pybind11来让我们的Cuda C/C++程序能够被链接到Python中使用。还是以向量相加为例,

我们现在希望写一个kernel,接受两个浮点类型的torch::Tensor,并返回他们相加的结果。

项目结构

project_root/
│
├── setup.py
│
└── csrc/
    ├── vector_add/
    │   ├── vector_add_binding.cpp
    │   └── vector_add.cu

vector_add.cu

之前我们之间使用main函数来调用CUDA kernel。而这次不同,我们使用了一个函数vector_add_launcher来调用CUDA kernel。该函数接受两个torch::Tensor类型作为输入,并返回他们的和。

  • 首先,我们校验了a和b的维度,确保他们是一维向量。如果dim=2则为矩阵,dim=3则为三维张量。我们定义的kernel只进行一维向量的加法;
  • 其次,校验了a和b的数据类型,确保他们是float32类型;
  • 使用torch::Tensor c = torch::empty_like(a);创建了返回值c,empty_like可以创建一个与a的形状、device等参数都相同,值未初始化的Tensor,用于保存结果;
  • 使用vector_add<<<blocksPerGrid, BLOCK_SIZE>>>(a.data_ptr<float>(), b.data_ptr<float>(), c.data_ptr<float>(), N)调用CUDA kernel,其中a.data_ptr<float>()返回float类型的数组指针,指向Tensor的首个元素。
#include <torch/extension.h>

#define BLOCK_SIZE 256


__global__ void vector_add(const float* a, const float* b, float* c, int N) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N) {
        c[i] = a[i] + b[i];
    }
}


torch::Tensor vector_add_launcher(torch::Tensor& a, torch::Tensor& b) {
    // Check the Tensor Dimension
    if (a.dim() != 1 || b.dim() != 1) {
        throw std::runtime_error("Input tensor must be one-dimensional");
    }
    // Check the Tensor Data Type
    if (a.scalar_type() != torch::kFloat32 || b.scalar_type() != torch::kFloat32) {
        throw std::runtime_error("Input tensor must be type of float32");
    }
    torch::Tensor c = torch::empty_like(a);
    auto N = a.sizes()[0]; // array length
    // Launch kernel
    int blocksPerGrid = (N + BLOCK_SIZE - 1) / BLOCK_SIZE;

    vector_add<<<blocksPerGrid, BLOCK_SIZE>>>(a.data_ptr<float>(), 
        b.data_ptr<float>(), c.data_ptr<float>(), N);

    return c;
}

vector_add_binding.cpp

然后我们创建一个 .cpp 文件vector_add_binding.cpp,其中包含了特殊的PYBIND11_MODULE声明,这样做可以允许我们使用Python代码来调用C++的代码。

#include <torch/extension.h>

torch::Tensor vector_add_launcher(torch::Tensor& a, torch::Tensor& b);

// Write the C++ function that we will call from Python
torch::Tensor vector_add_binding(torch::Tensor& a, torch::Tensor& b) {
    return vector_add_launcher(a, b);
}

PYBIND11_MODULE(example_kernels, m) {
  m.def(
    "vector_add", // Name of the Python function to create
    &vector_add_binding, // Corresponding C++ function to call
    "Launches the vector_add kernel" // Docstring
  );
}

对于C语言不熟悉的朋友,在这里特别解释一下,torch::Tensor vector_add_launcher(torch::Tensor& a, torch::Tensor& b);是函数的声明,编译完成后链接器会自动地在vector_add.cu中找到对应的函数实现。

在这里我们定义了函数vector_add_binding,这是Python代码调用C++函数的入口。在PYBIND11_MODULE中,我们定义了python模块example_kernels。并且定义了Python模块example_kernels的函数vector_add,对应于C++函数为vector_add_binding,并且可以在这里添加函数的说明。

setup.py

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

__version__ = "0.0.1"

# Define the C++ extension modules
ext_modules = [
    CUDAExtension('example_kernels', [
        'csrc/vector_add/vector_add_binding.cpp',
        'csrc/vector_add/vector_add.cu',
    ]),
]

setup(
    name="cuda_basics",
    version=__version__,
    ext_modules=ext_modules,
    cmdclass={"build_ext": BuildExtension}
)

最后我们创建setup.py。利用Python的setuptools,这样每次我们在项目目录运行pip install .时,都会生成可以import的Python模块example_kernels。我们在项目目录下依次import torchimport example_kernels后测试一下。使用torch.tensor创建a、b两个一维的向量,指定类型为torch.float32,然后使用example_kernels.vector_add(a, b)调用我们的手写CUDA C++ kernel:

总结