跳转至

硬件基础

在编写SIMD或GPU代码之前,你需要了解所编程的硬件。本文涵盖为什么并行性取代了时钟速度、现代CPU如何执行指令、什么是SIMD、用于推理性能的屋顶线模型,以及芯片架构的全景

  • 几十年来,软件免费变快:购买一个时钟频率更高的新CPU,你的程序无需修改一行代码就能运行得更快。这个时代大约在2005年结束。理解它为何结束以及什么替代了它,对任何想编写快速代码的人都至关重要。

免费性能的终结

  • 摩尔定律(1965年)观察到芯片上的晶体管数量大约每两年翻一番。这一规律维持了60年。更多晶体管意味着更小的晶体管,进而意味着更高的时钟频率,从而意味着更快的程序。

  • 但在2005年左右,时钟频率在大约4 GHz处撞上了墙壁。问题是功耗。芯片消耗的功率大约为:

\[P \propto C \cdot V^2 \cdot f\]
  • 其中 \(C\) 是电容(与晶体管数量成正比),\(V\) 是电压,\(f\) 是时钟频率。要提高频率,必须提高电压(以使晶体管更快地切换)。但功耗与 \(V^2 \cdot f\) 成比例,所以频率的小幅增加会导致功耗(和热量)的大幅增加。在4 GHz时,芯片已经达到100+瓦。达到8 GHz需要不切实际的冷却方案。

  • 解决方案:不让单个核心更快,而是在同一芯片上放置多个核心。一个4核芯片在3 GHz下使用与单个核心在4.5 GHz下相似的功耗,但可以做4倍的并行工作。这就是为什么每个现代CPU都有多个核心,以及为什么并行性(SIMD、多线程、GPU计算)是获得更高性能的唯一途径。

  • 对ML的影响:一个在单核上需要10分钟的训练步骤,无法通过购买更快的CPU来加速。只能通过使用更多核心(数据并行性,第6章)、更宽的SIMD单元(本章)或GPU(数千个核心)来加速。

现代CPU如何执行指令

  • 现代CPU核心远比第13章中简单的取指-译码-执行模型复杂。它使用几种技巧来每周期执行更多指令:

  • 超标量执行:CPU有多个执行单元(ALU、FPU、加载/存储单元),可以同时执行多个独立的指令。如果指令不相互依赖,现代核心每周期可能执行4-6条指令。

  • 乱序执行(OoO):CPU不按程序顺序执行指令。它向前看指令流,找到输入已准备好的指令,并立即执行,不论其位置。这隐藏了延迟:当一条指令等待来自内存的数据时(100+周期),CPU执行其他已准备好的指令。

  • 分支预测:条件分支(if语句、循环条件)造成不确定性:CPU在条件被评估之前不知道走哪条路径。为了避免停顿,CPU预测结果并沿预测路径投机执行。如果预测正确(使用现代预测器超过95%),则没有时间损失。如果错误,投机工作被丢弃,执行正确路径(约15周期惩罚)。

  • 投机执行:分支预测的延伸。CPU执行可能不需要的指令,赌它们会被需要。这填充了流水线并保持执行单元忙碌。

  • 所有这些都是自动的——CPU无需任何程序员干预即可完成。但它们只帮助指令级并行性(ILP):单条流内相互独立的指令。对于数据级并行性(对许多数据元素执行相同操作),我们需要SIMD。

SIMD:单指令多数据

  • SIMD是将一条指令同时应用于多个数据元素的思想。不是将两个数相加,而是在一条指令中将两个4(或8、或16)元素向量相加。

  • 无SIMD(标量):

// 逐元素相加两数组:4条加法指令
for (int i = 0; i < 4; i++) {
    c[i] = a[i] + b[i];  // 每次迭代一次加法
}
  • 有SIMD(向量化):
// 两数组相加:1条SIMD指令完成所有4次加法
#include <immintrin.h>  // x86 SIMD内联函数

__m128 va = _mm_load_ps(a);    // 加载4个浮点数到128位寄存器
__m128 vb = _mm_load_ps(b);    // 加载4个浮点数到另一个寄存器
__m128 vc = _mm_add_ps(va, vb); // 同时相加所有4对
_mm_store_ps(c, vc);            // 存储4个结果
  • SIMD版本用1/4的指令完成相同工作。这是理论上的4倍加速,通过每条指令处理4个浮点数而非1个实现。

