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%.

我不明白为什么指定parallel并没有改变任何东西而且需要使用prange.在doc中,有一个范围示例.

Answer 1

来自 Numba 文档：

@jit 的实验性 parallel=True 选项将尝试优化数组操作并并行运行它们。它还添加了对 prange() 的支持以显式并行化循环。

现在，由于您不在函数中执行任何数组操作，因此如果不使用 prange 显式标记循环，Numba 就无法并行化。

所以只是为了确保没有混淆。仅当您在装饰中设置parallel=True 时，Numba 才会将循环拆分为线程，并通过更改显式标记循环；范围 -> 范围。

在您的 f4() 中，您已将 prange 放在所有 for 循环上，我建议仅将 prange 放在最外层循环上，因为您不想冒险从线程中生成线程。IE：

@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