Как использовать numba parallel?

Question

Я пытаюсь распараллелить или оптимизировать следующий пример кода с помощью numba. Тем не менее, я до сих пор не понимаю логики c за этим.

Параллельная диагностика показывает, что параллельная структура уже оптимальна.

from numba import njit
@njit(parallel=True)
def SwVer(ResSwc,SwNew,Ia,Ja,Ka):
    for k in (Ia):
        for j in (Ja):
            for i in (Ka):
                if(SwNew[i,j,k]<ResSwc):
                    SwNew[i,j,k]=ResSwc
                if(SwNew[i,j,k]>(1-ResSwc)):
                    SwNew[i,j,k]=(1-ResSwc)
    
    return SwNew

import numpy as np
Imax=100
Jmax=100
Kmax=5

Ia=np.arange(0,Imax,1)
Ja=np.arange(0,Jmax,1)
Ka=np.arange(0,Kmax,1)

SwNew=np.random.random((Imax,Jmax,Kmax))
SwNew=SwVer(0.25,SwNew,Ia,Ja,Ka)

Как я могу по-настоящему распараллелить эту функцию? L oop unrolling улучшает время выполнения? Есть ли другой способ улучшить мой тройной л oop?

Спасибо

Answer 1

Q : «Ускоряет ли l oop развертывание время выполнения?»

L oop развертывание можно избежать накладные расходы на цикл, но есть больше проблем для более крупных и многоуровневых циклов, учитывая, что сглаженная «длина» A[N,N,N] скоро вырастет для N выше ~ 1E3, 1E4, 1E5 намного выше нескольких GB.

Q : «Есть ли другой способ улучшить мой тройной l oop?»

Избегайте передачи избыточных «параметров». Это дорого. Чем больше, тем больше они становятся. Ia, Ja, Ka - повторно представить индексы естественной области из A[i,j,k], которые уже присутствуют внутри A -пример, не так ли? Передача и получение больших избыточных параметров - это роскошь, которую мы часто предпочитаем избегать.

@njit(parallel=True)
def SwVer( ResSwc,   SwNew,           Ia, Ja, Ka ):
    #                     [: : :]     ||  ||  ||
    for                        k in (         Ka ):
        for                  j   in (     Ja ):
            for            i     in ( Ia ):
                if ( SwNew[i,j,k] <        ResSwc ):
                     SwNew[i,j,k]  =       ResSwc
                if ( SwNew[i,j,k] >  ( 1 - ResSwc ) ):
                     SwNew[i,j,k]  = ( 1 - ResSwc )
    #                SwNew[:,:,:]
    return           SwNew

Как было СОСТОЯНИЕ-0:

Исходный код был выполнен примерно в ~ 1,824,168 [us] ~ 1.8 [s].

Небольшой ШАГ для человека:

Модификации на месте всегда быстрее, чем наличие множества промежуточных экземпляров того же размера для сбора окончательных результатов, что является еще одним анти-шаблоном высокопроизводительного кода.

Просто удаление последней строки return SwNew дает ~ 714,977 [us] ~ 0.7 [s]

Теперь
а маленький ШАГ для человека,
но ГИГАНСКИЙ Прыжок для ЧЕЛОВЕЧЕСТВА _{(производительность)} :

Для уникальной производительности на тривиальных [i,j,k] преобразованиях можно попробовать смелые numpy -векторизованные вычисления.

Весь фокус в следующем: np.where( A > T_hi, T_hi, np.where( A < T_lo, T_lo, A ))

Ваша функция может работать где угодно между 12086 [us] и 97810 [us] в numpy -векторный режим. Детали физической ОЗУ, кеша, ЦП и операционных нагрузок будут вызывать различные эффекты, но размеры ОЗУ имеют значение для умно-векторизованного кода, а numpy умный и много, самый :
A[1000,1000,1000] ~ 8 [GB] -RAM-footprints. Но детали имеют значение. Очень много.
A[ 100, 100, 100] ~ 8 [MB] -ОЗУ-следы. Теперь он помещается в кэш L3 / L2 ... это важно ...
A[ 10, 10, 10] ~ 8 [kB] -RAM-footprints. Теперь он помещается в L1d-кеш ... это важно. Много ...

Давайте начнем с очень начала:

Мы начнем с 2D-уменьшенной размерности из-за ОЗУ.

>>> from zmq import Stopwatch;  aClk = Stopwatch()     # a [us]-resolution clock
>>> pass;    import gc;               gc.disable()     #   ...  elementary
>>> pass;    import numpy as np
#_______________________________________________________________________________
>>> N = int( 1E3 ); aClk.start(); A = np.random.random( ( N, N, N ) ); aClk.stop()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "mtrand.pyx", line 861, in mtrand.RandomState.random_sample
  File "mtrand.pyx", line 167, in mtrand.cont0_array
