Numba CUDA动态共享内存：能否使用多种数据类型？

直接说结论：不能通过你给出的方式定义两个不重叠的不同类型动态共享内存数组。

因为用cuda.shared.array(0, dtype)声明的动态共享内存，不管你声明多少次，都会指向内核启动时分配的共享内存块的起始地址。也就是说你定义的a和b会完全重叠在同一块内存上，切片操作也没法改变它们的起始位置，只会导致数据互相覆盖。

正确的实现方式 #

要在同一块动态共享内存中分配不同类型的数组，需要手动计算内存偏移，把整块内存划分成不同区域：

步骤1：计算总共享内存大小并启动内核

假设你需要n_int个int32元素和n_float个float32元素，先计算总字节数，再传入内核：

import numba as nb
from numba import cuda

n_int = 256
n_float = 256

# 计算总共享内存字节数
shared_bytes = n_int * nb.int32.itemsize + n_float * nb.float32.itemsize

# 启动内核，传入总共享内存大小
foo_kernel[grid_dim, block_dim, stream, shared_bytes](n_int, n_float)

步骤2：在内核中划分内存区域

可以通过两种方式实现：

方式一：通过无类型字节数组切片转换

@cuda.jit
def foo_kernel(n_int, n_float):
    # 声明无类型的原始共享内存块
    shared_mem = cuda.shared.array(0, nb.uint8)
    
    # 划分int32数组区域
    a_size_bytes = n_int * nb.int32.itemsize
    a = cuda.as_array(shared_mem[:a_size_bytes], dtype=nb.int32)
    
    # 划分float32数组区域（接在int数组之后）
    b_start = a_size_bytes
    b_size_bytes = n_float * nb.float32.itemsize
    b = cuda.as_array(shared_mem[b_start:b_start+b_size_bytes], dtype=nb.float32)
    
    # 示例操作：给a和b赋值
    tid = cuda.threadIdx.x
    if tid < n_int:
        a[tid] = tid
    if tid < n_float:
        b[tid] = tid * 1.5

方式二：通过指针偏移转换

@cuda.jit
def foo_kernel(n_int, n_float):
    # 先声明第一个类型的数组
    a = cuda.shared.array(0, nb.int32)
    
    # 计算第一个数组的字节长度，偏移指针得到第二个数组的起始地址
    a_bytes = n_int * nb.int32.itemsize
    b_ptr = cuda.pointer(a) + a_bytes
    b = cuda.as_array(b_ptr, shape=(n_float,), dtype=nb.float32)
    
    # 后续正常使用a和b即可

注意事项 #

• 必须保证内核启动时传入的shared_bytes等于所有数组的内存总和，否则会出现内存越界，导致未定义行为。
• 划分区域时要注意数据类型的字节对齐，不过Numba的itemsize已经考虑了对齐要求，直接用它计算即可。

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