如何用CuPy计算CUDA协作组启动时每SM最大活动块数

CUDA协作组启动必须确保所有线程瞬时启动，才能使网格同步正常工作。以下是CUDA C++中计算每SM最大活动块数并启动协作内核的示例代码：

/// 此代码将在默认流上启动一个能最大化利用GPU的网格，并传入内核参数
int numBlocksPerSm = 0;
// my_kernel启动时使用的线程数
int numThreads = 128;
cudaDeviceProp deviceProp;
cudaGetDeviceProperties(&deviceProp, dev);
cudaOccupancyMaxActiveBlocksPerMultiprocessor(&numBlocksPerSm, my_kernel, numThreads, 0);
// 启动内核
void *kernelArgs[] = { /* 添加内核参数 */ };
dim3 dimBlock(numThreads, 1, 1);
dim3 dimGrid(deviceProp.multiProcessorCount*numBlocksPerSm, 1, 1);
cudaLaunchCooperativeKernel((void*)my_kernel, dimGrid, dimBlock, kernelArgs);

针对你希望在CuPy中实现上述cudaOccupancyMaxActiveBlocksPerMultiprocessor功能的需求，目前CuPy没有直接封装该高层API，但可以通过以下两种方式实现：

方法一：调用CuPy封装的CUDA Runtime底层接口 #

CuPy提供了对CUDA Runtime API的直接封装，可以直接调用cudaOccupancyMaxActiveBlocksPerMultiprocessor：

import cupy as cp
from cupy.cuda.runtime import cudaOccupancyMaxActiveBlocksPerMultiprocessor

# 获取当前设备
device = cp.cuda.Device()

# 定义块大小
block_size = 256

# 编译示例内核
kernel_code = '''
extern "C" __global__ void mykernel(float *x, float *y, float *z) {
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    z[tid] = x[tid] + y[tid];
}
'''
module = cp.RawModule(code=kernel_code, backend='nvcc')
kernel = module.get_function('mykernel')

# 获取内核函数指针
kernel_ptr = kernel.ptr

# 计算每SM最大活动块数
num_blocks_per_sm = cudaOccupancyMaxActiveBlocksPerMultiprocessor(kernel_ptr, block_size, 0)

# 打印结果
print(f'每SM的最大活动块数：{num_blocks_per_sm}')

方法二：手动基于设备属性计算 #

通过获取设备SM的资源限制，结合内核的资源消耗手动计算最大块数：

import cupy as cp

# 获取当前设备及属性
device = cp.cuda.Device()
device_props = device.attributes

# 定义块大小
block_size = 256

# 编译示例内核并获取资源消耗
kernel_code = '''
extern "C" __global__ void mykernel(float *x, float *y, float *z) {
    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    z[tid] = x[tid] + y[tid];
}
'''
module = cp.RawModule(code=kernel_code, backend='nvcc')
kernel = module.get_function('mykernel')

regs_per_thread = kernel.attributes['num_regs']
shared_mem_per_block = kernel.shared_size_bytes

# 设备SM资源限制
max_threads_per_sm = device_props['maxThreadsPerMultiProcessor']
max_regs_per_sm = device_props['regsPerMultiprocessor']
max_shared_mem_per_sm = device_props['sharedMemPerMultiprocessor']

# 分别计算不同资源限制下的最大块数
block_limit_threads = max_threads_per_sm // block_size
block_limit_regs = max_regs_per_sm // (regs_per_thread * block_size)
block_limit_shared = max_shared_mem_per_sm // (shared_mem_per_block if shared_mem_per_block != 0 else 1)

# 取最小值作为最终结果
num_blocks_per_sm = min(block_limit_threads, block_limit_regs, block_limit_shared)

print(f'每SM的最大活动块数：{num_blocks_per_sm}')

内容的提问来源于stack exchange，提问作者Marko Grdinić