MemoryError

Тем не менее, Трюк с векторизованным кодом остается в принципе таким же для обработки 1D- ... nD-тензора:

#_______________________________________________________________________________
#       [us]
#_______________________________________________________________________________
>>> N = int( 1E4 ); aClk.start(); A = np.random.random( ( N, N ) ); aClk.stop()
17801992
>>> N = int( 1E3 ); aClk.start(); A = np.random.random( ( N, N ) ); aClk.stop()
  184895
>>> N = int( 1E2 ); aClk.start(); A = np.random.random( ( N, N ) ); aClk.stop()
    1585
>>> N = int( 1E1 ); aClk.start(); A = np.random.random( ( N, N ) ); aClk.stop()
      44
>>> N = int( 1E2 ); aClk.start(); A = np.random.random( ( N, N ) ); aClk.stop()
     465
>>> N = int( 1E3 ); aClk.start(); A = np.random.random( ( N, N ) ); aClk.stop()
   48651
>>> N = int( 1E4 ); aClk.start(); A = np.random.random( ( N, N ) ); aClk.stop()
 4954694
>>> N = int( 1E4 ); aClk.start(); A = np.random.random( ( N, N ) ); aClk.stop()
25549190
#_______________________________________________________________________________
#       [us] SEE THE GROWING COSTS FOR ram-ALLOCATIONS & STORAGE OF RANDOM num-s
#_______________________________________________________________________________


>>> N = int( 1E3 ); aClk.start(); A = np.random.random( ( N, N ) ); aClk.stop()
  471956
   50067
   49184
   42891
   48897
   52639
   45464
   48828
#_______________________________________________________________________________
#       [us] SEE THE 1st, resp. WAY-LOWER COSTS FOR 1st, resp. OTHER ram-ACCESS
#_______________________________________________________________________________


>>> T_lo = 0.25; T_hi = 1 - T_lo; aClk.start(); A = np.where( A > T_hi, T_hi, np.where( A < T_lo, T_lo, A ) ); aClk.stop()
   71044
   12086
   16277
   28192
#_______________________________________________________________________________
#       [us] SEE ALSO THE "noise" IN THE LATENCY-COSTS IN 2nd+ RE-RUN-s
#_______________________________________________________________________________

>>> N = int( 1E3 );               aClk.start(); A = np.random.random( ( N, N ) );                              aClk.stop()
   45759
>>> T_lo = 0.25; T_hi = 1 - T_lo; aClk.start(); A = np.where( A > T_hi, T_hi, np.where( A < T_lo, T_lo, A ) ); aClk.stop()
   38362    
   46640
   37927

>>> N = int( 1E4 );               aClk.start(); A = np.random.random( ( N, N ) );                              aClk.stop()
 4885472
>>> T_lo = 0.25; T_hi = 1 - T_lo; aClk.start(); A = np.where( A > T_hi, T_hi, np.where( A < T_lo, T_lo, A ) ); aClk.stop()
35747022
#_______________________________________________________________________________
#       [us] SEE THE SHIFT IN LATENCY-COSTS AS ram-i/o COSTS DOMINATE > 1 [GB]
#_______________________________________________________________________________


>>> N = int( 1E3 );               aClk.start(); A = np.random.random( ( N, N ) );                              aClk.stop()
 2307509
   50446
   49464
   43006
   50958
   54800
   43418
   57091
   52135
   46451

>>> T_lo = 0.25; T_hi = 1 - T_lo; aClk.start(); A = np.where( A > T_hi, T_hi, np.where( A < T_lo, T_lo, A ) ); aClk.stop()
   97810
   20117
   14480
   22696
   31172
   14846
#_______________________________________________________________________________
#       [us] SEE THE "noise" FOR 2nd+ RE-RUN-s
#_______________________________________________________________________________

>>> N = int( 1E3 );               aClk.start(); A = np.random.random( ( N, N ) );                              aClk.stop()
   47437
>>> T_lo = 0.25; T_hi = 1 - T_lo; aClk.start(); A = np.where( A > T_hi, T_hi, np.where( A < T_lo, T_lo, A ) ); aClk.stop()
   39298
   19422
#_______________________________________________________________________________
#       [us] SEE THE "noise" FOR 2nd+ RE-RUN
#_______________________________________________________________________________

>>> N = int( 1E3 ); aClk.start();               A = np.random.random( ( N, N ) );                              aClk.stop()
   44814
