CUB单次调用API：简化GPU并行算法开发

用户576110558132
3 阅读7分钟

CUB单次调用API：简化GPU并行算法开发

C++模板库CUB是实现高性能GPU基础算法（如扫描、排序、归约等）的首选工具。然而，其传统的“两阶段”API需要先估算临时内存大小，再显式分配内存并执行内核，这种分离虽然提供了灵活性，但也导致了大量重复的样板代码。

本文介绍CUDA 13.1中引入的新CUB单次调用API。该API通过在底层自动管理内存，在完全不损失性能的前提下，极大简化了开发流程。

CUB是什么？

当需要在GPU上运行标准算法（如扫描、直方图或排序）时，CUB通常是最快的实现方式。作为NVIDIA CUDA核心计算库的核心组件，CUB旨在抽象手动CUDA线程管理的复杂性，同时不牺牲性能。

像Thrust这样的库提供了类似于C++标准模板库的高级“主机端”接口，适合快速原型开发。而CUB则提供了一套“设备端”基础算法，使开发者能够将高度优化的算法直接集成到自己的自定义内核中。

传统的CUB两阶段API

CUB被广泛推荐用于充分发挥NVIDIA GPU的计算能力，但其使用方式存在一些复杂性。一个直观的、单次执行的流程通常期望：单次调用函数就能执行底层算法并立即获得结果。函数的副作用，如修改变量或返回结果，预期对下一条语句立即可见。

CUB的执行模型与此不同。调用CUB原语需要两步：首先计算所需的设备内存大小，然后显式分配内存并执行内核。

以下是传统的CUB调用方式：

// 第一次调用：确定临时存储大小
cub::DeviceScan::ExclusiveSum(nullptr, temp_storage_bytes, d_input, d_output, num_items);

// 分配所需的临时存储
cudaMalloc(&d_temp_storage, temp_storage_bytes);

// 第二次调用：执行实际的扫描算法
cub::DeviceScan::ExclusiveSum(d_temp_storage, temp_storage_bytes, d_input, d_output, num_items);

这种两阶段API的一个显著缺点是，在估算和执行步骤之间，哪些参数必须保持一致并不清晰。从上面的代码片段来看，哪些参数影响内部状态、哪些可以在两次调用之间改变，在编程上并不明确，因为两个阶段的函数签名完全相同。例如，d_input 和 d_output 参数实际上只在第二次调用时使用。

尽管存在这些复杂性，但传统设计有其根本目的：通过将内存分配与执行分离，用户可以分配一块内存并多次重用，甚至在不同算法之间共享。

虽然这一设计对一部分用户很重要，但充分利用此功能的整体用户群相当有限。因此，许多用户选择封装CUB调用，将两步调用抽象为单次调用，并提供自动内存管理。PyTorch就是一个例子，它使用宏来封装CUB调用。

以下是来自pytorch/pytorch GitHub仓库的源代码：

// 处理临时存储和CUB API的"两次"调用
#define CUB_WRAPPER(func, ...) do {                                       \
  size_t temp_storage_bytes = 0;                                          \
  AT_CUDA_CHECK(func(nullptr, temp_storage_bytes, __VA_ARGS__));          \
  auto& caching_allocator = *::c10::cuda::CUDACachingAllocator::get();    \
  auto temp_storage = caching_allocator.allocate(temp_storage_bytes);     \
  AT_CUDA_CHECK(func(temp_storage.get(), temp_storage_bytes, __VA_ARGS__));\
} while (false)

使用宏也有其缺点，它会模糊控制流和参数传递，导致代码难以理解和调试。

新的单次调用CUB API

鉴于许多生产级代码库中广泛使用封装，CUB引入了新的单次调用API：

// 单次调用：分配和执行在一步中完成
cub::DeviceScan::ExclusiveSum(d_input, d_output, num_items);

示例显示，不再需要显式内存分配。但请注意，分配过程仍在底层进行。图1显示，与两阶段API相比，单次调用接口（包括临时存储估算、内存分配和调用算法）引入了零开销。

图1. 单次与两阶段算法的ExclusiveSum耗时对比

