如何通过共享内存和寄存器溢出优化CUDA内核性能

如何通过共享内存和寄存器溢出优化CUDA内核性能

当一个CUDA内核需要的硬件寄存器数量超过可用上限时，编译器会被迫将多余的变量转移到本地内存（local memory）中，这个过程被称为寄存器溢出（register spilling）。寄存器溢出严重影响性能，因为内核必须访问物理上位于全局内存（global memory）中的本地内存来读写这些溢出的数据，而全局内存的访问延迟远高于寄存器。

在CUDA Toolkit 13.0中，NVIDIA在编译流程中引入了一项新的优化特性：针对CUDA内核的共享内存寄存器溢出（shared memory register spilling）。本文将详细解释这项新特性，阐述其背后的动机，并指导您如何启用它。同时，我们还将探讨何时应该考虑使用它，以及如何评估其潜在影响。

特性	共享内存寄存器溢出
特性详情	将共享内存作为寄存器的后备存储，优先将高成本的寄存器溢出到共享内存中。
影响平台	所有支持PTXAS全程序编译模式（-rdc=false，默认模式）的平台。
用户影响	降低寄存器密集型内核的溢出延迟和L2缓存压力；共享内存使用量会增加。
启用方式 (CUDA 13.0+)	在内核定义处通过内联汇编使用.pragma enable_smem_spilling。在CUDA 13.0中默认为关闭。

表1. CUDA 13.0中PTXAS的变更摘要，该变更启用了针对CUDA内核的共享内存寄存器溢出功能

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

在CUDA 13.0中，PTXAS编译器增加了将寄存器溢出到共享内存的支持。当启用此功能时，编译器会优先将寄存器溢出到共享内存中。如果可用的共享内存不足，任何剩余的溢出仍会像以前一样回退到本地内存。这一变化利用了延迟更低的片上内存（on-chip memory）来存储溢出的值，从而实现性能优化。

问题背景与示例

在CUDA 13.0之前的工具包中，所有的寄存器溢出都被放置在本地内存中，该内存位于设备全局内存（off-chip device global memory）中。尽管更大的L1缓存帮助许多应用程序减少了溢出开销，但溢出的数据仍可能被写入L2缓存。这可能导致重要的缓存行被驱逐，从而对整体性能产生负面影响。这种影响在性能关键区域（如循环和频繁执行的代码段）中尤其明显，因为这些区域的寄存器压力通常很高。

在许多工作负载中，很大一部分共享内存实际上在运行时并未被使用。这可能是因为每个线程块的共享内存需求很低，或者内核的设计并未旨在最大化占用率（occupancy）。例如，如果每个SM的线程块数量受限于启动边界或寄存器压力而非共享内存使用量，那么每个块最终分配到的共享内存可能比实际需要的多。如果没有办法利用这些额外的共享内存，大部分都会被浪费掉。

思考以下代码示例。您不必理解每一行代码，只需注意这是一个特意设计用来消耗大量寄存器以引发溢出的内核。

