GPU寄存器缓存的魔法秀 —— 当 CUDA 代码化身为闪电侠

不会弹吉他的布鲁克
134 阅读8分钟

在这个技术魔法秀中,我们带来了一招“寄存器缓存”大法，颠覆了传统的共享内存操作，让 CUDA 程序从慢吞吞变成快如闪电！别急，接下来我们就一探这门绝技的神秘奥义

image.png

warp level cache

首先 warp 级别的缓存位于哪里,如下图

image.png

  • 全局内存：犹如整个城市，资源多但交通拥堵；

  • 共享内存：像个小区，邻里之间互通有无，速度飞快；

  • 寄存器：则是每个线程的私人花园，空间虽小但效率爆棚。

遗憾的是，单个 warp（就像一群打工兄弟）竟然没有专属的邻里公园！这正是寄存器缓存登场的最佳时机，让这些兄弟们能共享自家花园里的果实，而不必再排队等共享内存的绿化车来配送。

魔法工具：Shuffle 指令的绝世好用

image.png

从 CUDA 的 Kepler 架构（2012 年出品）开始，NVIDIA 就给我们送上了神兵利器—— __shfl_sync 指令。借助这把利器，每个线程可以在 warp 内直接偷取（哦，不，分享）彼此寄存器中的数据，完全省去了那令人头疼的线程同步麻烦。
优势显而易见：寄存器带宽杠杠的，延迟低得让人想跳舞，同时还可将本来就捉襟见肘的共享内存资源，巧妙“增容”！

但世上无完美法术，直接将现有共享内存代码改写成 shuffle 魔法，往往需要重构整个算法,下面我们通过一个案例来演示我们这场魔法秀

__global__ void one_stencil (int *A, int *B, int sizeOfA)
{
    extern __shared__ int s[];
    
    // Id of thread in the block.
    int localId = threadIdx.x;

    // The first index of output element computed by this block.
    int startOfBlock = blockIdx.x * blockDim.x;

    // The Id of the thread in the scope of the grid.
    int globalId = localId + startOfBlock;

    if (globalId >= sizeOfA)
        return;

    // Fetching into shared memory.
    s[localId] = A[globalId];
    if (localId < 2 && blockDim.x + globalId < sizeOfA) {
        s[blockDim.x + localId] =  A[blockDim.x + globalId];
    }

    // We must sync before reading from shared memory.
    __syncthreads();

    // Each thread computes a single output.
    if (globalId < sizeOfA - 2)
        B[globalId] = (s[localId] + s[localId + 1] + s[localId + 2]) / 3;
}

一个经典的 1D stencil 计算为例（相当于给每个输出值“熬汤”时调配了三种调料），详细展示了如何将共享内存版本华丽转身为寄存器缓存版本。整个过程分为四大戏码：

  1. 共享内存版的初登场
    各线程先把全局内存中的数据搬进共享内存，等待“全员集合”后，大家齐心协力计算出结果，再挥师写回全局内存。这一流程虽稳重，但略显冗长。具体步骤如下:

    • 每个线程从全局内存拷贝数据到共享内存
    • 调用__syncthreads() 保证所有线程数据就绪后，计算输出并写回全局内存
    • 整体计算步骤如下图 image.png

  2. 锁定 Warp 内的秘密原料
    分析发现，一个 warp 需要连续的输入数据——比如 32 个输出竟需 34 份数据，简直比大厨备料还要精细！

  3. 数据分派：寄存器里的轮流上菜
    采用轮转（round-robin）策略，将各个输入巧妙地分配到 warp 内各线程的寄存器中，有的线程甚至要接待,例如：对于 32 线程的 warp，每个线程按照索引 i % 32 来存放相应的输入数据，部分线程可能需要存放两个数据。两道菜（数据），完美模拟共享内存的银行调度机制。

  4. 分工协作：沟通与计算的华丽对决
    整个过程分为“通信”和“计算”两个阶段。通信阶段中，每个线程先“发布”自己的数据，再“读取”队友的“特供”，最后在计算阶段将所有数据拼成一锅美味大餐。为实现这一连串精妙操作，__shfl_sync 指令化身为最佳助攻，一键搞定所有数据调度，还贴心解决了“寄存器缓存冲突”的尴尬局面。具体步骤如下:

    • 将 kernel 分为通信和计算两个阶段。通信阶段通过两个抽象原语：

      • Read(src_tid, remote_reg) ：从指定线程 src_tid 的寄存器中读取数据。
      • Publish(local_reg) ：将本地寄存器数据发布，以便其他线程读取。

    • 通过多阶段的通信与计算，完成了 1-stencil 中每个输出累加所需输入的操作。

    • 为实现这一过程，使用 __shfl_sync 将 Publish 和 Read 操作合并，解决了线程间数据共享的问题，同时注意处理“寄存器缓存冲突”（当同一线程需要发布多个不同数据时可能发生的问题，需要多次访问才能满足所有请求）。

    image.png

通过闪电侠替换 Publish-Reads

__shfl_sync版本如下