图1比较了原始两阶段ExclusiveSum调用与新引入的单次调用的GPU运行时间。从性能数据可以得出两个主要结论：

  1. 新API引入了零开销。
  2. 内存分配在新API下仍然发生，只是在底层自动进行。

第二点可以通过查看新API的内部实现来验证。异步分配被嵌入到设备原语中：

cub::DeviceScan::ExclusiveSum(d_input, d_output, num_items, env = {}) {
    . . .
    d_temp_storage = mr.allocate(stream, bytes);
    mr.deallocate(stream, d_temp_storage, bytes);
    . . .
}

两阶段API并未被移除，它们仍然是有效的CUB API调用。单次调用API是在现有API之上新增的，预计大多数用户将转向使用这些新接口。

环境与内存资源

除了解决上述问题，新的单次调用CUB API还扩展了所调用原语的执行配置能力。它引入了一个环境参数，既可以用于定制内存分配，也可以仅用于指定执行的CUDA流。

内存资源是一种用于分配和释放内存的新工具。单次调用API的环境参数可以可选地包含一个内存资源。当未通过环境参数提供内存资源时，API将使用CCCL提供的默认内存资源。相反，可以选择传递CCCL代码库中提供的非默认内存资源，甚至传递自定义的内存资源。

// 使用CCCL提供的内存资源类型
cuda::device_memory_pool mr{cuda::devices[0]};
cub::DeviceScan::ExclusiveSum(d_input, d_output, num_items, mr);

// 创建并使用自定义内存资源
my_memory_resource my_mr{cuda::experimental::devices[0]};
// 与CUB一起使用
cub::DeviceScan::ExclusiveSum(d_input, d_output, num_items, my_mr);

使用新API，CUDA流的处理并没有被简化，而是被封装到了新的环境变量中。当然，即使临时内存分配的处理被移除，仍然可以像以前一样显式传递流。CUB现在也提供了类型安全的 cuda::stream_ref，应优先使用。也可以传递拥有底层执行流的 cuda::stream。

组合执行选项

单次调用API不仅能传递内存资源或流作为最后一个参数。未来，环境参数将成为所有执行相关配置的中心，包括确定性要求、保证、用户自定义调优等。

通过引入单次调用API，CUB解锁了大量执行配置特性。面对众多新特性，如何最佳地组合它们成为一个问题。

解决方案在于新的env参数。通过利用cuda::std::execution，CUB提供了一个中心端点，作为算法的灵活“控制面板”。env参数使您能够创建所需任何特性的组合，而不是使用固定定义的函数参数。无论您是想将自定义流与特定内存池配对，还是将严格的确定性要求与自定义调优策略结合，env参数都能在一个类型安全的对象中处理所有情况。

cuda::stream custom_stream{cuda::device_ref{0}};
auto memory_prop  = cuda::std::execution::prop{cuda::mr::get_memory_resource,
          cuda::device_default_memory_pool(cuda::device_ref{0})};
auto env = cuda::std::execution::env{custom_stream.get(), memory_prop};
DeviceScan::ExclusiveSum(d_input, d_output, num_items, env);

CUB目前提供以下支持环境接口的算法，后续将支持更多算法：

  • cub::DeviceReduce::Reduce
  • cub::DeviceReduce::Sum
  • cub::DeviceReduce::Min/Max/ArgMin/ArgMax
  • cub::DeviceScan::ExclusiveSum
  • cub::DeviceScan::ExclusiveScan

关于基于环境的新重载的最新进展，请参阅NVIDIA/cccl GitHub仓库中的CUB设备原语跟踪议题。

开始使用CUB

通过将冗长的两阶段模式替换为精简的单次调用接口，CUB提供了一个现代化的API，消除了样板代码且不增加开销。通过利用可扩展的env参数，您将获得一个统一的控制面板，可以无缝地混合使用内存资源、流和其他工具。鼓励开发者采用这一新标准来简化代码库，并充分利用GPU的计算能力。下载CUDA 13.1或更高版本，开始使用这些单次调用API。FINISHED

avatar
文章
阅读
粉丝