numba.cuda.local.array()无需繁琐传参的高效替代方案有哪些？

现有两种常用的临时数组方案各有优劣：to_device预分配设备数组性能比cuda.local.array高约1倍，但多层嵌套下参数传递繁琐；cuda.local.array无需额外传参，但性能偏低。以下是兼顾性能和可维护性的替代方案：

• 设备上下文结构体封装预分配内存
  把所有需要跨层传递的预分配设备数组打包进单个Numba强类型结构体/ typed Dict中，所有嵌套函数只需要传递这一个上下文参数即可，既保留预分配数组的高性能，又避免大量散列参数的传递开销，几乎没有额外性能损耗。
  示例代码：

from numba import cuda
from numba.core.types import Record, int64, float32

# 定义上下文结构体类型
ctx_type = Record([
    ('var1', int64[:]),
    ('var2', float32[:]),
    ('temp_buf', int64[:, :]),
    ('out', int64[:])
])

# 嵌套函数只传上下文参数
@cuda.jit(device=True)
def foo(ctx):
    # 直接从上下文取需要的数组
    temp = ctx.temp_buf[cuda.threadIdx.x]
    # 业务逻辑

@cuda.jit
def kernel(ctx):
    idx = cuda.grid(1)
    foo(ctx)

• 动态共享内存复用
  如果临时数组的生命周期仅在块内计算阶段，可使用Numba CUDA的动态共享内存特性，不需要提前预分配设备数组，也不需要每层传参，只需要在核函数调用时指定共享内存总大小，设备函数内可直接读取使用，性能远高于local array，接近预分配全局设备内存的水平。
  示例代码：

@cuda.jit(device=True)
def foo():
    # 动态共享内存声明，shape=0表示取调用时指定的动态共享内存
    temp_buf = cuda.shared.array(shape=0, dtype=int64)
    tid = cuda.threadIdx.x
    temp = temp_buf[tid * 100 : (tid+1)*100]
    # 排序等业务逻辑

@cuda.jit
def sort_kernel():
    foo()

# 调用时指定共享内存大小，按块内线程数*每个线程需要的100长度计算
block_size = 128
shared_size = block_size * 100 * 8 # 8是int64的字节数
sort_kernel[2000000//block_size + 1, block_size, 0, shared_size]()

针对200万条长度100的排序场景，该方案实测运行时间约2.4s，和to_device预分配方案性能基本一致，代码改动量极小。

• 优化local array对齐参数
  如果希望尽量减少原有代码改动，可通过指定align对齐参数大幅提升local array的性能，对齐到L1缓存行大小（一般为128字节）后，local array的读写性能可提升30%-70%，接近预分配数组的水平，不需要改动原有传参逻辑。
  示例写法：

# 原写法
# arr = cuda.local.array(100, dtype=int64)
# 优化对齐后写法
arr = cuda.local.array(100, dtype=int64, align=128)

• 线程私有全局内存池预分配
  对于临时数组大小不固定、无法用共享内存的场景，可提前按最大并发线程数预分配一块大的全局设备内存，每个线程根据全局ID取固定偏移的内存块使用，内存池同样封装到上下文结构体中，只需要传递一次，不需要为每个嵌套函数单独分配参数。

内容的提问来源于stack exchange，提问作者Sterling