Это может быть удивительно, но суммирование некоторых значений в матрице не простая задача. Это и это мои ответы дают некоторое представление, поэтому я не буду повторяться в деталях, но для достижения максимальной производительности необходимо использовать кэш и SIMD / pipe- * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *} * * * * Понимание всех вышеперечисленных правил наилучшим образом: Это не невозможно, но нужно быть очень близким с низкоуровневой оптимизацией, чтобы быть успешным - и не стоит ожидать большого улучшения. Наивная реализация вообще не сможет побить 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:
Вот моя реализация с 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% быстрее. Использование -c=/arch:AVX512 (иначе скомпилировано только с /Ox с MSV C), как предложено @ max9111, сделало однопоточную версию примерно на 20% быстрее, но принесло только около 5% для параллельной версии.
Numba:
Существует тот же алгоритм, скомпилированный с Numba (который часто может побить Cython из-за лучшей производительности компилятора clang) - но результат немного медленнее, чем 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 можно по-прежнему видеть издержки для создание / управление потоками).
Существует еще одна более быстрая (по крайней мере, для размеров около 2000) попытка с Cython, которая выполняет суммирование не для небольших частей из 50 элементов, а для всего строка, таким образом, уменьшая издержки (но может иметь больше пропусков кэша, когда строка больше чем 1-2 КБ):
%%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
, который уже быстрее однопоточный (2 мс) и около 60% (1,27 мс ± 50,8 мкс на л oop (среднее ± стандартное отклонение из 7 циклов, по 100 циклов в каждом)) быстрее с распараллеливанием.