__global__ void one_stencil_with_rc (int *A, int *B, int sizeOfA)
{
    // Declaring local register cache.
    int rc[2];

    // Id of thread in the warp.
    int localId = threadIdx.x % WARP_SIZE;

    // The first index of output element computed by this warp.
    int startOfWarp = blockIdx.x * blockDim.x + WARP_SIZE*(threadIdx.x / WARP_SIZE);

    // The Id of the thread in the scope of the grid.
    int globalId = localId + startOfWarp;

    if (globalId >= sizeOfA)
        return;

    // Fetching into shared memory.
    rc[0] = A[globalId];
    if (localId < 2 & & WARP_SIZE + globalId < sizeOfA)
    {
        rc[1] =  A[WARP_SIZE + globalId];
    }

    // Each thread computes a single output.
    int ac = 0;
    int toShare = rc[0];

    for (int i = 0 ; i < 3 ; ++i)
    {
        // Threads decide what value will be published in the following access.
        if (localId < i)
            toShare = rc[1];

        // Accessing register cache.
        unsigned mask = __activemask();
        ac += __shfl_sync(mask, toShare, (localId + i) % WARP_SIZE);
    }

    if (globalId < sizeOfA - 2)
        B[globalId] = ac/3;
}

随着 k 值增加，寄存器缓存相对于共享内存实现的 k- stencil 的速度提升情况,对于较小的 k 值，数据重用较小或可以忽略不计，因此加速效果也较小，但随着 k 的增大，寄存器缓存实现的重用增加，加速效果也随之增加。

image.png

当 k 达到 12 时，加速效果趋于稳定，因为由于输入数据边缘重叠，寄存器缓存需要执行更多的全局内存访问，导致两个连续的 warp 重复读取相同的数据。

一种常见的减少全局内存访问的方法是线程粗化。该技术通过增加每个线程生成的输出数量，使得部分输入数据可以在多个迭代中复用，并通过存储在寄存器中来实现这一目的。

对于寄存器缓存而言，线程粗化在达到预期性能提升时显得尤为关键。这是由于共享缓存的线程数量较少所致。由于寄存器缓存仅限于单个 warp 内的线程，因此仅预取并缓存了该 warp 所需的输入数据。然而，输入数据的复用可能会跨越连续的 warp。以本文开发的 1-stencil kernel 为例，warp i 中第一个线程的输入与 warp i-1 中最后一个线程的输入是相同的，因此该值会从全局内存中读取两次。对于每个线程块中 32 个 warp 的情况，寄存器缓存实现需要进行 34 * 32 = 1088 次全局内存访问，比标准共享内存实现多出约 6%。需要注意的是，随着 k 值的增加，冗余全局内存访问的次数会显著上升，例如在 k = 16 时，几乎有一半的访问是冗余的。

image.png

线程粗化有助于减少冗余全局内存访问带来的影响。图6展示了在不同 k 值下，每个线程输出数量（1 至 8 个）变化时，与共享内存实现相比所获得的加速比。得益于线程粗化，寄存器缓存版本的加速比最高可达 1.8 倍。但对于较大的 k 值，寄存器缓存所需的寄存器数量过多，无法全部装入物理寄存器，编译器会将部分数据溢出到局部内存（通常存放在 L1 缓存中，如果空间不足，也可能溢出到全局内存），从而导致性能下降

何时使用寄存器缓存？
寄存器缓存并非在所有情况下都适用。首先，数据的访问模式必须在编译期就已确定。其次，寄存器缓存的高效运行依赖于足够的空闲寄存器；否则，寄存器溢出到全局内存将严重影响性能，就像图6中 k=25 的情况那样。

使用寄存器缓存的适用条件与 CUDA 9 新特性

  • 适用条件
    寄存器缓存并非在所有场景下都适用：

    1. 输入数据的访问模式必须在编译期确定即输入数据的访问模式必须在编译期就“写进剧本”；
    2. 必须有足够的空闲寄存器，否则会出现寄存器溢出问题，严重时性能会大幅下降（如文中 k=25 的实验情况）。

  • CUDA 9 与 Cooperative Groups
    CUDA 9 引入了 Cooperative Groups，鼓励程序员显式地组织和同步线程组，特别是在 warp 级别的同步。在新模型中，传统的 __shfl() 已被废弃，取而代之的是必须显式传递活动线程掩码的 __shfl_sync()。文章中通过实例展示了如何利用 __activemask() 获得活动线程掩码，以及如何使用 Cooperative Groups 创建静态大小的线程块分区，从而支持更安全的 warp 内数据共享。即使在不同 GPU 架构（如 Pascal 和 Volta）下，使用 __shfl_sync() 能确保代码的正确执行。

尾声：魔法秀背后的思考

通过这一寄存器缓存的魔法，我们带来了一种全新的思维方式——将硬件资源抽象成虚拟缓存，突破了传统内存层次的限制。对于那些追求极致性能的开发者来说，这无疑是一剂强心针，助你在 CUDA 世界中大展拳脚！

avatar
文章
阅读
粉丝