向量寄存器

  • SIMD指令操作向量寄存器:保存多个数据元素的宽寄存器。
寄存器宽度 浮点数(32位) 双精度浮点数(64位) 名称
128位 4 2 SSE(x86)、NEON(ARM)
256位 8 4 AVX/AVX2(x86)
512位 16 8 AVX-512(x86)
可变(128-2048) 可变 可变 SVE/SVE2(ARM)
  • 更宽的寄存器 = 更多并行性。一条512位AVX-512指令一次处理16个浮点数,是标量代码理论上的16倍加速。实际上,由于内存带宽限制(计算速度可能超过向CPU输送数据的速度),加速比更低。

  • 对于ML:float32值的矩阵乘法从SIMD中获益巨大。内循环(两个向量的点积)直接映射到SIMD乘加指令。这就是为什么BLAS库(NumPy和PyTorch调用的)用SIMD进行了如此深度优化。

屋顶线模型

  • 你如何知道你的代码是否快速?屋顶线模型提供了一个框架,根据两个硬件限制来描述性能:

  • 峰值计算能力(FLOPS):每秒最大浮点运算次数。对于一个4 GHz CPU,配备256位AVX(每条指令8个浮点数)和2个FMA单元:\(4 \times 10^9 \times 8 \times 2 = 64\) GFLOPS。

  • 峰值内存带宽(字节/秒):数据从内存到CPU的最大传输速度。现代CPU可能有50 GB/s的内存带宽。

  • 代码的算术强度是计算与内存访问的比率:

\[\text{算术强度} = \frac{\text{FLOPS}}{\text{传输的字节数}}\]
  • 如果算术强度低(每加载字节的操作数少),你的代码是内存受限的:大部分时间花在等待数据上。让计算更快(更宽的SIMD、更高的时钟)不会有帮助。

  • 如果算术强度高(每字节多次操作),你的代码是计算受限的:大部分时间花在计算上。更快的内存不会有帮助。

  • 屋顶线:

\[\text{可达FLOPS} = \min\left(\text{峰值FLOPS}, \; \text{带宽} \times \text{算术强度}\right)\]
  • 矩阵乘法具有高算术强度:\(O(n^3)\) 次操作作用于 \(O(n^2)\) 数据,因此强度 \(\approx O(n)\)。对于大矩阵,它是计算受限的。这就是为什么GPU(高计算能力)主导矩阵密集型的ML工作负载。

  • 逐元素操作(ReLU、加法、乘法)具有低算术强度:每加载一个元素1次操作。这些是内存受限的。让GPU更快没有帮助;你需要更快的内存(或者将这些操作与计算密集型操作融合,以避免独立的内存往返)。

  • 屋顶线模型解释了为什么核函数融合如此重要:将matmul与偏置加法和ReLU组合成一个核函数,避免了将中间结果写入内存并重新读取,将三个内存受限操作转化为一个计算受限操作。

延迟与吞吐量

  • 延迟是完成一个操作所需的时间。吞吐量是单位时间内完成的操作数量。

  • 打个比方:公交车延迟高(每站都停),但吞吐量高(一次搭载50人)。出租车延迟低(直达你的目的地),但吞吐量低(搭载1-4人)。

  • GPU是公交车:每次操作延迟高(每条指令需要许多周期完成),但吞吐量巨大(数千个核心同时处理)。CPU是出租车:延迟低(乱序执行、分支预测、深层缓存最小化延迟),但吞吐量有限(4-64个核心)。

  • 这就是为什么GPU更适合ML训练(吞吐量重要:处理数百万个样本)而CPU更适合操作系统任务(延迟重要:立即响应按键)。

  • 流水线将延迟转化为吞吐量。如果一条指令需要5个周期,但流水线每周期开始一条新指令,则吞吐量是每条指令1周期(即使每条指令需要5个周期完成)。这和第13章的CPU流水线是同一原理,但它适用于每个层面:SIMD单元、内存控制器和GPU核心都是流水线化的。

芯片架构全景

  • 你编写代码的硬件决定了哪些SIMD指令可用:

