Question

我在Numba中有一个函数，它具有几个可以并行化的循环。该循环写入一个公共数组K，因此我知道编译器可能没有尽其所能进行优化。但是，我对使numba的jit编译器创建最高效的代码没有任何感觉。文档中的示例过于简单，无法提供帮助。

我尝试将每个range更改为prange。在k_x上并行化循环时，可获得最佳结果，但是在具有4个内核的机器上，我只得到了2.8倍的改进。我知道我不应该期望线性性能有所提高，但是我有一种感觉，在这种情况下我应该获得更好的结果。例如，使用dask x_cond.map_blocks(cond_expect_kernel, x_tr, *args)可以获得更好的结果，考虑到调度程序的开销，这是奇怪的。

除了简单地将range更改为prange之外，还有什么方法可以改善此功能的并行性？

原始功能

@jit(float64[:,:](float64[:,:], float64[:,:], int64, int64), nopython=True, nogil=True)
def cond_expect_kernel(x_cond, x_tr, degree, amount_non_cond_vars):
    size = x_cond.shape[1]
    x_tr_cond = x_tr[:, :size]
    samples_x = x_cond.shape[0]
    samples_tr = x_tr.shape[0]
    K = (1+np.dot(x_cond, x_tr_cond.T))**degree

    for j in range(size, size+amount_non_cond_vars):
        for k_x in range(samples_x):
            for k_x_tr in range(samples_tr):
                K[k_x, k_x_tr] += x_tr[k_x_tr, j]**2*3 
                for j_left in range(size):
                    K[k_x, k_x_tr] += x_cond[k_x, j_left]*x_tr[k_x_tr, j_left]*x_tr[k_x_tr, j] ** 2 *3  
    return K

到目前为止最好的并行版本：

@jit(float64[:,:](float64[:,:], float64[:,:], int64, int64), nopython=True, parallel=True)
def cond_expect_kernel_parallel(x_cond, x_tr, degree, amount_non_cond_vars):
    size = x_cond.shape[1]
    x_tr_cond = x_tr[:, :size]
    samples_x = x_cond.shape[0]
    samples_tr = x_tr.shape[0]
    K = (1+np.dot(x_cond, x_tr_cond.T))**degree
    for j in range(size, size+amount_non_cond_vars):
        for k_x in prange(samples_x):
            for k_x_tr in range(samples_tr):
                K[k_x, k_x_tr] += x_tr[k_x_tr, j]**2*3 
                for j_left in range(size):
                    K[k_x, k_x_tr] += x_cond[k_x, j_left]*x_tr[k_x_tr, j_left]*x_tr[k_x_tr, j] ** 2 *3  
    return K

作为参考，我正在使用4核和16核的计算机。 samples_x约100,000，samples_tr约50000，size约3，amount_non_cond_vars约100。

谢谢！

Answer 1

您的代码有一些性能关键问题。

您正在将输入和输出数组声明为非连续的，例如。这将是C连续数组nb.float64[:,::1]的败笔。这通常会阻止SIMD向量化，并且在许多情况下会导致性能降低。如果不确定数组是否是C邻接的，只需不声明即可，Numba可以自行完成。
如果代码中有一些简化，请使用标量求和变量，然后将结果复制到数组中。必须避免不必要地读取/写入共享阵列。
请考虑您的内存访问模式/循环顺序。如果您在这里做错了，那将很早就陷入内存瓶颈。
如果类似size之类的数字经常为3，则可以编写此问题的专用版本。（在这种情况下，手动循环展开）。您可以使用小型包装函数来检查是否发生特殊情况。

示例

import numpy as np
import time
import llvmlite.binding as llvm
llvm.set_option('', '--debug-only=loop-vectorize')
@nb.njit(nb.float64[:,:](nb.float64[:,:], nb.float64[:,:], nb.int64, nb.int64),fastmath=True,parallel=True)
def cond_expect_kernel_gen(x_cond, x_tr, degree, amount_non_cond_vars):
    x_tr_cond = x_tr[:,:x_cond.shape[1]]
    K =  np.dot(x_cond, x_tr_cond.T)

    for k_x in nb.prange(x_cond.shape[0]):
        for k_x_tr in range(x_tr.shape[0]):
          sum=(K[k_x, k_x_tr]+1)**degree
          for j in range(x_cond.shape[1], x_cond.shape[1]+amount_non_cond_vars):
            sum += x_tr[k_x_tr, j]**2*3 
            for j_left in range(x_cond.shape[1]):
                sum += x_cond[k_x, j_left]*x_tr[k_x_tr, j_left]*x_tr[k_x_tr, j] ** 2 *3
          K[k_x, k_x_tr]=sum
    return K

@nb.njit(nb.float64[:,::1](nb.float64[:,::1], nb.float64[:,::1], nb.int64, nb.int64),fastmath=True,parallel=True)
def cond_expect_kernel_3(x_cond, x_tr, degree, amount_non_cond_vars):
    assert x_cond.shape[1]==3
    x_tr_cond = x_tr[:,:x_cond.shape[1]]
    K = np.dot(x_cond, x_tr_cond.T)

    for k_x in nb.prange(x_cond.shape[0]):
        for k_x_tr in range(x_tr.shape[0]):
          sum=(K[k_x, k_x_tr]+1)**degree
          for j in range(x_cond.shape[1], x_cond.shape[1]+amount_non_cond_vars):
            sum += x_tr[k_x_tr, j]**2*3
            sum_2=0.
            sum_2 += x_cond[k_x, 0]*x_tr[k_x_tr, 0]
            sum_2 += x_cond[k_x, 1]*x_tr[k_x_tr, 1]
            sum_2 += x_cond[k_x, 2]*x_tr[k_x_tr, 2]
            sum+=sum_2*x_tr[k_x_tr, j] ** 2 *3
          K[k_x, k_x_tr]=sum
    return K

性能

x_cond=np.random.rand(10_000,3)
x_tr=np.random.rand(5_000,103)
amount_non_cond_vars=100
degree=3

t1=time.time()
res_1=cond_expect_kernel_gen(x_cond, x_tr, degree, amount_non_cond_vars)
print(time.time()-t1)
t1=time.time()
res_2=cond_expect_kernel_3(x_cond, x_tr, degree, amount_non_cond_vars)
print(time.time()-t1)

(Quadcore i7, Numba 0.40dev)
your version, single threaded: 40s
your version, parallel: 8.61s

mod_general:3.8s
mod_3: 1.35s

改善Numba中的并行化

1 个答案: