什么是 numba.cuda.local.array() 的有效替代方案，它们不像通过 to_device() 传递许多参数那么麻烦？

Posted 2023-02-19

【中文标题】什么是 numba.cuda.local.array() 的有效替代方案，它们不像通过 to_device() 传递许多参数那么麻烦？

【英文标题】：What are efficient alternatives to numba.cuda.local.array() that aren't as cumbersome as passing many arguments via to_device()?

【发布时间】：2021-10-22 11:44:20

【问题描述】：

cuda.local.array()

在 How is performance affected by using numba.cuda.local.array() compared with numba.cuda.to_device()? 中，简单快速排序算法的基准测试表明，使用 to_device 传递预分配数组的效率可以提高约 2 倍，但这需要更多内存。

单独排序 2,000,000 行每行包含 100 个元素的基准测试结果如下：

2000000
Elapsed (local: after compilation) = 4.839058876037598
Elapsed (device: after compilation) = 2.2948694229125977
out is sorted
Elapsed (NumPy) = 4.541851282119751

使用to_device() 的虚拟示例

如果您有一个复杂的程序，它有许多 cuda.local.array() 调用，等效的 to_device 版本可能会开始看起来像这样并且变得相当麻烦：

def foo2(var1, var2, var3, var4, var5, var6, var7, var8, var9, var10, out):
    for i in range(len(var1)):
        out[i] = foo(var1, var2, var3, var4, var5, var6, var7, var8, var9, var10, out)

def foo3(var1, var2, var3, var4, var5, var6, var7, var8, var9, var10, out):
    idx = cuda.grid(1)
    foo(var1, var2, var3, var4, var5, var6, var7, var8, var9, var10, out[idx])

在真实的代码库中，可能有 3-4 级函数嵌套，跨越数十个函数和数百到数千行代码。 这两种方法的替代方法是什么？

【参考方案1】：

替代方案

以下是cuda.local.array() 和通过cuda.to_device() 单独传递参数的一些替代方案：

分配一个串联的向量/矩阵（例如local_args），它实际上代表了 15 个变量。这样做的缺点是需要不断对其进行切片，并希望您不会意外使用来自不同“子变量”的索引或通过稍后添加新变量、更改大小等来破坏排序。 将操作拆分为顺序调用的 Numba/CUDA 内核，或 Numbacuda.jit()、CuPycupy.fuse() 调用和/或其他 CUDA 实现的组合。例如，如果您对一组向量进行操作，否则这些操作将在成对距离矩阵计算中（例如 10,000² 而不是 10,000 次）重复多次（昂贵且冗余），然后考虑预先执行这些操作并将它们作为参数传入（可以与 1. 或 3. 结合使用） 我遇到的一个方便的替代方法是define a custom NumPy dtype，尽管这可能会导致issues with the NVCC compiler（希望永久修复）。 GitHub issue 有一个例子如下：

import numpy as np
np_int = np.int32
np_float = np.float32
cuda_const_arrays_type = np.dtype([
('a1', (np_int,(7776, 13))),
('a2', (np_int,(7776, 2, 5))),
('a3', (np_int,(16494592))),
('a4', (np_int,13)),
('a5', (np_float,(22528, 64))),
('a6', (np_int,(522523, 64))),
('a7', (np_int,(32,5))),
('a8', (np_int,(66667))),
('a9', (np_int,(252, 64, 3, 2, 2, 2, 2, 2, 2, 13))),
('a10', (np_int,(7776)))
])
cuda_const_arrays = np.zeros(1, dtype=cuda_const_arrays_type)
for txt in cuda_const_arrays_type.names: # i.e. ("a1", "a2", ...)
    cuda_const_arrays[0][txt] = np.loadtxt(open(txt+".csv", "rb"), delimiter=",", skiprows=1)
gpu_const_arrays = cuda.to_device(cuda_const_arrays[0])

@cuda.jit(device=True)
def cuda_doSomething(gpu_const_arrays,...):
    gpu_const_arrays.a1

可以在Gitlab 上找到来自同一用户的示例（可以删除import keras as ks 行）。虽然这会导致以前的 Numba 版本出现零星错误，但它在 numba 0.53.1 和 cudatoolkit 11.2.2 上运行良好，表明“自定义 dtype”方法可能是 OK now。

为了防止不必要地将大量数据传递给堆栈跟踪中较低的函数，在此自定义 dtype 中仅传递参数的子集可能是合适的，但我不知道该怎么做这个。

其他一般有用的例子

虽然我们正在等待 CuPy 或 NumPy 支持 Numba/CUDA 7 9 10 11，但以下是我发现在编写 Numba/CUDA 脚本的工作流程中相关/有用的示例.

Why numba cuda is running slow after recalling it several times? accelerated FFT to be invoked from Python Numba CUDA kernel Numba Discourse: Optimizing Code Further, CUDA Jit?（Graham Markall 提供了很好的建议和示例） Cuda Optimize Jaro Distance（Graham Markall 的实施示例和解释） Numba convolutions 和 user's implementations in NumPy, CuPy, and Numba How to generalize fast matrix multiplication on GPU using numba（扩展/更正 Numba Docs matmul 示例）

其中一些示例非常好，因为您可以看到原始的低效方法以及如何对其进行修改以提高效率，类似于 Numba Docs: CUDA: Matrix Multiplication 示例，并了解其他人如何处理数组分配和 Numba 中的参数传递/ CUDA。