x86(Intel, AMD)

  • 主导台式机、笔记本电脑和数据中心CPU。SIMD:SSE(128位)、AVX/AVX2(256位)、AVX-512(512位)。Intel AMX提供专用的矩阵乘法单元用于AI工作负载。

  • 优势:最高单核性能、最宽SIMD、成熟的软件生态系统(MKL、oneDNN)。

  • 弱点:高功耗、复杂指令集、昂贵。

ARM

  • 主导移动设备(每一部智能手机),在服务器(AWS Graviton、Ampere Altra)和笔记本电脑(Apple M系列)中增长。SIMD:NEON(128位)、SVE/SVE2(可伸缩,128-2048位)。

  • 优势:出色的功耗效率(每瓦性能)、自定义核心(Apple M4在单核性能上媲美Intel,功耗仅为其一小部分)。

  • 弱点:较窄的SIMD(NEON仅为128位,虽SVE可更宽)、用于HPC的软件生态系统较小。

Apple Silicon(M1/M2/M3/M4)

  • 基于ARM并带有自定义扩展。包含AMX(Apple矩阵扩展)——未公开的矩阵乘法单元,Accelerate框架将其用于BLAS操作。统一内存架构:CPU和GPU共享同一物理内存,消除了CPU↔GPU拷贝的瓶颈。

  • 对于ML:Apple的神经网络引擎(16核,专用ML加速器)和统一内存使M系列芯片在本地ML推理和小规模训练方面出奇地强大。不过没有CUDA:你必须使用Metal(Apple的GPU API)或MLX(Apple的ML框架)。

RISC-V

  • 开源ISA。无许可费用(不像ARM)。在嵌入式系统、物联网和研究领域增长。SIMD:"V"(向量)扩展提供类似于ARM SVE的可伸缩向量处理。

  • 对于ML:在ML工作负载上尚不能与x86/ARM竞争,但值得关注。几个AI加速器初创公司使用RISC-V核心。

GPU(NVIDIA、AMD、Intel)

  • 在文件04-05中深入介绍。数千个为吞吐量优化的简单核心。NVIDIA以CUDA主导ML;AMD以ROCm竞争;Intel以Arc GPU和Gaudi加速器进入市场。

TPU(Google)

  • 专门为ML设计的自定义ASIC。为矩阵乘法优化的脉动阵列。在文件05中介绍。

热与功耗约束

  • 性能最终受限于功耗和散热:

  • TDP(热设计功耗):芯片可以持续消耗的最大功率。笔记本电脑CPU可能有15W TDP;服务器CPU 250W;数据中心GPU 700W(NVIDIA B200)。

  • 暗硅:在任何给定时刻,为了保持在热预算内,必须关闭芯片的相当大一部分晶体管。理论上芯片可以同时使用所有晶体管,但会熔化。

  • 功耗效率(FLOPS/瓦)日益成为重要指标,而非原始FLOPS。这就是为什么:

    • ARM正在接管数据中心(相比x86更好的FLOPS/瓦)。
    • TPU尽管峰值FLOPS较低,但仍与GPU竞争(对于ML工作负载,FLOPS/瓦好得多)。
    • 量化(INT8、FP8)不仅关乎内存:它也降低了每次操作的功耗。
  • 对于大规模ML:训练前沿LLM数月消耗兆瓦级功率。电费可能超过硬件成本。功耗效率直接影响AI研究的经济性。

实践:在C++中测量性能

  • 要推理性能,你需要测量它。以下是一个最小的C++基准测试设置:
#include <iostream>
#include <chrono>
#include <vector>

// 标量加法
void add_scalar(const float* a, const float* b, float* c, int n) {
    for (int i = 0; i < n; i++) {
        c[i] = a[i] + b[i];
    }
}

