数据中心/云端

如何通过共享内存寄存器溢出来提高 CUDA 内核性能

当 CUDA 内核所需的硬件寄存器数量超过可用数量时,编译器会将多余的变量溢出到本地内存中,这一过程称为寄存器溢出。由于本地内存实际上位于全局内存中,内核在读取和写入溢出数据时必须访问物理内存,因此寄存器溢出会降低性能。

CUDA 工具包 13.0 中,NVIDIA 在编译流程中引入了一项新的优化功能:CUDA 内核的共享内存寄存器溢出。本文将介绍该功能的设计初衷,详细说明如何启用该功能,并提供关于何时应考虑使用该功能以及如何评估其潜在影响的指导建议。

特征 共享内存寄存器溢出
特征详细信息 启用共享内存作为寄存器的备份存储,优先考虑将高成本寄存器溢出到共享内存。
受影响的平台 所有平台均处于 PTXAS 完整程序编译模式 (-rdc = false) 。这是默认的 PTXAS 模式
用户影响 减少寄存器密集型内核的溢出延迟和 L2 压力;增加共享内存使用量
Opt-in (CUDA 13.0+) .pragma enable_smem_spilling 内联组件在内核定义处。CUDA 13.0 中的默认值:false (阅读 PTX 文档)
表 1. CUDA 13.0 中 PTXAS 的更改摘要,这些更改允许共享内存寄存器溢出以防止 CUDA 内核中的寄存器溢出

共享内存寄存器溢出如何优化性能?

在 CUDA 13.0 中,PTXAS 增加了将寄存器溢出数据存储到 CUDA 内核共享内存的支持。启用该功能后,编译器会优先将溢出的寄存器数据分配到共享内存中;当共享内存不足时,剩余的溢出数据将回退至本地内存,行为与此前版本保持一致。此项改进通过利用片上低延迟的共享内存来存放溢出数据,从而实现了性能优化。

通过示例概述问题

在 CUDA 13.0 之前的工具包中,所有寄存器溢出的数据都会被存放在本地内存中,而本地内存位于芯片外的设备全局内存里。尽管较大的 L1 缓存能在一定程度上降低许多应用程序的溢出开销,但溢出数据仍可能写入 L2 缓存,从而导致重要的缓存行被逐出,对整体性能产生负面影响。这种影响在寄存器压力较高的性能关键区域尤为显著,例如循环体或代码中频繁执行的部分。

在许多工作负载中,很大一部分共享内存在运行时往往处于未使用状态。这种情况通常发生在每个线程块对共享内存的需求较低,或内核设计无法充分提升占用率时。例如,当每个 SM 上的线程块数量受限于启动边界或寄存器压力,而非共享内存的使用量时,每个线程块最终分配的共享内存可能远超其实际所需。如果这些额外分配的共享内存无法被有效利用,其大部分将被浪费。

请参考以下代码示例。理解每一行的具体含义并非关键,但请注意,该内核旨在使用大量寄存器,以达到引发寄存器溢出的效果。

/-- main.cu --
#include <cuda_runtime.h>
#include <stdio.h>

extern "C" __launch_bounds__(256)
__global__ void foo(float *output_tensor, int num_elements) {

    int thread_id = blockIdx.x * blockDim.x + threadIdx.x;
    if (thread_id >= num_elements) return;

    volatile float input_feature[89], weight_scaled[89], bias_added[89], pre_activation[89];
    volatile float activation_sin[89], activation_cos[89], output_accum[89];

    #pragma unroll
    for (int i = 0; i < 89; ++i) {
        input_feature[i] = (float)thread_id + i;
        weight_scaled[i] = input_feature[i] * 2.0f;
        bias_added[i] = 5 + weight_scaled[i];
        activation_sin[i] = __sinf(bias_added[i] * pre_activation[i]);
        activation_cos[i] = __cosf(activation_sin[i % 2] + pre_activation[i]);
        float product = input_feature[i] * weight_scaled[i];
        float squared = product * product;
        float biased = squared + bias_added[i % 4];
        float shifted_sin = __sinf(biased * 0.5f);
        float shifted_cos = __cosf(shifted_sin + 1.0f);
        float amplified = shifted_cos * bias_added[i % 5];
        float combined = amplified + activation_cos[i];
        output_accum[i] = combined;
    }
    volatile float sum = 0.0f;
    #pragma unroll
    for (int i = 0; i < 89; ++i) {
        sum += input_feature[i] + weight_scaled[i] + bias_added[i] + pre_activation[i]
             + activation_sin[i] + activation_cos[i] + output_accum[i];
    }

    output_tensor[thread_id] = sum;
}

