Можно использовать numba.guvectorize для эмуляции параллельного forall/prange?

Question

Как пользователь Python для анализа данных и численных расчетов, а не реального «кодера», мне не хватало действительно низкозатратного способа распределения неловко параллельных вычислений на нескольких ядрах. Как я узнал, в Numba раньше была конструкция prange, но это было abandoned because of "instability and performance issues".

Играя с newly open-sourced @guvectorize decorator Я нашел способ использовать его для практически без накладных эмуляции функциональности позднего prange.

Я очень рад, что теперь этот инструмент оказался под рукой, благодаря ребятам из Continuum Analytics, и ничего не нашел в Интернете, явно упоминая об этом использовании @guvectorize. Хотя это может быть тривиальным для людей, которые раньше использовали NumbaPro, я публикую это для всех тех не-кодировщиков (см. Мой ответ на этот «вопрос»).

Answer 1

Рассмотрим приведенный ниже пример, где двухуровневый цикл вложенного цикла с ядром, выполняющим численный расчет с использованием двух входных массивов и функцией индексов цикла, выполняется четырьмя различными способами. Каждый вариант приурочен %timeit магии IPython в:

1. наивным цикл, составлен с использованием numba.jit
2. FORALL-подобной конструкции с использованием numba.guvectorize, выполненных в одном потоке (target = "cpu")
3. FORALL-подобный построить с использованием numba.guvectorize, выполненных в столько же потоков, что и «core» ядра процессора (в моем случае гиперпотоки) (target = "parallel")
4. то же, что и 3., однако называя «guvectorized» forall последовательностью «параллельного» цикла индексы случайным образом переставляются

Последнее сделано, потому что (в данном конкретном примере) диапазон внутреннего цикла зависит от значения индекса внешнего контура. Я не знаю, как именно диспетчеризация вызовов gufunc организована внутри numpy, но похоже, что рандомизация «параллельных» индексов цикла обеспечивает немного лучшую балансировку нагрузки.

На моих (медленных) машинах (первое поколение Core i5, 2 ядра, 4 hyperthreads) Я получаю тайминги:

1 loop, best of 3: 8.19 s per loop 
1 loop, best of 3: 8.27 s per loop 
1 loop, best of 3: 4.6 s per loop 
1 loop, best of 3: 3.46 s per loop 

Примечания: мне было бы интересно, если этот рецепт легко относится к целевому =» gpu "(он должен делать, но сейчас у меня нет доступа к подходящей графической карте), и что такое ускорение. Сообщение

И вот пример:

import numpy as np 
from numba import jit, guvectorize, float64, int64 

@jit 
def naive_for_loop(some_input_array, another_input_array, result): 
    for i in range(result.shape[0]): 
     for k in range(some_input_array.shape[0] - i): 
      result[i] += some_input_array[k+i] * another_input_array[k] * np.sin(0.001 * (k+i)) 

@guvectorize([(float64[:],float64[:],int64[:],float64[:])],'(n),(n),()->()', nopython=True, target='parallel') 
def forall_loop_body_parallel(some_input_array, another_input_array, loop_index, result): 
    i = loop_index[0]  # just a shorthand 
    # do some nontrivial calculation involving elements from the input arrays and the loop index 
    for k in range(some_input_array.shape[0] - i): 
     result[0] += some_input_array[k+i] * another_input_array[k] * np.sin(0.001 * (k+i)) 

@guvectorize([(float64[:],float64[:],int64[:],float64[:])],'(n),(n),()->()', nopython=True, target='cpu') 
def forall_loop_body_cpu(some_input_array, another_input_array, loop_index, result): 
    i = loop_index[0]  # just a shorthand 
    # do some nontrivial calculation involving elements from the input arrays and the loop index 
    for k in range(some_input_array.shape[0] - i): 
     result[0] += some_input_array[k+i] * another_input_array[k] * np.sin(0.001 * (k+i)) 

arg_size = 20000 

input_array_1 = np.random.rand(arg_size) 
input_array_2 = np.random.rand(arg_size) 
result_array = np.zeros_like(input_array_1) 

# do single-threaded naive nested for loop 
# reset result_array inside %timeit call 
%timeit -r 3 result_array[:] = 0.0; naive_for_loop(input_array_1, input_array_2, result_array) 
result_1 = result_array.copy() 

# do single-threaded forall loop (loop indices in-order) 
# reset result_array inside %timeit call 
loop_indices = range(arg_size) 
%timeit -r 3 result_array[:] = 0.0; forall_loop_body_cpu(input_array_1, input_array_2, loop_indices, result_array) 
result_2 = result_array.copy() 

# do multi-threaded forall loop (loop indices in-order) 
# reset result_array inside %timeit call 
loop_indices = range(arg_size) 
%timeit -r 3 result_array[:] = 0.0; forall_loop_body_parallel(input_array_1, input_array_2, loop_indices, result_array) 
result_3 = result_array.copy() 

# do forall loop (loop indices scrambled for better load balancing) 
# reset result_array inside %timeit call 
loop_indices_scrambled = np.random.permutation(range(arg_size)) 
loop_indices_unscrambled = np.argsort(loop_indices_scrambled) 
%timeit -r 3 result_array[:] = 0.0; forall_loop_body_parallel(input_array_1, input_array_2, loop_indices_scrambled, result_array) 
result_4 = result_array[loop_indices_unscrambled].copy() 


# check validity 
print(np.all(result_1 == result_2)) 
print(np.all(result_1 == result_3)) 
print(np.all(result_1 == result_4))