如何为CuPy Raw Kernel传递动态分配共享内存大小参数？

正确传递动态共享内存大小的方式 #

CuPy中调用Raw Kernel时，动态共享内存的大小是通过调用接口的第三个位置参数（或shared_mem关键字参数）指定的，不需要将其作为Kernel的输入参数传递。

实际调用示例：

import cupy as cp

# 定义带动态共享内存的Raw Kernel
add_kernel = cp.RawKernel(r'''
extern "C" __global__ void add(int* x, int* y, int* z) {
    extern __shared__ int sharedValues[];
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    sharedValues[threadIdx.x] = x[idx];
    __syncthreads();
    z[idx] = sharedValues[threadIdx.x] + y[idx];
}
''', 'add')

# 准备数据
n = 1024
x = cp.arange(n, dtype=cp.int32)
y = cp.ones(n, dtype=cp.int32)
z = cp.zeros(n, dtype=cp.int32)

# 计算所需共享内存大小（按每个block的线程数分配int型空间）
block_size = 256
shared_mem_size = block_size * cp.int32().itemsize

# 调用Kernel：grid维度、block维度、Kernel参数、共享内存大小
add_kernel((n//block_size,), (block_size,), (x, y, z), shared_mem_size)
# 或用关键字参数更直观
# add_kernel((n//block_size,), (block_size,), (x, y, z), shared_mem=shared_mem_size)

两种错误操作的原因 #

• 添加额外参数传递共享内存大小：这是CUDA C++的写法，但CuPy的Raw Kernel调用接口中，共享内存大小是独立于Kernel参数的特殊配置项，不能放到Kernel的参数列表里，否则会被当成普通输入参数，导致参数不匹配报错。
• 修改shared_size_bytes属性：该属性是Kernel对象的只读属性，仅用于查看当前实例关联的共享内存大小，无法直接修改，必须通过调用时的参数指定。

关于max_dynamic_shared_size_bytes的说明 #

add_kernel.max_dynamic_shared_size_bytes是只读属性，它返回当前CUDA硬件支持的最大动态共享内存容量（单位字节），这个值由硬件和CUDA版本决定，无法手动修改。你实际使用的动态共享内存大小只要不超过这个最大值即可，具体大小仍通过调用Kernel时的shared_mem参数指定。

内容的提问来源于stack exchange，提问作者Uwe.Schneider