Как отметил @Antonio, использование pow
для простого умножения не очень разумно и приводит к значительным накладным расходам:
Таким образом, замена pow(arr[i], 2)
на arr[i]*arr[i]
приводит к довольно большому ускорению:
cython-pow-version 356 µs
numba-version 11 µs
cython-mult-version 14 µs
Оставшееся различие, вероятно, связано с разницей между компиляторами и уровнями оптимизации (в моем случае, llvm против MSVC). Возможно, вы захотите использовать clang для соответствия производительности numba (см., Например, SO-answer )
Чтобы упростить оптимизацию для компилятора, вы должны объявить входные данные как непрерывный массив, то есть double[::1] arr
(см. этот вопрос почему это важно для векторизации), используйте @cython.boundscheck(False)
( используйте опцию -a
, чтобы увидеть, что там меньше желтого), а также добавьте флаги компилятора (например, -O3
, -march=native
или аналогичные в зависимости от вашего компилятора, чтобы включить векторизацию, обратите внимание на флаги сборки, используемые по умолчанию, которые могут запрещать некоторая оптимизация, например -fwrapv ). В конце вы можете написать рабочий цикл в C, скомпилировать с правильной комбинацией flags / compiler и использовать Cython, чтобы обернуть его.
Кстати, введя параметры функции как nb.float64[:](nb.float64[:])
, вы уменьшите производительность numba - больше нельзя допускать, что входной массив является непрерывным, что исключает векторизацию. Пусть numba обнаружит типы (или определит его как непрерывный, т. Е. nb.float64[::1](nb.float64[::1]
), и вы получите лучшую производительность:
@nb.jit(nopython=True)
def nb_vec_f(arr):
res=np.zeros(len(arr))
for i in range(len(arr)):
res[i]=(arr[i])**2
return res
Приводит к следующему улучшению:
%timeit f(arr) # numba version
# 11.4 µs ± 137 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
%timeit nb_vec_f(arr)
# 7.03 µs ± 48.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
И, как указывает @ max9111, нам не нужно инициализировать полученный массив с нулями, но мы можем использовать np.empty(...)
вместо np.zeros(...)
- эта версия даже превосходит numpy's np.square()
Характеристики различных подходов на моей машине:
numba+vectorization+empty 3µs
np.square 4µs
numba+vectorization 7µs
numba missed vectorization 11µs
cython+mult 14µs
cython+pow 356µs