为何Numba中该循环的2D版本比1D版本效率高5倍以上？

测试环境 #

• Python 3.12.11
• Numba 0.61.2
• 系统：Ubuntu 22.04.5
• CPU：AMD Ryzen 7 3800X

基准测试代码 #

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, len(img)-w):
    p = img[i]
    fl[i] = 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]-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}')  

测试结果 #

True
Minimum, median and maximum execution time in us:
f2d(img2d)             170.701   172.806   180.690
f1d(img1d, w)          853.687   864.637   873.464

2D版本的执行速度比1D版本快5倍以上，通过分析汇编代码发现两者SIMD指令使用差异极大：f2d中ymm指令出现559次，而f1d仅出现6次。

核心疑问 #

为何逻辑本质相同的循环，f2d能被Numba成功向量化，而f1d却几乎没有使用SIMD优化？

尝试的无循环优化版本 #

针对建议尝试了无循环实现的f1da：

@nb.njit(fastmath=True)
def f1da(img, w):
  p = img[w:-w]
  l = img[w-1:-w-1]
  r = img[w+1:-w+1]
  d = img[:-2*w]
  u = img[2*w:]
  fl = np.maximum(np.maximum(np.abs(p-l), np.abs(p-r)), np.maximum(np.abs(p-u), np.abs(p-d)))

  return fl

优化后的测试结果 #

f1da(img1d, w)         713.653   714.284   727.849

仅获得小幅性能提升，远未达到2D版本的水平。

实际场景说明 #

上述示例是对实际业务场景的简化，实际需求无法通过纯NumPy数组操作的无循环方式实现，必须依赖循环逻辑。

内容的提问来源于stack exchange，提问作者Paul Jurczak