cython的prange中的线程局部数组,无需分配大量内存

时间:2018-08-30 15:48:20

标签: multithreading numpy memory-management cython gil

我想使用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的限制,我是否需要学习适当的C,或者我在做/理解某些非常错误的事情?

编辑:我正在寻找的功能是openmp.omp_get_max_threads()openmp.omp_get_thread_num(),以创建合理大小的临时数组。我必须先cimport openmp

1 个答案:

答案 0 :(得分:2)

这是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被设置为线程本地的。如果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%确切地知道您要做什么,但我会尝试猜测一下...

  1. 我认为无法创建线程本地的内存视图。
  2. 您可以使用mallocfree创建一个线程局部C数组,但是除非您对C有好的的知识,否则我不建议您这样做。
  3. 最简单的方法是分配一个2D数组,其中每个线程都有一列。该数组是共享的,但由于每个线程仅接触其自己的列,所以没有关系。一个简单的例子:

    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)