Параллельные циклы с Numba - нет параллелизации с prange - PullRequest
Новая
       16

Параллельные циклы с Numba - нет параллелизации с prange

0 голосов
/ 18 февраля 2019

Я реализую последовательный алгоритм (фильтр Калмана) с определенной структурой, в которой много внутренних циклов может выполняться параллельно.Мне нужно получить как можно больше производительности от этой функции.В настоящее время он работает около 600 мс на моей машине с типичными входными данными (n, p = 12, d = 3, T = 3000)

Я использовал @numba.jit с nopython=True, parallel=True и аннотировал свои диапазоны с помощьюnumba.prange.Тем не менее, даже при очень больших входных данных (n> 5000) параллелизма явно не происходит (основываясь только на взгляде на ядра с top).

Здесь довольно много кода, япоказывая только основной кусок.Есть ли причина, по которой Numba не сможет распараллелить операции с массивами в prange?Я также проверил numba.config.NUMBA_NUM_THREADS (это 8) и играл с другими numba.config.THREADING_LAYER (это в настоящее время 'tbb').Я также пробовал с версиями numy + scipy как для openblas, так и для MKL, версия MKL выглядит немного медленнее, и все еще не распараллеливается.

Аннотация: 

@numba.jit(nopython=True, cache=False, parallel=True,
           fastmath=True, nogil=True)

И основная часть функции: 

P = np.empty((T + 1, n, p, d, d))
m = np.empty((T + 1, n, p, d))

P[0] = P0
m[0] = m0

phi = 0.0
Xt = np.empty((n, p)

for t in range(1, T + 1):
    sum_P00 = 0.0
    v = y[t - 1]

    # Purely for convenience, little performance impact
    for tau in range(1, p + 1):
        Xt[:, tau - 1] = X[p + t - 1 - tau]

    # Predict
    for i in numba.prange(n):
        for tau in range(p):
            # Prediction step
            m[t, i, tau] = Phi[i, tau] @ m[t - 1, i, tau]
            P[t, i, tau] = Phi[i, tau] @ P[t - 1, i, tau] @ Phi[i, tau].T

    # Auxiliary gain variables
    for i in numba.prange(n):
        for tau in range(p):
            v = v - Xt[i, tau] * m[t, i, tau, 0]
            sum_P00 = sum_P00 + P[t, i, tau, 0, 0]

    # Energy function update
    s = np.linalg.norm(Xt)**2 * sum_P00 + sv2
    phi += np.pi * s + 0.5 * v**2 / s

    # Update
    for i in numba.prange(n):
        for tau in range(p):
            k = Xt[i, tau] * P[t, i, tau, :, 0]  # Gain
            m[t, i, tau] = m[t, i, tau] + (v / s) * k
            P[t, i, tau] = P[t, i, tau] + (k / s) @ k.T

1 Ответ

0 голосов
/ 18 февраля 2019

Кажется, что просто проблема с интерактивным запуском в Ipython.Запуск тестового сценария из консоли приводит к параллельному выполнению, как и ожидалось.

...