int main() {
    const int N = 1 << 24;  // 约1600万个元素
    std::vector<float> a(N, 1.0f), b(N, 2.0f), c(N);

    // 预热(填充缓存,触发频率缩放)
    add_scalar(a.data(), b.data(), c.data(), N);

    // 基准测试
    auto start = std::chrono::high_resolution_clock::now();

    for (int trial = 0; trial < 100; trial++) {
        add_scalar(a.data(), b.data(), c.data(), N);
    }

    auto end = std::chrono::high_resolution_clock::now();
    double elapsed = std::chrono::duration<double>(end - start).count();

    double total_bytes = 3.0 * N * sizeof(float) * 100;  // 读a、读b、写c
    double bandwidth = total_bytes / elapsed / 1e9;        // GB/s

    std::cout << "时间: " << elapsed << " s\n";
    std::cout << "带宽: " << bandwidth << " GB/s\n";

    return 0;
}
# 启用优化编译
g++ -O3 -march=native -o bench bench.cpp
./bench
  • 这段代码中的关键C++概念

    • #include <vector>:动态数组(std::vector<float>)——类似Python的list但带类型且在内存中连续。
    • a.data():返回底层数组的原始指针(float*)——SIMD内联函数需要。
    • std::chrono:用于基准测试的高分辨率计时器。
    • -O3:最高编译器优化级别。编译器可能自动向量化你的循环(自动使用SIMD)。-march=native 启用你的CPU支持的所有SIMD指令。
  • 为什么需要预热:首次运行填充缓存并可能触发CPU频率缩放(睿频加速)。后续运行更具代表性。

  • 为什么测量带宽:对于内存受限的操作(如逐元素加法),有意义的度量是带宽(GB/s),而不是FLOPS。如果你的测量带宽接近硬件极限(DDR5约50 GB/s),你是内存受限的,SIMD不会有多大帮助(瓶颈是内存,而非计算)。

编程任务(使用CoLab或笔记本)

  1. 计算常见ML操作的算术强度,并将它们分类为内存受限或计算受限。

    import jax.numpy as jnp
    
    def arithmetic_intensity(flops, bytes_transferred):
        return flops / bytes_transferred
    
    # 逐元素ReLU:每元素1次比较,读取+写入
    n = 1024
    relu_flops = n  # 每元素1次操作
    relu_bytes = 2 * n * 4  # 读取输入+写入输出(float32)
    print(f"ReLU: {arithmetic_intensity(relu_flops, relu_bytes):.2f} FLOPS/byte → 内存受限")
    
    # 矩阵乘法:2*n^3次操作,读取2*n^2 + 写入n^2个浮点数
    matmul_flops = 2 * n**3
    matmul_bytes = 3 * n**2 * 4  # 读取A + 读取B + 写入C
    print(f"矩阵乘法 ({n}×{n}): {arithmetic_intensity(matmul_flops, matmul_bytes):.0f} FLOPS/byte → 计算受限")
    
    # 层归一化:约5n次操作(均值、方差、归一化),读取+写入
    ln_flops = 5 * n
    ln_bytes = 2 * n * 4
    print(f"LayerNorm: {arithmetic_intensity(ln_flops, ln_bytes):.2f} FLOPS/byte → 内存受限")
    
    # 3x3卷积:2*9*C_in*C_out*H*W,读取卷积核+特征图+写入输出
    C_in, C_out, H, W = 64, 128, 32, 32
    conv_flops = 2 * 9 * C_in * C_out * H * W
    conv_bytes = (9 * C_in * C_out + C_in * H * W + C_out * H * W) * 4
    print(f"Conv3x3: {arithmetic_intensity(conv_flops, conv_bytes):.0f} FLOPS/byte → 计算受限")
    

  2. 演示为什么并行性重要。比较顺序执行与并行(NumPy)执行随数据规模增长的表现。

    import numpy as np
    import time
    
    for n in [1000, 10000, 100000, 1000000, 10000000]:
        a = np.random.randn(n).astype(np.float32)
        b = np.random.randn(n).astype(np.float32)
    
        # "顺序执行"(Python循环)
        start = time.time()
        c = [a[i] * b[i] for i in range(min(n, 100000))]  # 上限10万以确保合理
        seq_time = time.time() - start
        if n > 100000:
            seq_time *= n / 100000  # 外推
    
        # "并行"(NumPy,内部使用SIMD+多线程)
        start = time.time()
        c = a * b
        par_time = time.time() - start
    
        print(f"n={n:>10,}  顺序={seq_time:.4f}s  并行={par_time:.6f}s  "
              f"加速比={seq_time/par_time:.0f}x")