具有动态分配共享内存的Cupy

Question

考虑以下通过 CUPY 在 python 中使用的 CUDA 内核，来自 此链接

add_kernel = cp.RawKernel(r'''
extern "C" __global__
void my_add(const float* x1, const float* x2, float* y) {
    
    extern __shared__ int sharedValues[];
    
    int tid = blockDim.x * blockIdx.x + threadIdx.x;
    y[tid] = x1[tid] + x2[tid];
}
''', 'my_add')
x1 = cp.arange(25, dtype=cp.float32).reshape(5, 5)
x2 = cp.arange(25, dtype=cp.float32).reshape(5, 5)

n = 1
shared_memory_size = n*64; #*sizeof(int)

# Output Results
y = cp.zeros((5, 5), dtype=cp.float32)

add_kernel((512,), (1024,), (x1, x2, y))  # grid, block and arguments

我想在（cupy raw）CUDA 内核中分配共享内存 但我不知道如何将参数shared_memory_size给add_kernel。 示例取自共享内存 CUDA

add_kernel<<<512,1024,n*sizeof(int)>>>(x1,x2,y);

---

当我尝试使用附加参数调用它时，出现错误

add_kernel((512,), (1024,), (shared_memory_size,), (x1, x2, y))  # grid, block and arguments

当我尝试设置属性时

assign_importance_into_dense_array_kernel_int32.shared_size_bytes = shared_memory_size;

我收到错误

AttributeError: attribute 'shared_size_bytes' of 'cupy._core.raw.RawKernel' objects is not writable

还有参数

add_kernel.max_dynamic_shared_size_bytes

但这会改变动态大小吗？

Answer 1

正如@Robert所说，您在

cp.RawKernel.__call__

方法中指定动态共享内存。根据文档

https://docs.cupy.dev/en/stable/reference/ generated/cupy.RawKernel.html

__call__(self, grid, block, args, *, shared_mem=0)

最后一个（可选）参数可以设置为

n*sizeof(int)

，然后就可以开始了。

...
shared_memory_size = n * sizeof(int)
add_kernel((512,), (1024,), (x1, x2, y), shared_memory_size)