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

Here is a fix based on the good advice of Homer512 and its related past answer. I corrected the indices so to avoid any relative offset in memory accesses (even having fl[i+w] = ... prevent Numba to optimise the code).

@nb.njit
def f1d(img, w):
  fl = np.zeros_like(img)
  pr = img[w+1:]
  pl = img[w-1:]
  pc = img[w:]
  pt = img[0:]
  pb = img[w+w:]
  out = fl[w:]
  size = len(img) - 2 * w

  for i in range(size):
    p = pc[i]
    out[i] = max(abs(p-pr[i]), abs(p-pl[i]), abs(p-pt[i]), abs(p-pb[i]))

  return fl

I think Numba fails to optimise the initial code because it does not see that there is no overlapping between the input and output arrays (pointer aliasing issue). AFAIK, LLVM can add some checks to generate a more specialised/optimised variant of the loop, but this kind of optimisation is very brittle (and require aggressive optimisations to be enabled internally). Numba could theoretically do that at a higher level (this is obvious here since fl is created in the function) but AFAIK it is not done yet.

The above code is about equally fast to the 2D version. As stated in the comments, both implementation are far from being optimal compared to a naive SIMD-optimised native code (about 6 times faster still with a single thread).

By the way, please note I removed the fastmath flag based it does not apply on integer operations (only floating point ones).

Update: faster implementation

As pointed out by max9111 in comments, the code can be optimised further. Here is a much faster implementation (equally fast with an optimised native code and about 7 times faster than the above one):

@nb.njit
def f1d_fastest(img, w):
  assert w > 0
  fl = np.zeros_like(img)
  pr = img[w+1:]
  pl = img[w-1:]
  pc = img[w:]
  pt = img[0:]
  pb = img[w+w:]
  out = fl[w:]
  size = len(img) - 2 * w

  for i in range(np.uint64(size)):
    p = pc[i]
    out[i] = max(abs(np.int16(p-pr[i])), abs(np.int16(p-pl[i])), abs(np.int16(p-pt[i])), abs(np.int16(p-pb[i])))

  return fl

This code assumes there is not overflow when doing subtractions. If this is the case, then you can replace 16-bit explicit conversions to 32-bit. This is totally fine with the provided random input array.

The same thing applies on the 2D code.

Please note that the 2D code tends to be slower than the 1D one once optimised on machines not supporting the AVX-512 instruction set (e.g. AMD Zen5 CPUs or most Intel server-side CPUs). This is because the SSE/AVX1/AVX2/Neons instruction sets do not provide a way to do masked load/stores (of 16-bit integers) which are critical to compute the borders efficiently (as opposed to AVX-512). The 1D version is not limited by that.

Actually, the 1D version compute the borders differently: it fills them contrary to the 2D version. This means the 1D and 2D version are not strictly equivalent. This is not a problem if you do not care about the borders. If you want to fill them with zeros, then you should include this in the 1D version so to perform a fair comparison (it is not that fast to fill the borders because this can only be done efficiently with scalar instructions).

Memory Access Patterns

f2d (2D Array): When you iterate with img[y, x], img[y, x+1], img[y, x-1], these accesses are largely contiguous in memory (or at least within the same cache line). Even img[y+1, x] and img[y-1, x] access adjacent rows, but within the inner x loop, the data is processed sequentially. Modern CPUs and SIMD instructions (like AVX, using ymm registers) are highly optimized for loading and processing data in contiguous blocks. This allows Numba's LLVM backend to generate highly efficient vectorized code.

f1d (1D Flattened Array with w stride): The issue here is the w (width) parameter, which is 500 in your example. When you access img[i+w] or img[i-w], you are jumping 500 elements (1000 bytes for i2 data type) away from img[i]. These are large, non-contiguous memory strides.

SIMD units work by loading a block of data (e.g., 256 bits or 32 bytes for ymm registers) into a register and then performing operations on all elements in that block simultaneously.

If the data needed for an operation (p, img[i+1], img[i-1], img[i+w], img[i-w]) is scattered across memory, the CPU has to perform multiple, less efficient memory loads (potentially cache misses) to bring the data into the registers. This overhead quickly negates the benefits of vectorization.

Compiler (Numba/LLVM) Optimizations

Numba leverages LLVM, which is a powerful optimizing compiler. LLVM is excellent at vectorizing "embarrassingly parallel" loops where memory accesses are predictable and contiguous. However, loops with complex or large strided memory access patterns are notoriously difficult for compilers to vectorize automatically. The compiler might fall back to scalar (element-by-element) operations for f1d because it cannot efficiently map the strided loads to SIMD instructions. The ymm count difference you observed strongly supports this: f2d is heavily vectorized, f1d is not.

When you perform np.abs(p-u) or np.abs(p-d), NumPy's internal C code will be highly optimized for these element-wise operations. However, the initial creation of the d and u arrays (which involve large strides relative to the original img1d's conceptual 2D structure) still involves non-local memory accesses. While NumPy's C routines are better at handling this than a Numba-compiled Python loop for some cases, the underlying memory fetching bottleneck remains. The operations like np.maximum and np.abs themselves are vectorized, but the overall efficiency is still limited by how the data is laid out and accessed.