cython 的 prange 中的线程局部数组没有大量内存分配

Question

我想使用 Cython 并行执行一些独立计算。

现在我正在使用这种方法：

import numpy as np
cimport numpy as cnp
from cython.parallel import prange

[...]

cdef cnp.ndarray[cnp.float64_t, ndim=2] temporary_variable = \
    np.zeros((INPUT_SIZE, RESULT_SIZE), np.float64)
cdef cnp.ndarray[cnp.float64_t, ndim=2] result = \
    np.zeros((INPUT_SIZE, RESULT_SIZE), np.float64)

for i in prange(INPUT_SIZE, nogil=True):
    for j in range(RESULT_SIZE):
        [...]
        temporary_variable[i, j] = some_very_heavy_mathematics(my_input_array)
        result[i, j] = some_more_maths(temporary_variable[i, j])

这种方法可行，但我的问题是我实际上需要多个 temporary_variable。当 INPUT_SIZE 增长时，这会导致 巨大的 内存使用。但我相信真正需要的是每个线程中的临时变量。

我是否面临 Cython 的 prange 的限制，我是否需要学习正确的 C 或者我 doing/understanding 有什么非常错误的地方？

编辑：我正在寻找的函数是 openmp.omp_get_max_threads() 和 openmp.omp_get_thread_num() 来创建一个合理大小的临时数组。我必须先cimport openmp。

Answer 1

这是 Cython 试图检测的东西，实际上大部分时间都是正确的。如果我们拿一个更完整的示例代码：

import numpy as np
from cython.parallel import prange

cdef double f1(double[:,:] x, int i, int j) nogil:
    return 2*x[i,j]

cdef double f2(double y) nogil:
    return y+10

def example_function(double[:,:] arr_in):
    cdef double[:,:] result = np.zeros(arr_in.shape)
    cdef double temporary_variable
    cdef int i,j
    for i in prange(arr_in.shape[0], nogil=True):
        for j in range(arr_in.shape[1]):
            temporary_variable = f1(arr_in,i,j)
            result[i,j] = f2(temporary_variable)
    return result

（这和你的基本一样，但是可以编译）。这编译为 C 代码：

#pragma omp for firstprivate(__pyx_v_i) lastprivate(__pyx_v_i) lastprivate(__pyx_v_j) lastprivate(__pyx_v_temporary_variable)
                #endif /* _OPENMP */
                for (__pyx_t_8 = 0; __pyx_t_8 < __pyx_t_9; __pyx_t_8++){

可以看到temporary_variable设置为thread-local。如果 Cython 没有正确检测到这一点（我发现它通常太热衷于减少变量）那么我的建议是将循环的（部分）内容封装在一个函数中：

cdef double loop_contents(double[:,:] arr_in, int i, int j) nogil:
    cdef double temporary_variable
    temporary_variable = f1(arr_in,i,j)
    return f2(temporary_variable)

这样做会强制 temporary_variable 对函数而言是局部的（因此对线程而言）

关于创建线程本地数组：我不是 100% 清楚您想要做什么，但我会尝试猜测...

我认为不可能创建线程局部内存视图。
您可以使用 malloc 和 free 创建线程局部 C 数组，但除非您对 C 有良好的理解，否则我不会推荐它。

最简单的方法是分配一个二维数组，其中每个线程都有一列。该数组是共享的，但由于每个线程只触及它自己的无关紧要的列。一个简单的例子：

cdef double[:] f1(double[:,:] x, int i) nogil:
    return x[i,:]

def example_function(double[:,:] arr_in):
    cdef double[:,:] temporary_variable = np.zeros((arr_in.shape[1],openmp.omp_get_max_threads()))
    cdef int i
    for i in prange(arr_in.shape[0],nogil=True):
        temporary_variable[:,openmp.omp_get_thread_num()] = f1(arr_in,i)

cython 的 prange 中的线程局部数组没有大量内存分配

Thread-local arrays in cython's prange without huge memory allocation

multithreading

memory-management

numpy

cython

gil