摘要：最近工作时出现了需要对稀疏矩阵的非零元素进行频繁的增删改查的场景。为了提高对稀疏矩阵的操作的速度，一个可行的做法就是使用哈希表构建稀疏矩阵，使用键 Key 和值 Value 分别存储矩阵的行列 index 与非零元的实际值。一般在 CPU 上实现的话，C++
最近工作时出现了需要对稀疏矩阵的非零元素进行频繁的增删改查的场景。为了提高对稀疏矩阵的操作的速度，一个可行的做法就是使用哈希表构建稀疏矩阵，使用键 Key 和值 Value 分别存储矩阵的行列 index 与非零元的实际值。一般在 CPU 上实现的话，C++ 标准库的 std::unordered_map 是个不错的选择。然而到了老黄家的 GPU 上的 cuda 编程上，想使用能支持并行处理的 on-device 的哈希表就没有那么方便了。目前用 CUDA 开发的 GPU 上的哈希表比较 sota 的也有几个，我最终选择了 cuCollections 进行开发。这篇文章主要介绍一下我最近 cuCollections 的使用心得。
1. cuCollection 简介cuCollections（git repo：）是一个开源的 CUDA 库，专为 GPU 上的高性能哈希表操作而设计，旨在提供高效且易用的哈希表实现。它有如下几个特点：
高性能（√）： 利用 GPU 的并行计算能力，cuCollections 能够显著加速哈希表操作，例如插入、查找和删除。易用性（其实也没有那么易用）： 提供简洁的 API，与 C++ STL 风格类似，方便开发者快速上手。内存效率： 采用高效的内存管理机制，最大限度地利用 GPU 内存资源。cuCollections 专注于哈希表操作，提供了较为丰富的功能和基于 NVIDIA GPU 的高性能。目前还在很密集地更新与维护，还有相当一部分功能正在开发（比如一些 kernel 函数中的 device reference）。可以期待之后 CUDA 将会有更强大的哈希表的支持。这篇文章基于 cuCollections 截止到 e3aec27 commit 的版本进行讨论。
2. 使用方法2.1. 安装
对 cuCollections 的依赖安装比较简单。它是一个头文件库，依赖 nvcc 11.5+ 编译 （用于 Volta 及以上的 GPU 架构），使用 C++ 17 标准。推荐使用 CMake Package Manager (CPM) 搭建环境：
cmake_minimum_required(VERSION 3.23.1 FATAL_ERROR)include(path/to/CPM.cmake)CPMAddPackage( NAME cuco GITHUB_REPOSITORY NVIDIA/cuCollections GIT_TAG dev OPTIONS "BUILD_TESTS OFF" "BUILD_BENCHMARKS OFF" "BUILD_EXAMPLES OFF")target_link_libraries(my_library cuco)2.2. On-Host 操作
哈希表的初始化/键值对插入等 bulk 级别的操作都可以通过 host API 完成。
例如构建一个静态容量的哈希表 cuco::static_map:
在这里，capacity 是你所需的容量，empty_key_sentinel/empty_value_sentinel 是占位值，一般使用不会在待插入的键值对中出现的数字作为占位值。冲突解决使用的是 linear probing，cuco 目前并不支持太多其他 collision resolutions （比如 separate chaining）。初始化好一个 static map 后，插入键值对：
// Combine keys and values into pairs {{0,0}, {1,1}, ... {i,i}}auto pairs = thrust::make_transform_iterator( thrust::counting_iterator{0}, cuda::proclaim_return_type>( [keys = insert_keys.begin, values = insert_values.begin] __device__(auto i) { return cuco::pair{keys[i], values[i]}; }));// Inserts all pairs into the mapmap.insert(pairs, pairs + num_keys);使用 thrust 的 make_transform_iterator 将作为输入的 on-device 的 key array (insert_key.begin) 和 value array (insert_value.begin) 转变成 on-device 的键值对，然后通过 insert 可以插入 static map。如果想将 map 中存在的键值对再提取到 array 中，可以使用 retrieve：
map.retrieve_all(keys.begin, values.begin);2.3. On-Device 操作 (kernel)
对于构建好的 cuCollections 哈希表，对其增删改查某些键值对的操作可以通过在 kernel function 里使用 on-device reference 完成。一些使用 on-device reference 的例子可以在 cuCollections git 库里提供的 examples 中找到，比如 。在这里我展示一个简单的 kernel 的写法：在 kernel function 中查找某一键值对并更新其值。
template __global__ void upd_hash_val(Map map_ref, KeyIter hash_key, ValIter upd_val, int num_keys) { // Thread numberint tid = threadIdx.x + blockIdx.x * blockDim.x; while (tid {found->second}; ref.fetch_add(upd_val[tid], cuda::memory_order_relaxed); } // Next gridtid += blockDim.x * gridDim.x; }}在这个 kernel function 中，我们从哈希表中寻找输入的 hash_key 中的每一个键，如果某个线程找到了对应 hash_key[tid] 的键，那么对其值进行累加（fetch_add），累加值是输入的 upd_val[tid]。要注意，map_ref 是一个我们在 host 中提前构造的一个 cuco 哈希表的 find reference（如何构造详见上面那个链接）。也就是说，cuco 的哈希表是不作为 kernel function 的输入的，作为替代，是专门负责某项功能的（insert/find/...）reference 作为输入在 on-device 的函数上进行操作。
另外值得一提的是，目前 cuCollections 还在活跃开发维护，on-device reference 其实还支持的不是很好。目前试下来，static_map 是最完善的。其他的哈希表像是 multi/dynamic...之类的目前对 on-device 操作支持还不是很好。并且根据我尝试的经验来看，目前两个 reference 同时在一个 kernel function 中使用不保证原子性，比如说 insert 和 find 的 reference 同时使用有可能造成结果正确性问题。关于这一点，我还没有做详细验证，欢迎大家讨论。
性能经过一些简单实验，我发现 cuCollections 确实在 GPU 上迸发出不俗的并行性能。不过不同的 map ...
来源：鼠meme