>>> T_lo = 0.25; T_hi = 1 - T_lo; aClk.start(); A = np.where( A > T_hi, T_hi, np.where( A < T_lo, T_lo, A ) ); aClk.stop()
   42565

>>> N = int( 1E3 ); aClk.start();               A = np.random.random( ( N, N ) );                              aClk.stop()
   43120
>>> T_lo = 0.25; T_hi = 1 - T_lo; aClk.start(); A = np.where( A > T_hi, T_hi, np.where( A < T_lo, T_lo, A ) ); aClk.stop()
   38039

>>> N = int( 1E3 ); aClk.start();               A = np.random.random( ( N, N ) );                              aClk.stop()
   45296
>>> T_lo = 0.25; T_hi = 1 - T_lo; aClk.start(); A = np.where( A > T_hi, T_hi, np.where( A < T_lo, T_lo, A ) ); aClk.stop()
   41898
#_______________________________________________________________________________
#       [us] SEE THE "noise" FOR block-RE-randomised-RE-RUN-s
#_______________________________________________________________________________

Answer 2

У вас как минимум две проблемы с производительностью в коде.

• Numpy массивы по умолчанию C -упорядочены (последнее измерение изменяется быстрее всего), поэтому внутренний l oop должен выполнять итерацию по последней оси.
• Вы используете массивы индексов для итерации данных. Это требует значительных накладных расходов (памяти и производительности) и совершенно не требуется.

Три шага

Помимо очень простых проблем обычно необходимо использовать nb.parfor явно. Также обратите внимание, что параллельная версия работает медленнее при довольно мелких проблемах со временем работы в диапазоне мкс.

@njit(parallel=False) #True for parallel
def SwVer(ResSwc,SwNew):
    for i in range(SwNew.shape[0]):
        for j in range(SwNew.shape[1]):
            for k in range(SwNew.shape[2]):
                if(SwNew[i,j,k]<ResSwc):
                    SwNew[i,j,k]=ResSwc
                if(SwNew[i,j,k]>(1-ResSwc)):
                    SwNew[i,j,k]=(1-ResSwc)
    return SwNe

Один l oop

#only works on contigous arrays, otherwise reshape will fail
@njit(parallel=False)
def SwVer_unrolled(ResSwc,SwNew):
    shape_orig=SwNew.shape
    SwNew_flat=SwNew.reshape(-1) #create a 1d-view
    for i in nb.prange(SwNew_flat.shape[0]):
        if(SwNew_flat[i]<ResSwc):
            SwNew_flat[i]=ResSwc
        if(SwNew_flat[i]>(1-ResSwc)):
            SwNew_flat[i]=(1-ResSwc)
    return SwNew_flat.reshape(shape_orig)

Numpy векторизованная версия

def SwVer_np(ResSwc,SwNew):
    SwNew[SwNew<ResSwc]    =ResSwc
    SwNew[SwNew>(1-ResSwc)]=(1-ResSwc)
    return SwNew

Тайминги

#very small array
Imax=100
Jmax=100
Kmax=5

Ia=np.arange(0,Imax,1)
Ja=np.arange(0,Jmax,1)
Ka=np.arange(0,Kmax,1)

#your version
%timeit SwVer_orig(0.25,SwNew,Ia,Ja,Ka)
#181 µs ± 1.92 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit SwVer(0.25,SwNew)
#parallel=False
#44.6 µs ± 213 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
#parallel=True
#104 µs ± 5.61 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit SwVer_unrolled(0.25,SwNew)
#parallel=False
#11.4 µs ± 96.8 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
#parallel=True
#116 µs ± 4.45 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
%timeit SwVer_np(0.25,SwNew)
#30.1 µs ± 568 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

#bigger arrays -> parallelization beneficial
Imax=1000
Jmax=1000
Kmax=5
SwNew=np.random.random((Imax,Jmax,Kmax))

#your version
%timeit SwVer_orig(0.25,SwNew,Ia,Ja,Ka)
#17.7 ms ± 165 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit SwVer(0.25,SwNew)
#parallel=False
#4.73 ms ± 96.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
#parallel=True
#1.3 ms ± 63.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit SwVer_unrolled(0.25,SwNew)
#parallel=False
#2.03 ms ± 18.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
#parallel=True
#1.17 ms ± 30.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit SwVer_np(0.25,SwNew)
#7.9 ms ± 100 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Заключение

Если эта часть не очень важна для производительности. Я бы предпочел векторизованную версию Numpy из-за ее простоты. Вы можете превзойти его, но я не думаю, что это того стоит.