AI Infra和算子开发(1)

发表于 本文字数： 1.7k 阅读时长 ≈ 2 分钟
无

GPU

整体而言，GPU 相比于 CPU 牺牲了复杂计算上的能力，取而代之的是大规模并行处理数据的能力。

gpu-structure

GPU 存在三个并行层级：

  1. thread：GPU 上最小执行单位，对应一次函数调用的一个实例。每个 thread 有自己的寄存器和程序计数器，每个 thread 执行逻辑相同，处理的数据不同（Single Instruction Multiple Data）。
  2. block：由多个能互相合作的 thread 组成，每个 block 有自己的共享内存，其中的线程可以交换数据，可以与其他 block 并行执行。
  3. grid：是一次 kernel launch 的所有 block 的集合，对应整个 GPU 上跑的任务，不同 block 之间完全独立，调度顺序无所谓。

warp

GPU 不是以 thread 为单位执行，而是以 warp（32 个 thread）为单位执行的。

SIMT（Single Instruction, Multiple Thread）的含义：

  • 这 32 个 thread 在同一个时钟周期执行完全相同的指令
  • 每个 thread 操作自己那份数据（不同寄存器值）

算子开发里要尽量保证同一 warp 内的 thread 走相同控制流。比如，

1
2
3
4
5
6
// 危险！同一个 warp 内不同 thread 走不同分支
if (thread_idx % 2 == 0) {
    do_A();   // 偶数 thread 走这
} else {
    do_B();   // 奇数 thread 走这
}

由于 warp 必须执行同一条指令，硬件会串行执行两个分支：先让走 A 的 thread 执行 do_A()，其他 thread 闲置；再让走 B 的执行 do_B()。原本的并行变成串行，性能直接砍半。

当 warp 访问全局内存时，32 个 thread 各自请求一个地址。硬件会把这些请求合并，此时连续的寻址会更好。

Streaming Multiprocessor

SM（Streaming Multiprocessor） 是 GPU 内部一个完整的“小处理器”，负责调度和执行 block。

一次 kernel launch 有成千上万个 block，但 SM 数量有限。规则是：

  1. 每个 SM 可以同时驻留多个 block
  2. SM 用 warp 调度在这些 block 的 warp 之间快速切换，掩盖延迟
  3. 一个 SM 同时能跑多少 warp，取决于资源限制

一个 SM 同时能驻留多少 warp，受三种资源限制：

  1. 每个 SM 最多支持的 warp 数
  2. 每个 SM 最多支持的 block 数
  3. 每个 block 用的寄存器/共享内存

Occupancy（占用率）记为实际驻留的 warp 数 / SM 最大 warp 数。

算子的性能指标

1. 核心利用率

Bandwidth Utilization（带宽利用率） = 实际带宽 / 硬件峰值带宽

Compute Utilization（MFU） = 实际 FLOPS / 硬件峰值 FLOPS

FLOPS：floating-point operations per second（每秒浮点算力）

MACs：Multiply ACcumulate operations（乘加累积操作次数），一个MACs包含一个乘法操作与一个加法操作，因此1个MACs约等价于2个FLOPs。

Roofline Model

算力 $\pi$ 是一个计算平台倾尽全力每秒钟所能完成的浮点运算数（FLOPS）。

带宽 $\beta$ 是一个计算平台倾尽全力每秒所能完成的内存交换量（Byte/s）。

计算强度上限 $I_{\text{max}}=\frac{\pi}{\beta}$ 单位是FLOPs/Byte。

roofline-model

Compute-bound：整个计算的耗时主要在于计算本身，对显存的读写耗时较低。

Memory-bound：整个计算的耗时主要集中在存储的访问上，计算本身耗时较低。

Memory-bound 的优化方向：减少访存

  1. 算子融合（Fusion）
  2. 数据布局 / 向量化访问
  3. 用更小的 dtype
  4. 减少中间结果落盘（不写回 HBM）

Compute-bound 的优化方向：提高算力利用率

  1. 用 Tensor Core
  2. Tiling（分块）
  3. 双缓冲 / 流水线（Pipelining）