Локальные по потоку массивы в prange Cython без огромного выделения памяти - PullRequest
Новая
       44

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

0 голосов
/ 30 августа 2018

У меня есть несколько независимых вычислений, которые я хотел бы сделать параллельно с использованием 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 голосов
/ 01 сентября 2018

Это то, что 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. Вы можете создать локальный для потока массив 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)
...