Python不同线程间如何传递变量？CUDA代码获取最小值失败求助

问题根源 #

• 你当前代码中，pp作为普通Python整数传递给CUDA核函数，是值传递，核函数内部的pp只是副本，修改后不会同步回主机端。
• 要在CUDA线程间协作计算并返回结果，必须使用CUDA能读写的内存空间（全局内存、共享内存），还要处理线程间的同步。

修正方案 #

要实现线程间计算并获取最小值，我们可以用全局内存数组存储中间结果，通过线程同步完成归约计算，最后将结果拷贝回主机。

修正后的代码 #

import numpy as np
from numba import cuda

@cuda.jit
def find_min(result_arr):
    # 获取全局线程索引
    global_idx = cuda.grid(1)
    # 每个线程将自己的索引写入全局内存对应位置（若计算自定义数据，替换此处即可）
    result_arr[global_idx] = global_idx
    
    # 同步所有线程，确保所有线程完成写入操作
    cuda.syncthreads()
    
    # 归约计算最小值（假设线程数为2的幂，适配当前16线程的场景）
    stride = cuda.blockDim.x // 2
    while stride > 0:
        if global_idx < stride:
            if result_arr[global_idx] > result_arr[global_idx + stride]:
                result_arr[global_idx] = result_arr[global_idx + stride]
        cuda.syncthreads()
        stride = stride // 2

if __name__ == '__main__':
    threads_per_block = 16
    blocks_per_grid = 1
    total_threads = threads_per_block * blocks_per_grid
    
    # 分配设备端全局内存数组，用于存储计算结果
    d_result = cuda.device_array((total_threads,), dtype=np.int32)
    
    # 调用CUDA核函数
    find_min[blocks_per_grid, threads_per_block](d_result)
    
    # 将设备端结果拷贝回主机端
    h_result = d_result.copy_to_host()
    
    # 归约后数组的第一个元素即为最小值
    print("最小值为:", h_result[0])

关键说明 #

• 用cuda.device_array创建设备可访问的内存空间，作为线程间共享数据的载体
• 通过归约算法在CUDA线程间协作计算最小值，cuda.syncthreads()用于确保线程间操作同步
• 最终通过copy_to_host()将设备端结果同步回主机，获取最终计算值

若你需要计算的是自定义数据的最小值，只需修改线程写入result_arr的内容，归约逻辑可以直接复用。

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