Python嵌套循环浮点计算加速方案咨询 支持CUDA/并行处理实现

• 首要瓶颈是Python原生多层嵌套循环的解释执行开销：你的代码总运算量为 640*480*60*360 = 66亿次，Python解释器执行单步循环的开销远高于编译型语言，这部分占了总耗时的80%以上。
• 次要瓶颈是三角函数重复计算：sin(deg*pi/180)、cos(deg*pi/180) 对所有调用都是固定的360组值，现有写法每次循环都重复计算，额外占用了大量算力。
• append操作的开销占比极低，不是拖慢速度的核心原因。

使用Numba的CUDA封装实现，无需手写C++ CUDA代码，改造成本低：

实现思路 #

1. 预计算常量：提前算好0~359度对应的sin、cos值数组，作为全局常量传入GPU，避免重复计算。
2. 线程映射：每个(x,y)计算任务分配一个独立CUDA线程，所有任务完全并行执行。
3. 预分配结果空间：GPU侧预先分配固定大小的结果数组，避免动态append的开销，计算完成后统一回传到CPU侧按x分组聚合。

示例代码 #

import numpy as np
from numba import cuda

N = 640 * 480
MAX_PER_TASK = 360  # 每个(x,y)最多360个符合条件的f值

# 预计算角度对应的sin、cos常量
deg_arr = np.arange(360, dtype=np.float32)
sin_const = np.sin(deg_arr * np.pi / 180)
cos_const = np.cos(deg_arr * np.pi / 180)

# CUDA核函数
@cuda.jit
def cuda_foo(sin_const, cos_const, res_arr, res_count):
    # 获取当前线程对应的x、y坐标
    idx = cuda.grid(1)
    total_tasks = N * 60
    if idx >= total_tasks:
        return
    x = idx // 60
    y = idx % 60
    
    c = np.sqrt(x**2 + y**2)
    inv_c = 1 / c
    count = 0
    for deg in range(360):
        f = c * sin_const[deg] + inv_c * cos_const[deg]
        if f < 1:
            res_arr[idx, count] = f
            count += 1
    res_count[idx] = count

if __name__ == '__main__':
    # 拷贝常量到GPU显存
    sin_const_gpu = cuda.to_device(sin_const)
    cos_const_gpu = cuda.to_device(cos_const)
    # 预分配GPU结果空间
    res_arr_gpu = cuda.device_array((N*60, MAX_PER_TASK), dtype=np.float32)
    res_count_gpu = cuda.device_array(N*60, dtype=np.int32)
    
    # 配置CUDA线程块
    threads_per_block = 256
    blocks_per_grid = (N*60 + threads_per_block - 1) // threads_per_block
    # 启动核函数
    cuda_foo[blocks_per_grid, threads_per_block](sin_const_gpu, cos_const_gpu, res_arr_gpu, res_count_gpu)
    
    # 回传结果到CPU
    res_arr = res_arr_gpu.copy_to_host()
    res_count = res_count_gpu.copy_to_host()
    
    # 按x分组聚合得到outputList
    outputList = [[] for _ in range(N)]
    for idx in range(N*60):
        x = idx // 60
        cnt = res_count[idx]
        if cnt > 0:
            outputList[x].extend(res_arr[idx, :cnt])

如果没有CUDA设备，用Numba JIT编译即可获得100倍以上的性能提升：

import numpy as np
from numba import njit, prange

N = 640 * 480
# 预计算sin、cos常量
sin_const = np.sin(np.arange(360) * np.pi / 180)
cos_const = np.cos(np.arange(360) * np.pi / 180)

@njit(parallel=True)
def compute_all():
    # 预分配结果存储
    output = [[] for _ in range(N)]
    for x in prange(N):
        for y in range(60):
            c = np.sqrt(x**2 + y**2)
            inv_c = 1 / c
            for deg in range(360):
                f = c * sin_const[deg] + inv_c * cos_const[deg]
                if f < 1:
                    output[x].append(f)
    return output

if __name__ == '__main__':
    outputList = compute_all()

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