int main() {
    const int num_elements = 896;
    const int ARRAY_BYTES = num_elements * sizeof(float);
    float host_output[num_elements];
    float *device_output;

    cudaMalloc(&device_output, ARRAY_BYTES);

    const int blockSize = 256;
    const int gridSize = (num_elements + blockSize - 1) / blockSize;

    foo<<<gridSize, blockSize>>>(device_output, num_elements);
    cudaDeviceSynchronize();

    cudaMemcpy(host_output, device_output, ARRAY_BYTES,   cudaMemcpyDeviceToHost);

    for (int i = 0; i < num_elements; ++i) {
        printf("host_output[%d] = %f\n", i, host_output[i]);
    }

    cudaFree(device_output);
    return 0;
}
nvcc -arch=sm_90 -Xptxas -v main.cu

当程序正常编译(未出现共享内存寄存器溢出)时,输出如下:

ptxas info    : Compiling entry function 'foo' for 'sm_90'
ptxas info    : Function properties for foo
    176 bytes stack frame, 176 bytes spill stores, 176 bytes spill loads
ptxas info    : Used 255 registers, used 0 barriers, 176 bytes cumulative stack size

请注意,输出显示了存储和加载操作的“溢出”,这意味着寄存器将溢出至本地内存。

此外,在本示例中,编译后的内核未使用任何共享内存,因此按块分配的共享内存也完全未被使用。

CUDA 13.0 中引入的寄存器溢出附加解决方案是什么?

为提升寄存器受限内核的性能,CUDA 13.0 引入了一项新优化,可将寄存器溢出数据重定向至共享内存,而非本地内存。通过利用片上共享内存,编译器能使溢出数据更接近流多处理器,从而显著降低访问延迟,并减轻 L2 缓存的压力。在本地内存溢出通常成为性能瓶颈的场景下,该优化可带来显著的性能提升。

启用优化后,编译器会优先尝试将寄存器溢出到可用的共享内存中;如果共享内存空间不足,则退回到使用本地内存,从而确保程序的正确性。

编译前一个内核以启用共享内存寄存器溢出时,输出结果如下所示:

ptxas info    : Compiling entry function 'foo' for 'sm_90'
ptxas info    : Function properties for foo
    0 bytes stack frame, 0 bytes spill stores, 0 bytes spill loads
ptxas info    : Used 255 registers, used 0 barriers, 46080 bytes smem

请注意,与之前的示例相比,本示例不会发生溢出。共享内存的使用由 46080 bytes smem 表示。

支持和不支持共享内存寄存器溢出优化的 Nsight 计算结果 

为体现该优化的价值,对上述 CUDA 内核在启用与不启用共享内存溢出功能两种情况下进行了基准测试。表 2 展示了 Nsight Compute 在启用共享内存寄存器溢出优化前后的内核性能对比结果,重点呈现了三个关键指标的变化:执行时间、经过周期数以及 SM 活动周期数,充分反映了改进寄存器溢出处理机制后带来的效率提升。

指标 不含优化的值 (基准) 优化后的值 改进
持续时间【us】 8.35 7.71 7.76%
运行周期【cycle】 12477 11503 7.8%
SM 活动周期【cycle】 218.43 198.71 9.03%
表 2。 Nsight Compute 结果对比了启用共享内存寄存器溢出优化前后的内核性能

如何选择共享内存寄存器溢出 

CUDA 13.0 引入了共享内存寄存器溢出功能,但该功能在早期版本的工具包中不可用。针对 CUDA 13.0 及更高版本进行开发的开发者,必须在函数声明后的函数体内通过内联组件直接添加 PTX pragma enable_smem_spilling 通过内联组件 直接添加 PTX pragma enable_smem_spilling,以显式启用此功能。

