使用Cython内存视图实现并行化

Question

这里是我试图完成的最小示例。我正在计算堆叠数组中矩阵之间的成对距离。 idx 数组保存每个子矩阵的偏移量。

当我删除 parallel() 并将 prange 替换为 range 时，代码按预期工作（当然不是并行的）。

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

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
def all_pairs_distance(float[:,::1] stacked, int[::1] idx):
  cdef int n = idx.shape[0] - 1
  cdef float[:,::1] D = np.zeros((n,n), dtype='float32')
  cdef int i,j
  cdef float[:,::1] t1,t2
  cdef float d
  with nogil, parallel():
    for i in prange(n):
      t1 = stacked[idx[i]:idx[i+1],:]
      for j in range(i+1, n):
        t2 = stacked[idx[j]:idx[j+1],:]
        d = nogil_cython_function(t1, t2)
        D[i,j] = d
        D[j,i] = d
  return D

cdef float nogil_cython_function(float[:,::1] a, float[:,::1] b) nogil:
  # Function abbreviated for simplicity
  return 0.0

尝试编译这个文件时，每次给 t1 或 t2 赋值时会出现错误。

Memoryview slices can only be shared in parallel sections

1. 我不确定如何解决这个错误。那些赋值语句不在并行部分吗？
2. 我该如何告诉编译器我的 stacked 内存视图是只读的，并且可以在线程之间共享？

- perimosocordiae

---

我认为你可以给每个线程只分配他们实际使用的数组部分。请查看这个答案，其中使用原始指针实现了这一点（我认为这对于此任务更加直观）。 - Saullo G. P. Castro

Answer 1

经过一番尝试，我终于弄清楚了：

for i in prange(n, nogil=True):
  for j in range(i+1, n):
    d = nogil_cython_function(stacked[idx[i]:idx[i+1],:],
                              stacked[idx[j]:idx[j+1],:])

结论是在prange内部对内存视图进行切片是可以的，但将这些切片赋值给变量是不行的。

- perimosocordiae

---

这正是我需要的，以便使用memoryview切片来工作。显然，您的解决方案在其他地方没有记录，只需不将切片分配给变量，直接使用切片即可。 - Matt

将prange循环内的代码“包装”到一个小的帮助函数中（只需重写您的nogil_cython_function，使其接受堆叠的自身而不是切片）。在helper函数内部，可以随意分配memoryviews。至少目前为止，对我进行了简单测试。 - oli