这可能令人惊讶,但总结矩阵中的一些值并不是一件容易的事。我的This 和this 的答案给出了一些见解,所以我不打算重复太多细节,但为了获得最佳性能,必须在现代 CPU 上利用缓存和 SIMD/管道性质的浮动操作最好的办法。
Numpy 获得了上述所有权利,并且很难被击败。这不是不可能的,但是一个人必须非常熟悉低级优化才能成功——而且人们不应该期望有太多的改进。幼稚的实现根本无法击败 Numpy。
这是我使用 cython 和 numba 的尝试,所有的加速都来自并行化。
让我们从基线算法开始:
def bin2d(a,K):
m_bins = a.shape[0]//K
n_bins = a.shape[1]//K
return a.reshape(m_bins, K, n_bins, K).sum(3).sum(1)
并衡量它的性能:
import numpy as np
N,K=2000,50
a=np.arange(N*N, dtype=np.float64).reshape(N,N)
%timeit bin2d(a,K)
# 2.76 ms ± 107 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
赛通:
这是我使用 Cython 的实现,它使用 OpenMP 进行并行化。为了简单起见,我在原地执行求和,因此传递的数组将被更改(numpy 的版本不是这种情况):
%%cython -a -c=/openmp --link-args=/openmp -c=/arch:AVX512
import numpy as np
import cython
from cython.parallel import prange
cdef extern from *:
"""
//assumes direct memory and row-major-order (i.e. rows are continuous)
double calc_bin(double *ptr, int N, int y_offset){
double *row = ptr;
for(int y=1;y<N;y++){
row+=y_offset; //next row
//there is no dependencies, so the summation can be done in parallel
for(int x=0;x<N;x++){
ptr[x]+=row[x];
}
}
double res=0.0;
for(int x=0;x<N;x++){
//could be made slightly faster (summation is not in parallel), but is it needed?
res+=ptr[x];
}
return res;
}
"""
double calc_bin(double *ptr, int N, int y_offset) nogil
@cython.boundscheck(False)
@cython.wraparound(False)
def cy_bin2d_parallel(double[:, ::1] a, int K):
cdef int y_offset = a.shape[0]
cdef int m_bins = a.shape[0]//K
cdef int n_bins = a.shape[1]//K
cdef double[:,:] res = np.empty((m_bins, n_bins), dtype=np.float64)
cdef int i,j,k
for k in prange(m_bins*n_bins, nogil=True):
i = k//m_bins
j = k%m_bins
res[i,j] = calc_bin(&a[i*K, j*K], K, y_offset)
return res.base
现在(验证结果相同后):
a=np.arange(N*N, dtype=np.float64).reshape(N,N)
%timeit cy_bin2d_parallel(a,K)
# 1.51 ms ± 25.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
# without parallelization: 2.93 ms
大约快 30%。使用@max9111 提出的-c=/arch:AVX512(否则只能使用/Ox 和MSVC 编译),使单线程版本的速度提高了约20%,但并行版本的速度仅提高了约5%。
Numba:
使用 numba 编译的算法相同(由于 clang-compiler 的性能更好,它通常可以击败 cython)——但结果比 cython 稍慢,但比 numpy 快 20%:
import numba as nb
@nb.njit(parallel=True)
def nb_bin2d_parallel(a, K):
m_bins = a.shape[0]//K
n_bins = a.shape[1]//K
res = np.zeros((m_bins, n_bins), dtype=np.float64)
for k in nb.prange(m_bins*n_bins):
i = k//m_bins
j = k%m_bins
for y in range(i*K+1, (i+1)*K):
for x in range(j*K, (j+1)*K):
a[i*K, x] += a[y,x]
s=0.0
for x in range(j*K, (j+1)*K):
s+=a[i*K, x]
res[i,j] = s
return res
现在:
a=np.arange(N*N, dtype=np.float64).reshape(N,N)
%timeit nb_bin2d_parallel(a,K)
# 1.98 ms ± 162 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
# without parallelization: 5.8 ms
简而言之:我想有可能击败上述方法,但不再有免费的午餐,因为 numpy 做得很好。最大的潜力可能是并行化,但由于问题的内存带宽限制性质而受到限制(并且可能应该像我一样使用更智能的策略 - 对于 50*50 的求和,一个可能的仍然可以看到开销创建/管理线程)。
Cython 有另一种更快的尝试(至少对于大约 2000 的大小),它不是对 50 个元素的小部分执行求和,而是对整行执行求和,从而减少开销(但一旦行大于 1-2k):
%%cython -a --verbose -c=/openmp --link-args=/openmp -c=/arch:AVX512
import numpy as np
import cython
from cython.parallel import prange
cdef extern from *:
"""
void calc_bin_row(double *ptr, int N, int y_offset, double* out){
double *row = ptr;
for(int y=1;y<N;y++){
row+=y_offset; //next row
for(int x=0;x<y_offset;x++){
ptr[x]+=row[x];
}
}
double res=0.0;
int i=0;
int k=0;
for(int x=0;x<y_offset;x++){//could be made faster, but is it needed?
res+=ptr[x];
k++;
if(k==N){
k=0;
out[i]=res;
i++;
res=0.0;
}
}
}
"""
void calc_bin_row(double *ptr, int N, int y_offset, double* out) nogil
@cython.boundscheck(False)
@cython.wraparound(False)
def cy_bin2d_parallel_rowise(double[:, ::1] a, int K):
cdef int y_offset = a.shape[1]
cdef int m_bins = a.shape[0]//K
cdef int n_bins = a.shape[1]//K
cdef double[:,:] res = np.empty((m_bins, n_bins), dtype=np.float64)
cdef int i,j,k
for k in prange(0, y_offset, K, nogil=True):
calc_bin_row(&a[k, 0], K, y_offset, &res[k//K, 0])
return res.base
单线程 (2ms) 已经更快了,并行化速度提高了大约 60%(每个循环 1.27 ms ± 50.8 µs(平均值 ± 标准偏差,7 次运行,每次 100 个循环))。