Cython：使prange並行化線程安全

Question

Cython啓動器在這裏。我試圖通過使用多個線程來加速計算某些成對統計量（在幾個分箱中）。特別是，我使用cython.parallel中的prange，它在內部使用openMP。

下面的最小例子說明了這個問題（通過Jupyter筆記本Cython magic編譯）。

筆記本設置：

%load_ext Cython 
import numpy as np 

用Cython代碼：

%%cython --compile-args=-fopenmp --link-args=-fopenmp -a 

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

@boundscheck(False) 
def my_parallel_statistic(double[:] X, double[:,::1] bins, int num_threads): 

    cdef: 
     int N = X.shape[0] 
     int nbins = bins.shape[0] 
     double Xij,Yij 
     double[:] Z = np.zeros(nbins,dtype=np.float64) 
     int i,j,b 

    with nogil, parallel(num_threads=num_threads): 
     for i in prange(N,schedule='static',chunksize=1): 
      for j in range(i): 
       #some pairwise quantities 
       Xij = X[i]-X[j] 
       Yij = 0.5*(X[i]+X[j]) 
       #check if in bin 
       for b in range(nbins): 
        if (Xij < bins[b,0]) or (Xij > bins[b,1]): 
         continue 
        Z[b] += Xij*Yij 

    return np.asarray(Z) 

模擬數據和箱

X = np.random.rand(10000) 
bin_edges = np.linspace(0.,1,11) 
bins = np.array([bin_edges[:-1],bin_edges[1:]]).T 
bins = bins.copy(order='C') 

時序經由

%timeit my_parallel_statistic(X,bins,1) 
%timeit my_parallel_statistic(X,bins,4) 

個

產量

1 loop, best of 3: 728 ms per loop 
1 loop, best of 3: 330 ms per loop 

這是不是一個完美的比例，但是這不是問題的重點。 （但不要讓我知道，如果你有超越添加常用的裝飾或微調PRANGE參數的建議。）

然而，這種計算顯然不是線程安全的：

Z1 = my_parallel_statistic(X,bins,1) 
Z4 = my_parallel_statistic(X,bins,4) 
np.allclose(Z1,Z4) 

透着顯著差異在這兩個結果之間（在這個例子中高達20％）。

我強烈懷疑，問題是，多個線程可以在同一時間做

Z[b] += Xij*Yij 

。但是我不知道如何在不犧牲加速的情況下解決這個問題。

在我的實際使用情況中，Xij和Yij的計算更加昂貴，因此我希望每對執行一次。此外，預先計算和存儲所有對的Xij和Yij，然後簡單地循環通過bin不是一個好選擇，因爲N可以變得非常大，並且我不能在內存中存儲100,000 x 100,000個numpy數組（這實際上是在Cython中重寫它的主要動機！）。

系統信息（中添加註釋如下建議）：

CPU(s): 8 
Model name: Intel(R) Core(TM) i7-4790K CPU @ 4.00GHz 
OS: Red Hat Linux v6.8 
Memory: 16 GB 

Answer 1

是的，Z[b] += Xij*Yij確實是一種競賽條件。

有幾個選項可以使atomic或critical。拋開Cython的實施問題，由於共享的Z矢量上的虛假共享，您在任何情況下都會導致性能不佳。

所以更好的選擇是爲每個線程保留一個專用數組。還有一些（非）選項。可以使用私人的malloc'd指針，但我想堅持np。內存片不能被分配爲私有變量。二維(num_threads, nbins)數組可行，但由於某種原因會生成非常複雜的低效數組索引代碼。這有效，但速度較慢，不能縮放。

具有手動「2D」索引的平面numpy陣列效果很好。通過避免將數組的私有部分填充爲64字節（這是典型的高速緩存行大小），可以獲得一些額外的性能。這可以避免核心之間的錯誤共享。私密部分簡單地在平行區域外串聯。

%%cython --compile-args=-fopenmp --link-args=-fopenmp -a 
from cython cimport boundscheck 
import numpy as np 
from cython.parallel cimport prange, parallel 
cimport openmp 

@boundscheck(False) 
def my_parallel_statistic(double[:] X, double[:,::1] bins, int num_threads): 

    cdef: 
     int N = X.shape[0] 
     int nbins = bins.shape[0] 
     double Xij,Yij 
     # pad local data to 64 byte avoid false sharing of cache-lines 
     int nbins_padded = (((nbins - 1) // 8) + 1) * 8 
     double[:] Z_local = np.zeros(nbins_padded * num_threads,dtype=np.float64) 
     double[:] Z = np.zeros(nbins) 
     int i,j,b, bb, tid 

    with nogil, parallel(num_threads=num_threads): 
     tid = openmp.omp_get_thread_num() 
     for i in prange(N,schedule='static',chunksize=1): 
      for j in range(i): 
       #some pairwise quantities 
       Xij = X[i]-X[j] 
       Yij = 0.5*(X[i]+X[j]) 
       #check if in bin 
       for b in range(nbins): 
        if (Xij < bins[b,0]) or (Xij > bins[b,1]): 
         continue 
        Z_local[tid * nbins_padded + b] += Xij*Yij 
    for tid in range(num_threads): 
     for bb in range(nbins): 
      Z[bb] += Z_local[tid * nbins_padded + bb] 


    return np.asarray(Z) 

這表現相當好我的4核機器上，用720 ms/191 ms，3.6的加速。剩餘的差距可能是由於渦輪模式。我現在無法使用適當的機器進行測試。

Answer 2

你是正確的，將獲得Z是一個競爭條件下。

將num_threads的副本Z定義爲cdef double[:] Z = np.zeros((num_threads, nbins), dtype=np.float64)，並在prange循環後執行沿軸0的和可能會更好。

return np.sum(Z, axis=0) 

用Cython代碼可以在一個並行區域with gil聲明，但它只是記錄錯誤處理。您可以查看一般的C代碼，看看是否會觸發原子OpenMP操作，但我懷疑它。