Выделите промежуточные многомерные массивы в Cython без получения GIL

Я пытаюсь использовать Cython для распараллеливания дорогостоящей операции, которая включает генерацию промежуточных многомерных массивов.

Следующий очень упрощенный код иллюстрирует то, что я пытаюсь сделать:

import numpy as np
cimport cython
cimport numpy as np
from cython.parallel cimport prange
from libc.stdlib cimport malloc, free


@cython.boundscheck(False)
@cython.wraparound(False)
def embarrasingly_parallel_example(char[:, :] A):

    cdef unsigned int m = A.shape[0]
    cdef unsigned int n = A.shape[1]
    cdef np.ndarray[np.float64_t, ndim = 2] out = np.empty((m, m), np.float64)
    cdef unsigned int ii, jj
    cdef double[:, :] tmp

    for ii in prange(m, nogil=True):
        for jj in range(m):

            # allocate a temporary array to hold the result of
            # expensive_function_1
            tmp_carray = <double * > malloc((n ** 2) * sizeof(double))

            # a 2D typed memoryview onto tmp_carray
            tmp = <double[:n, :n] > tmp_carray

            # shove the intermediate result in tmp
            expensive_function_1(A[ii, :], A[jj, :], tmp)

            # get the final (scalar) output for this ii, jj
            out[ii, jj] = expensive_function_2(tmp)

            # free the intermediate array
            free(tmp_carray)

    return out


# some silly examples - the actual operation I'm performing is a lot more
# involved
# ------------------------------------------------------------------------
@cython.boundscheck(False)
@cython.wraparound(False)
cdef void expensive_function_1(char[:] x, char[:] y, double[:, :] tmp):

    cdef unsigned int m = tmp.shape[0]
    cdef unsigned int n = x.shape[0]
    cdef unsigned int ii, jj

    for ii in range(m):
        for jj in range(m):
            tmp[ii, jj] = 0
            for kk in range(n):
                tmp[ii, jj] += (x[kk] + y[kk]) * (ii - jj)


@cython.boundscheck(False)
@cython.wraparound(False)
cdef double expensive_function_2(double[:, :] tmp):

    cdef unsigned int m = tmp.shape[0]
    cdef unsigned int ii, jj
    cdef double result = 0

    for ii in range(m):
        for jj in range(m):
            result += tmp[ii, jj]

    return result

Кажется, есть по крайней мере две причины, по которым это не удается скомпилировать:

1. Основываясь на выводе cython -a, создание представления типизированной памяти здесь:

   cdef double[:, :] tmp = <double[:n, :n] > tmp_carray
   

   , кажется, включает вызовы API Python, и поэтому я не могу выпустить GIL, чтобы разрешить внешний цикл должен работать параллельно.

   У меня сложилось впечатление, что типизированные представления памяти не являются объектами Python, и поэтому дочерний процесс должен иметь возможность создавать их без предварительного получения GIL. Так ли это?

2. Даже если я заменю prange(m, nogil=True) на нормальный range(m), Cython все равно не нравится присутствие cdef во внутреннем цикле:

    Error compiling Cython file:
    ------------------------------------------------------------
    ...
                # allocate a temporary array to hold the result of
                # expensive_function_1
                tmp_carray = <double*> malloc((n ** 2) * sizeof(double))

                # a 2D typed memoryview onto tmp_carray
                cdef double[:, :] tmp = <double[:n, :n]> tmp_carray
                    ^
    ------------------------------------------------------------

    parallel_allocate.pyx:26:17: cdef statement not allowed here

 cdef double[:, :] tmp

 tmp = <double[:n, :n] > tmp_carray

Error compiling Cython file:
------------------------------------------------------------
...
            # allocate a temporary array to hold the result of
            # expensive_function_1
            tmp_carray = <double*> malloc((n ** 2) * sizeof(double))

            # a 2D typed memoryview onto tmp_carray
            tmp = <double[:n, :n]> tmp_carray
               ^
------------------------------------------------------------

parallel_allocate.pyx:28:16: Memoryview slices can only be shared in parallel sections