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

Question

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 层函数嵌套在数十个函数和数百到数千行代码中。 这两种方法的替代方案是什么？

Answer 1

备择方案
这里有一些替代方案 cuda.local.array()并通过 cuda.to_device() 单独传递参数:

分配一个连接的向量/矩阵(称为例如 local_args )，它实际上代表了 15 个变量。这样做的缺点是需要不断地对其进行切片，并希望您不会意外地使用来自不同“子变量”的索引或通过稍后添加新变量、更改大小等来破坏排序。

将操作拆分为顺序调用的 Numba/CUDA 内核，或 Numba 的组合 cuda.jit() , CuPy cupy.fuse() 调用和/或其他 CUDA 实现。例如，如果您对一组向量进行操作，否则这些操作将在成对距离矩阵计算中(例如 10,0002 次而不是 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 ，表明“自定义数据类型”方法可能是 OK now .
为了防止将大量数据不必要地传递到堆栈跟踪中较低的函数，在此自定义 dtype 中仅传递参数的子集可能是合适的。 ，但我不知道如何做到这一点。
其他一般有用的例子
我们正在等待 Numba/CUDA 的 CuPy 或 NumPy 支持 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 中的数组分配和参数传递。