[Python] Why is Numba more efficient with 2D vs 1D version of this loop?

I'm using Python 3.12.11 with Numba 0.61.2 on Ubuntu 22.04.5 and AMD Ryzen 7 3800X CPU. This benchmark:

import numpy as np, timeit as ti, numba as nb

@nb.njit(fastmath=True)
def f1d(img, w):
fl = np.zeros_like(img)

for i in range(w+1, len(img)-w-1):
p = img
fl = max(abs(p-img[i+1]), abs(p-img[i+w]), abs(p-img[i-w]), abs(p-img[i-1]))

return fl


@nb.njit(fastmath=True)
def f2d(img):
fl = np.zeros_like(img)

for y in range(1, img.shape[0]-1):
for x in range(1, img.shape[1]):
p = img[y, x]
fl[y, x] = max(abs(p-img[y, x+1]), abs(p-img[y, x-1]), abs(p-img[y+1, x]), abs(p-img[y-1, x]))

return fl


img2d = np.random.randint(256, size=(500, 500)).astype('i2')
img1d = img2d.ravel().astype('i2')
w = img2d.shape[1]
print(np.array_equal(f2d(img2d)[:, 1:-1], f1d(img1d, w).reshape((w, -1))[:, 1:-1]))
print(f'Minimum, median and maximum execution time in us:')

for fun in ('f2d(img2d)', 'f1d(img1d, w)'):
t = 10**6 * np.array(ti.repeat(stmt=fun, setup=fun, globals=globals(), number=1, repeat=99))
print(f'{fun:20} {np.amin(t):8,.3f} {np.median(t):8,.3f} {np.amax(t):8,.3f}')


produces:

True
Minimum, median and maximum execution time in us:
f2d(img2d) 171.773 173.927 182.363
f1d(img1d, w) 901.357 908.651 920.172


2D version is over 5x faster than 1D version, due to heavy use of SIMD: ymm is present 559 times in f2d assembly vs. 6 times in f1d. What is the reason that essentially the same loop is vectorized for f2d, but not for f1d?

Continue reading...