Numba не распараллеливает диапазон

Question

В моем коде есть циклы, которые я хочу распараллелить

from numba import njit, prange
from time import time


@njit
def f1(n):
    s = 0
    for i in range(n):
        for j in range(n):
            for k in range(n):
                s += (i * k < j * j) - (i * k > j * j)
    return s


@njit
def f2(n):
    s = 0
    for i in prange(n):
        for j in prange(n):
            for k in prange(n):
                s += (i * k < j * j) - (i * k > j * j)
    return s


@njit(parallel=True)
def f3(n):
    s = 0
    for i in range(n):
        for j in range(n):
            for k in range(n):
                s += (i * k < j * j) - (i * k > j * j)
    return s


@njit(parallel=True)
def f4(n):
    s = 0
    for i in prange(n):
        for j in prange(n):
            for k in prange(n):
                s += (i * k < j * j) - (i * k > j * j)
    return s


for f in [f1, f2, f3, f4]:
    d = time()
    f(2500)
    print('%.02f' % (time() - d))

Я получаю время:

27.44
27.34
26.83
13.05

Я проверил активность моего процессора, и хотя первые три функции были на уровне 100%, четвертая была на уровне ~ 300%.

Я не понимаю, почему указание параллели ничего не изменило, и нужно использовать prange. В документе есть пример с диапазоном.

Answer 1

Я создал проблему на Github, и автор очень точно ответил (https://github.com/numba/numba/issues/2960#issuecomment-388767318).

Если он захочет опубликовать здесь свой ответ, я приму его.

Answer 2

Из документации Numba:

Экспериментальная параллельная = True опция @jit попытается оптимизировать операции с массивами и запустить их параллельно. Также добавлена ​​поддержка prange () для явного распараллеливания цикла.

Теперь, поскольку вы не выполняете никаких операций с массивами в своей функции, Numba не может ничего распараллелить без явного пометки циклов с помощью prange.

Так что просто чтобы убедиться, что нет путаницы. Numba будет разбивать ваш цикл на потоки только тогда, когда вы установите параллельное = True в декорации, и явно помечает циклы, изменяя их; диапазон -> prange.

В вашей функции f4 () вы поместили prange на все циклы for, я бы рекомендовал помещать prange только в самый внешний цикл, потому что вы не хотите рисковать порождением потоков из потоков. I.e.:

@njit(parallel=True)
def f5(n):
    s = 0
    for i in prange(n):
        for j in range(n):
            for k in range(n):
                s += (i * k < j * j) - (i * k > j * j)
    return s