Python циклов for для массивов 3D Cupy на GPU, когда широковещательная передача массива невозможна - PullRequest
Новая
       86

Python циклов for для массивов 3D Cupy на GPU, когда широковещательная передача массива невозможна

0 голосов
/ 04 августа 2020

Я пытаюсь использовать Cupy с GPU (NVIDIA K6000 с 2880 ядрами) для ускорения некоторых циклов for над несколькими трехмерными массивами в Python; как показано в примере кода ниже. Я определяю функцию, которая принимает трехмерные массивы в качестве входного аргумента, выполняет циклы for и выводит массив. Я запускаю два случая и сравниваю время выполнения: (1) jit-скомпилированная функция, содержащая for-l oop, обрабатывающая два массива numpy (A_cpu, B_cpu), параллельно на 4 ядрах ЦП, и (2) for-l *1012* l oop вызывается на двух массивах Cupy (A_gpu, B_gpu) на 2880 ядрах графического процессора. Я получаю одинаковые значения в каждом случае; тем не менее, выполнение на графическом процессоре типа cupy значительно медленнее, чем выполнение на CPU numpy; как показано в результатах ниже. Я не могу использовать широковещательную рассылку массива (A [i,:,:]) для исключения любого из циклов for {i, j, k}, потому что индексация массива не является единообразной для разных операций.

I следует добавить, что математические операции в реальном коде намного сложнее и включают около 12 различных трехмерных массивов; однако для краткости здесь я просто показал взаимозависимость индексации массива. Кроме того, запуск фактического кода с @jit (parallel = True) значительно ускоряет код на 4-ядерном ЦП по сравнению с непараллельным jit, поэтому я подозреваю, что он должен поддаваться еще более массивному распараллеливанию на 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

Любые идеи или советы приветствуются заранее !

...