当阵列广播不可能时,Python在GPU上对3D cupy阵列进行循环

2024-03-28 18:05:44 发布

您现在位置:Python中文网/ 问答频道 /正文

我正在尝试使用cupy和GPU(NVIDIA K6000,拥有2880个内核)来加速Python中几个三维数组上的for循环;如下面的示例代码所示。我定义了一个函数,该函数将3D数组作为输入参数,执行for循环,并输出数组。我运行了两种情况并比较了执行时间:(1)jit编译的函数包含for循环,在4个cpu核上并行处理两个numpy数组(A_cpu,B_cpu),以及(2)在2880个gpu核上的两个cupy数组(A_gpu,B_gpu)上调用for循环。我从每个案例中得到相同的值;然而,cupy GPU的执行速度明显慢于numpy CPU的执行速度;如下面的结果所示。我无法使用数组广播(A[I,:,:])来消除任何{I,j,k}for循环,因为数组索引在各个操作中是不统一的

我应该补充的是,实际代码中的数学运算要复杂得多,涉及12种不同的3D阵列;然而,为了简洁起见,我在这里展示了数组索引的相互依赖性。此外,与非并行jit相比,使用@jit(parallel=True)运行实际代码可以显著加快4核CPU上的代码速度,这就是为什么我怀疑它应该适合在2880核GPU上进行更大规模的并行化

Python代码示例如下:

import numpy as np
import cupy as cp
from time import time
from numba import cuda, jit, prange

M = 100
N = 100
P = 10
A_cpu = np.random.rand(M, N, P)
B_cpu = np.random.rand(M, N, P)
A_gpu = cp.asarray(A_cpu) 
B_gpu = cp.asarray(B_cpu) 

print('Being CPU test:')
@jit(parallel=True)
def cpu_loop(A_cpu, B_cpu):
    for i in prange(1, M-1):
        for j in prange(1, N-1):
            for k in prange(1, P-1):
                A_cpu[i, j, k] = (A_cpu[i-1, j-1, k+1]*2 - A_cpu[i, j+1, k-1]/2) / (B_cpu[i+1, j+1, k-1]*2 - B_cpu[i, j-1, k+1]/2)
    return(A_cpu)
start = time()
cpu_res = cpu_loop(A_cpu, B_cpu)
end = time()

cpu_total_time = end - start
print('CPU result = ', cpu_res[M//2, N//2, P//2])
print('CPU total_time = {:2.3f} seconds'.format(cpu_total_time))

print('Begin GPU test:')
def gpu_loop(A_gpu, B_gpu):
    for i in prange(1, M-1):
        for j in prange(1, N-1):
            for k in prange(1, P-1):
                A_gpu[i, j, k] = (A_gpu[i-1, j-1, k+1]*2 - A_gpu[i, j+1, k-1]/2) / (B_gpu[i+1, j+1, k-1]*2 - B_gpu[i, j-1, k+1]/2)
    return(A_gpu)
start = time()
gpu_res = gpu_loop(A_gpu, B_gpu)
end = time()

gpu_total_time = end - start
print('GPU result = ', gpu_res[M//2, N//2, P//2])
print('GPU total_time = {:2.3f} seconds'.format(gpu_total_time))    

以上代码的输出如下:

Being CPU test:
CPU result =  208.3083524104198
CPU total_time = 2.117 seconds
Begin GPU test:
GPU result =  208.3083524104198
GPU total_time = 33.462 seconds

任何见解或建议,请提前告知


Tags: 代码inimportforgputime数组cpu