Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Parallelizing with Cython memoryviews

Here's a minimal example of what I'm trying to do. I'm computing pairwise distances between matrices that I've arranged into a stacked array. The idx array holds the offsets of each submatrix.

When I remove the parallel() and replace prange with range, the code works as expected (just not in parallel, of course).

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

When trying to compile this file, I get errors on each assignment to t1 or t2:

Memoryview slices can only be shared in parallel sections
  1. I'm not sure how to resolve this error. Are those assignments not in a parallel section?
  2. How do I tell the compiler that my stacked memoryview is read-only, and is fine to share between threads?
like image 453
perimosocordiae Avatar asked Dec 19 '22 13:12

perimosocordiae


1 Answers

I figured it out, after a bit of experimenting:

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],:])

The takeaway is that slicing memoryviews inside a prange is fine, but assigning those slices to variables is not.

like image 73
perimosocordiae Avatar answered May 10 '23 02:05

perimosocordiae