2. 硬件架构基础

各位同学，欢迎来到第二章。这一章我们聊聊硬件。我知道很多人一听到“硬件架构”就头大，觉得这是搞芯片设计的人该关心的事。但说实话，你写AI训练算子，如果不懂GPU怎么干活，那优化基本靠蒙。

我刚开始做算子优化那会儿，就犯过这种错。写了个看起来很漂亮的kernel，结果跑起来比PyTorch原生还慢。后来一分析，发现是全局内存访问模式太乱，把带宽全浪费了。嗯，从那以后，我老老实实把GPU架构手册翻了三遍。

2.1 GPU/加速器架构概览

先看一张整体架构图，帮你建立直观印象。

这张图展示的是现代GPU的典型层次结构。最下面是全局内存，也就是我们常说的显存。往上走是L2缓存，再往上就是一个个GPC（图形处理簇）。每个GPC里包含多个TPC，每个TPC里又有两个SM（流式多处理器）。

SM才是真正干活的地方。每个SM内部有Warp调度器、寄存器文件、共享内存，还有一堆计算单元。你写的CUDA kernel，最终就是被拆成一个个warp，在这些SM上执行的。

2.2 内存层次结构

内存层次是算子优化的核心。我个人习惯把GPU内存分成三个层级：全局内存、共享内存、寄存器。每个层级的速度和容量天差地别。

内存类型	容量	延迟	作用域	典型用途
全局内存	GB级别	~400-800周期	所有线程	存储输入数据、权重、输出
共享内存	几十KB/SM	~20-30周期	单个Block内	数据复用、线程间通信
寄存器	几百KB/SM	~1周期	单个线程	局部变量、中间结果

2.2.1 全局内存

全局内存是GPU里最大的存储区域，所有线程都能访问。但代价是延迟极高——一次全局内存访问可能要等几百个时钟周期。

我在项目中遇到过最典型的坑：写了一个矩阵乘法，每次计算都从全局内存读数据。结果计算只用了几个周期，等数据却等了上百个周期。性能惨不忍睹。

2.2.2 共享内存

共享内存是每个SM内部的一块高速缓存，同一个Block内的线程可以共享。它的延迟比全局内存低一个数量级。

你想想看，如果多个线程需要反复读取同一块数据，把它从全局内存搬到共享内存，是不是比每次都去全局内存读要快得多？这就是共享内存的核心用法——数据复用。

// 典型用法：将数据从全局内存加载到共享内存
__global__ void kernel(float* input, float* output, int N) {
    __shared__ float s_data[256];  // 共享内存声明
    
    int tid = threadIdx.x;
    int idx = blockIdx.x * blockDim.x + tid;
    
    // 从全局内存加载到共享内存
    if (idx < N) {
        s_data[tid] = input[idx];
    }
    __syncthreads();  // 同步，确保所有线程加载完成
    
    // 在共享内存上做计算
    // ... 计算逻辑 ...
    
    // 写回全局内存
    if (idx < N) {
        output[idx] = s_data[tid];
    }
}

2.2.3 寄存器

寄存器是速度最快的存储，延迟只有1个周期。每个线程都有自己的寄存器，其他线程无法访问。

但寄存器也是稀缺资源。每个SM的寄存器总数是固定的，比如A100每个SM有65536个32位寄存器。如果每个线程用太多寄存器，SM能同时容纳的线程数就会减少，这会降低并行度。

为什么会这样？因为寄存器是分给线程的。假设一个SM有65536个寄存器，每个线程用32个，那这个SM最多同时容纳2048个线程。如果每个线程用64个，就只能容纳1024个了。

2.3 计算单元与线程调度

说完了存储，我们聊聊计算。GPU的计算核心是SM里的CUDA Core和Tensor Core。CUDA Core负责通用计算，Tensor Core专门加速矩阵运算。

2.3.1 Warp与线程调度

GPU执行线程的最小单位是warp。一个warp包含32个线程，它们在同一时刻执行同一条指令。这就是所谓的SIMT（单指令多线程）模型。

你想想看，如果warp里的32个线程走了不同的分支（if-else），会发生什么？嗯，GPU会串行执行所有分支，导致性能下降。这就是warp divergence问题。

// 避免warp divergence的写法
// 不好的写法：warp内线程走不同分支
if (threadIdx.x % 2 == 0) {
    // 偶数线程执行
} else {
    // 奇数线程执行
}

// 好的写法：将数据重新组织，让同一warp走相同分支
// 或者用三元运算符替代分支
float result = (threadIdx.x % 2 == 0) ? a : b;

2.3.2 线程块与网格

线程的组织方式也很重要。一个Block内的线程可以共享共享内存，也可以同步。不同Block之间则完全独立。

Block的大小会影响SM的利用率。太小了，SM吃不饱；太大了，可能超过SM的资源限制。我一般建议Block大小设为256或512，这是经过大量实践验证的黄金值。

参数	建议值	说明
Block大小	128-512	太小则SM利用率低，太大则资源受限
Grid大小	SM数量的2-4倍	保证SM有足够的任务可调度
每个SM的Block数	4-8个	隐藏内存访问延迟

2.3.3 Tensor Core

从Volta架构开始，NVIDIA引入了Tensor Core。这东西专门为矩阵乘法设计，一次可以执行4×4矩阵的乘加运算。对于AI训练来说，Tensor Core是性能的关键。

我建议你在写算子时，尽量让计算满足Tensor Core的要求。比如矩阵维度最好是8的倍数，数据类型用half或bfloat16。这样编译器会自动生成Tensor Core指令。

好了，这一章的内容就到这里。硬件架构的知识点比较多，但都是后续优化的基础。下一章我们开始讲具体的算子实现，到时候你会感谢今天花时间搞懂了这些底层细节。