tl; dr Numpy использует C функций памяти, в то время как Numba должна присваивать нули
Я написал сценарий для построения графика времени, необходимого для завершения нескольких вариантов, и он появляется что Numba сильно падает в производительности, когда размер массива np.zeros достигает 2048*2048*8 = 32 MB на моей машине, как показано на диаграмме ниже.
Реализация Numba np.zeros столь же быстра, как и создание пустого массива и заполнение его нулями путем итерации по размерам массива (это вложенная Numba l oop зеленая кривая диаграммы). На самом деле это можно проверить дважды, установив переменную окружения NUMBA_DUMP_IR перед запуском скрипта (см. Ниже). При сравнении с дампом для numba_loop разница невелика.
Интересно, что np.zeros немного повысил порог 32 МБ.
Мое лучшее предположение, хотя я далеко По мнению эксперта, ограничение в 32 МБ - это узкое место в ОС или аппаратном обеспечении, возникающее из объема данных, которые могут поместиться в кэш для того же процесса. Если это превышено, операция перемещения данных в кеш и из него для работы с ним занимает очень много времени.
В отличие от этого, Numpy использует callo c, чтобы получить некоторый сегмент памяти с обещанием заполнить данные нулями при получении доступа к ним.
Это это то, как далеко я прошел, и я понимаю, что это только половина ответа, но, возможно, кто-то более знающий может пролить свет на то, что на самом деле происходит.
Numba IR dump:
---------------------------IR DUMP: pure_numba_zeros----------------------------
label 0:
n = arg(0, name=n) ['n']
$2load_global.0 = global(np: <module 'numpy' from '/lib/python3.8/site-packages/numpy/__init__.py'>) ['$2load_global.0']
$4load_attr.1 = getattr(value=$2load_global.0, attr=zeros) ['$2load_global.0', '$4load_attr.1']
del $2load_global.0 []
$10build_tuple.4 = build_tuple(items=[Var(n, script.py:15), Var(n, script.py:15)]) ['$10build_tuple.4', 'n', 'n']
$12load_global.5 = global(np: <module 'numpy' from '/lib/python3.8/site-packages/numpy/__init__.py'>) ['$12load_global.5']
$14load_attr.6 = getattr(value=$12load_global.5, attr=double) ['$12load_global.5', '$14load_attr.6']
del $12load_global.5 []
$18call_function_kw.8 = call $4load_attr.1($10build_tuple.4, func=$4load_attr.1, args=[Var($10build_tuple.4, script.py:15)], kws=[('dtype', Var($14load_attr.6, script.py:15))], vararg=None) ['$10build_tuple.4', '$14load_attr.6', '$18call_function_kw.8', '$4load_attr.1']
del $4load_attr.1 []
del $14load_attr.6 []
del $10build_tuple.4 []
mat = $18call_function_kw.8 ['$18call_function_kw.8', 'mat']
del $18call_function_kw.8 []
$24load_method.10 = getattr(value=mat, attr=reshape) ['$24load_method.10', 'mat']
del mat []
$const28.12 = const(int, 2) ['$const28.12']
$30binary_power.13 = n ** $const28.12 ['$30binary_power.13', '$const28.12', 'n']
del n []
del $const28.12 []
$32call_method.14 = call $24load_method.10($30binary_power.13, func=$24load_method.10, args=[Var($30binary_power.13, script.py:16)], kws=(), vararg=None) ['$24load_method.10', '$30binary_power.13', '$32call_method.14']
del $30binary_power.13 []
del $24load_method.10 []
$34return_value.15 = cast(value=$32call_method.14) ['$32call_method.14', '$34return_value.15']
del $32call_method.14 []
return $34return_value.15 ['$34return_value.15']
Скрипт для создания диаграммы:
import numpy as np
from numba import jit
from time import time
import os
import matplotlib.pyplot as plt
os.environ['NUMBA_DUMP_IR'] = '1'
def numpy_zeros(n):
mat = np.zeros((n,n), dtype=np.double)
return mat.reshape((n**2))
@jit(nopython=True)
def numba_zeros(n):
mat = np.zeros((n,n), dtype=np.double)
return mat.reshape((n**2))
@jit(nopython=True)
def numba_loop(n):
mat = np.empty((n * 2,n), dtype=np.float32)
for i in range(mat.shape[0]):
for j in range(mat.shape[1]):
mat[i, j] = 0.
return mat.reshape((2 * n**2))
# To compile
numba_zeros(10)
numba_loop(10)
os.environ['NUMBA_DUMP_IR'] = '0'
max_n = 4100
time_deltas = {
'numpy_zeros': [],
'numba_zeros': [],
'numba_loop': [],
}
call_count = 10
for n in range(0, max_n, 10):
for f in (numpy_zeros, numba_zeros, numba_loop):
start = time()
for i in range(call_count):
x = f(n)
delta = time() - start
time_deltas[f.__name__].append(delta / call_count)
print(f'{f.__name__:25} n = {n}: {delta}')
print()
size = np.arange(0, max_n, 10) ** 2 * 8 / 1024 ** 2
fig, ax = plt.subplots()
plt.xticks(np.arange(0, size[-1], 16))
plt.axvline(x=32, color='gray', lw=0.5)
ax.plot(size, time_deltas['numpy_zeros'], label='Numpy zeros (calloc)')
ax.plot(size, time_deltas['numba_zeros'], label='Numba zeros')
ax.plot(size, time_deltas['numba_loop'], label='Numba nested loop')
ax.set_xlabel('Size of array in MB')
ax.set_ylabel(r'Mean $\Delta$t in s')
plt.legend(loc='upper left')
plt.show()