Локальные массивы потока в cython's prange без огромного выделения памяти

Это то, что 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 быть локальным для функции (и, следовательно, для потока).

Что касается создания локального массива потока: я не совсем понимаю, что вы хотите сделать, но я попытаюсь предположить...

1. Я не верю, что можно создать локальное представление памяти потока.
2. Вы можете создать локальный для потока массив C с malloc и free, но если вы не хорошо понимаете 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)