python - 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。在 the doc , 有一个范围的例子。

Answer 1

来自 Numba 文档:

The experimental parallel=True option to @jit will attempt to optimize array operations and run them in parallel. It also adds support for prange() to explicitly parallelize a loop.

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

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

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

@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