#include <cuda_runtime.h>
#include <stdio.h>

extern "C" __launch_bounds__(256)
__global__ void foo(float *output_tensor, int num_elements) {
    asm volatile (".pragma \"enable_smem_spilling\";");

    int thread_id = blockIdx.x * blockDim.x + threadIdx.x;
    if (thread_id >= num_elements) return;
    volatile float input_feature[89], weight_scaled[89], bias_added[89], pre_activation[89];
    volatile float activation_sin[89], activation_cos[89], output_accum[89];
    #pragma unroll
    for (int i = 0; i < 89; ++i) {
        input_feature[i] = (float)thread_id + i;
        weight_scaled[i] = input_feature[i] * 2.0f;
        bias_added[i] = 5 + weight_scaled[i];
        activation_sin[i] = __sinf(bias_added[i] * pre_activation[i]);
        activation_cos[i] = __cosf(activation_sin[i % 2] + pre_activation[i]);

        float product = input_feature[i] * weight_scaled[i];
        float squared = product * product;
        float biased = squared + bias_added[i % 4];
        float shifted_sin = __sinf(biased * 0.5f);
        float shifted_cos = __cosf(shifted_sin + 1.0f);
        float amplified = shifted_cos * bias_added[i % 5];
        float combined = amplified + activation_cos[i];
        output_accum[i] = combined;
    }
    volatile float sum = 0.0f;

    #pragma unroll
    for (int i = 0; i < 89; ++i) {
        sum += input_feature[i] + weight_scaled[i] + bias_added[i] + pre_activation[i]
             + activation_sin[i] + activation_cos[i] + output_accum[i];
    }
    output_tensor[thread_id] = sum;
}

共享内存寄存器溢出有哪些限制?

该优化为设备代码和 PTXAS 编译器带来了性能提升的机会,但存在重要限制:它仅在函数范围内有效,无法用于以下场景,否则可能导致编译错误。

  • 每函数编译模式,例如 nvcc -rdc=trueptxas -c, nvcc -Gptxas -g, nvcc -ewpptxas -ewp。请注意,设备调试编译模式 ( nvcc -Gptxas -g) 也意味着每函数编译。
  • 使用动态分配共享内存的内核。
  • 执行跨线程束动态重新分配寄存器的内核。

如果未显式指定 启动范围,PTXAS 在估算共享内存使用量时,会假定每个线程块包含最大可能的线程数。当核函数启动时的实际线程数少于该假设值时,每个线程块分配的共享内存可能会超出实际所需,从而可能限制流多处理器上可并发运行的线程块数量。这可能导致占用率降低,进而影响性能。为确保行为更可预测并获得更好的性能,建议仅在明确定义启动边界时使用此功能。

对于实际工作负载,可以实现哪些性能提升?

我们在 QUDA 库(用于 GPU 上的格点 QCD 计算)中的多个 CUDA 内核中评估了该优化。如图所示,优化带来的性能提升通常在 5% 到 10% 之间。这一改进主要得益于将寄存器溢出重定向至共享内存,从而减少甚至完全消除了对本地内存的溢出访问。

Bar chart showing the percentage performance gain across QUDA kernel subtests after enabling shared memory register spilling. Most tests exhibit improvements in the 5–10% range.
图 1。启用共享内存寄存器溢出后,QUDA 内核子测试的性能提升幅度在 5 – 10% 之间

开始使用共享内存寄存器溢出优化 

CUDA 13.0 现已包含 PTXAS 优化,可使共享内存更高效地处理寄存器溢出,从而提升寄存器压力较高的内核性能。如果您的 CUDA 内核具有明确的启动边界和稳定的共享内存使用率,可尝试使用内联 pragma enable_smem_spilling 将寄存器溢出至共享内存。

致谢

感谢以下 NVIDIA 贡献者:Jerry Zheng、Kate Clark、Howard Chen、Neumann Hon、Jaewook Shin、Abhishek Patwardhan 和 Yufan Cheng。

 

标签