// main.cu
#include <cuda_runtime.h>
#include <stdio.h>// 定义内核启动边界，每个线程块最多256个线程
extern "C" __launch_bounds__(256)
__global__ void foo(float *output_tensor, int num_elements) {// 计算全局唯一的线程IDint thread_id = blockIdx.x * blockDim.x + threadIdx.x;if (thread_id >= num_elements) return;// 声明大量volatile数组，以强制编译器使用大量寄存器，从而引发溢出// volatile关键字防止编译器进行可能减少寄存器使用的优化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 unrollfor (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];// 这里的计算相互依赖，且使用了__sinf, __cosf等函数，增加了计算复杂度和寄存器需求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 unrollfor (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;// 启动CUDA内核foo<<<gridSize, blockSize>>>(device_output, num_elements);// 同步设备，等待内核执行完毕cudaDeviceSynchronize();// 将结果从设备复制回主机cudaMemcpy(host_output, device_output, ARRAY_BYTES, cudaMemcpyDeviceToHost);// 打印部分结果以验证for (int i = 0; i < 10; ++i) { // 只打印前10个结果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 foo176 bytes stack frame, 176 bytes spill stores, 176 bytes spill loads
ptxas info    : Used 255 registers, used 0 barriers, 176 bytes cumulative stack size

请注意输出中的 176 bytes spill stores, 176 bytes spill loads。这明确表示发生了寄存器溢出，并且数据被存储和加载到了本地内存中。

此外，在这个例子中，编译后的内核完全没有使用任何共享内存，导致每个线程块分配到的共享内存被完全浪费。

CUDA 13.0中针对寄存器溢出的新解决方案是什么？

为了提高受寄存器限制的内核的性能，CUDA 13.0引入了一项新优化，允许将寄存器溢出重定向到共享内存而不是本地内存。通过利用片上共享内存，编译器将溢出的数据保持在离流式多处理器（SM）更近的位置，显著降低了访问延迟，并减轻了L2缓存的压力。在本地内存溢出通常会成为瓶颈的场景中，这种增强带来了显著的性能提升。

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

当我们编译之前的内核并启用共享内存寄存器溢出时，输出如下：

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

与之前的例子相比，请注意 spill stores 和 spill loads 现在都是0字节，表示没有发生到本地内存的溢出。同时，46080 bytes smem 表明内核现在使用了共享内存来处理溢出的数据。

Nsight Compute性能分析结果

为了展示这项优化的价值，我们对上述CUDA内核在启用和未启用共享内存溢出功能的情况下进行了基准测试。下表展示了Nsight Compute的对比结果，突出了三个关键指标的变化：持续时间、总周期和SM活动周期，证明了改进寄存器溢出处理带来的效率提升。

指标	未优化基线	优化后	提升幅度
Duration [us]	8.35	7.71	7.76%
Elapsed cycles [cycle]	12477	11503	7.8%
SM active cycles [cycle]	218.43	198.71	9.03%

表2. Nsight Compute结果对比：启用共享内存寄存器溢出优化前后的内核性能

如何启用共享内存寄存器溢出

共享内存寄存器溢出功能是在CUDA 13.0中引入的，早期版本的工具包不提供此功能。开发者需要通过在函数内部、函数声明之后，使用内联汇编添加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) {// 使用内联汇编启用共享内存寄存器溢出优化// 这条指令告诉PTXAS编译器，如果发生寄存器溢出，优先使用共享内存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 unrollfor (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 unrollfor (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;
}// main函数保持不变
int main() {// ...
}

共享内存寄存器溢出的局限性

这项优化为设备代码和PTXAS编译器提供了性能提升的机会，但它也带有一些重要的限制。它只在函数作用域内有效，并且在以下场景中使用会导致编译错误：

• 按函数编译模式（Per-function compilation mode）: 例如 nvcc -rdc=true 或 ptxas -c，以及设备调试模式（nvcc -G 或 ptxas -g）。
• 使用动态分配共享内存的内核。
• 在warp之间执行动态寄存器重分配的内核。

如果未明确指定启动边界（launch bounds），PTXAS在估算共享内存使用量时会假设每个线程块可能拥有的最大线程数。如果内核启动时使用的线程数少于估算值，每个块实际分配的共享内存可能会超过必要量，这可能限制在单个SM上并发运行的线程块数量，从而导致占用率下降和性能衰退。为确保更可预测的行为和更好的性能，建议仅在明确定义了启动边界时使用此功能。

真实工作负载的性能增益

我们在QUDA库中的多种CUDA内核上评估了这项优化，该库用于GPU上的格子量子色动力学（Lattice QCD）计算。如下图所示，该优化带来的性能增益通常在**5-10%**的范围内。这些改进源于将溢出重定向到共享内存，从而减少或完全消除了到本地内存的寄存器溢出。

图1. 启用共享内存寄存器溢出后，QUDA内核子测试的性能增益在5-10%范围内

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

CUDA 13.0现在包含一项PTXAS优化，允许通过共享内存高效处理寄存器溢出，从而在经历高寄存器压力的内核中提升性能。如果您的CUDA内核具有明确定义的启动边界和一致的共享内存利用率，请尝试通过内联pragma enable_smem_spilling 来启用共享内存溢出，并评估其带来的性能提升。