Repository: numba/numba
Pull request: performance improvements to np.full and np.ones
Closes #7665. np.full and np.ones use the slow np.ndindex iterator, which results in both functions being about twice as slow as their numpy counterparts. Iterating over the flat array instead makes the functions as fast as their numpy versions. 

Problem statement:
Poor np.full performace with jit
numba version: 0.54.1

This is the testing code.
```python
import numpy as np
from numba import jit, prange, set_num_threads
import time

set_num_threads(4)

@jit
def f1(shape):
    return np.full(shape, np.nan, np.float64)

@jit
def f2(shape):
    ans = np.empty(shape, np.float64)
    for i in range(ans.shape[0]):
        for j in range(ans.shape[1]):
            ans[i, j] = np.nan

@jit(parallel=True)
def f3(shape):
    ans = np.empty(shape, np.float64)
    for i in prange(ans.shape[0]):
        for j in range(ans.shape[1]):
            ans[i, j] = np.nan

# Warm up
shape = (12, 10)
f1(shape)
f2(shape)
f3(shape)

shape = (128000, 1000)

t0 = time.time()
np.full(shape, np.nan, np.float64)
print(time.time() - t0)

t0 = time.time()
f1(shape)
print(time.time() - t0)

t0 = time.time()
f2(shape)
print(time.time() - t0)

t0 = time.time()
f3(shape)
print(time.time() - t0)
```
This is the output
```
0.24103260040283203 # numpy
0.5575239658355713 # np.full in jit
0.45681166648864746 # fill with for-loop
0.18029499053955078 # fill with parallel for-loop
```
The speed of `np.full` in jit is less than half compared with the normal one from my testing.

Hints from issue discussion:
This is similar to the problem outlined in https://numba.readthedocs.io/en/stable/user/faq.html#can-numba-speed-up-short-running-functions - there's nothing Numba can do to optimize inside `np.full`, so when you wrap it in a jitted function you're just adding the overhead of calling a jitted function on top of the time it takes to execute `np.full`.

If you're trying to understand whether Numba can help speed up your application, it would be more representative to benchmark something that looks a bit more like a portion of your application code.
But this does not look like just overhead. This is the output that I increase the shape 10 times up. (`shape = (1280000, 1000)`)
```
2.9751839637756348
5.439911603927612
4.502171516418457
1.7119011878967285
```

OK, I see. There's probably some difference in implementation between Numba and NumPy's implementation - we could compare with the NumPy implementation to see if there's a more efficient way we can implement it.
I can't exactly reproduce this, but the problem is that numba's np.full uses np.ndindex as the iterator, which is quite slow. Some of the other numba array creation routines that are fill element-by-element use array.flat and iterate over it using enumerate, which is a good bit faster. 

```python
@jit
def nb_full(shape, fill_value, dtype):
    ans = np.empty(shape, dtype)
    fl = ans.flat
    for idx, v in enumerate(fl):
      fl[idx] = fill_value
    return ans

@jit
def f4(shape):
  return nb_full(shape, np.nan, np.float64)

@jit
def f1(shape):
    return np.full(shape, np.nan, np.float64)

f1(shape)
f4(shape)

t0 = time.time()
np.full(shape, np.nan, np.float64)
print(time.time() - t0)

t0 = time.time()
f1(shape)
print(time.time() - t0)

t0 = time.time()
f4(shape)
print(time.time() - t0)
```
On my end, this outputs the following:
```python
0.1057584285736084 # numpy
0.2277970314025879 # current implementation
0.10829496383666992 # array.flat implementation 
```

This at least matches the numpy implementation, which should be as good as it gets. 
@rishi-kulkarni Many thanks for providing the additional info and example! Would you be keen to modify Numba's `flat` implementation so that it uses the same process as your `nb_full` function and submit a PR?
@gmarkall Yep, not a problem. I'll take a crack at it this weekend. 
> @gmarkall Yep, not a problem. I'll take a crack at it this weekend.

Many thanks, that's very kind of you - please though, don't feel any pressure to hurry on it... Whenever it's convenient for you to work on it, that will be more than great!

Commit message: performance improvements to np.full and np.ones

Patch to review:
diff --git a/numba/np/arrayobj.py b/numba/np/arrayobj.py
index a7046a404f4..6722193bb83 100644
--- a/numba/np/arrayobj.py
+++ b/numba/np/arrayobj.py
@@ -3789,8 +3789,9 @@ def numpy_full_nd(context, builder, sig, args):
 
     def full(shape, value):
         arr = np.empty(shape, type(value))
-        for idx in np.ndindex(arr.shape):
-            arr[idx] = value
+        arr_flat = arr.flat
+        for idx, _ in enumerate(arr_flat):
+            arr_flat[idx] = value
         return arr
 
     res = context.compile_internal(builder, full, sig, args)
@@ -3803,8 +3804,9 @@ def numpy_full_dtype_nd(context, builder, sig, args):
 
     def full(shape, value, dtype):
         arr = np.empty(shape, dtype)
-        for idx in np.ndindex(arr.shape):
-            arr[idx] = value
+        arr_flat = arr.flat
+        for idx, _ in enumerate(arr_flat):
+            arr_flat[idx] = value
         return arr
 
     res = context.compile_internal(builder, full, sig, args)
@@ -3816,8 +3818,9 @@ def numpy_full_like_nd(context, builder, sig, args):
 
     def full_like(arr, value):
         arr = np.empty_like(arr)
-        for idx in np.ndindex(arr.shape):
-            arr[idx] = value
+        arr_flat = arr.flat
+        for idx, _ in enumerate(arr_flat):
+            arr_flat[idx] = value
         return arr
 
     res = context.compile_internal(builder, full_like, sig, args)
@@ -3830,8 +3833,9 @@ def numpy_full_like_nd_type_spec(context, builder, sig, args):
 
     def full_like(arr, value, dtype):
         arr = np.empty_like(arr, dtype)
-        for idx in np.ndindex(arr.shape):
-            arr[idx] = value
+        arr_flat = arr.flat
+        for idx, _ in enumerate(arr_flat):
+            arr_flat[idx] = value
         return arr
 
     res = context.compile_internal(builder, full_like, sig, args)
@@ -3843,8 +3847,9 @@ def numpy_ones_nd(context, builder, sig, args):
 
     def ones(shape):
         arr = np.empty(shape)
-        for idx in np.ndindex(arr.shape):
-            arr[idx] = 1
+        arr_flat = arr.flat
+        for idx, _ in enumerate(arr_flat):
+            arr_flat[idx] = 1
         return arr
 
     valty = sig.return_type.dtype
@@ -3859,8 +3864,9 @@ def numpy_ones_dtype_nd(context, builder, sig, args):
 
     def ones(shape, dtype):
         arr = np.empty(shape, dtype)
-        for idx in np.ndindex(arr.shape):
-            arr[idx] = 1
+        arr_flat = arr.flat
+        for idx, _ in enumerate(arr_flat):
+            arr_flat[idx] = 1
         return arr
 
     res = context.compile_internal(builder, ones, sig, args)
@@ -3872,8 +3878,9 @@ def numpy_ones_like_nd(context, builder, sig, args):
 
     def ones_like(arr):
         arr = np.empty_like(arr)
-        for idx in np.ndindex(arr.shape):
-            arr[idx] = 1
+        arr_flat = arr.flat
+        for idx, _ in enumerate(arr_flat):
+            arr_flat[idx] = 1
         return arr
 
     res = context.compile_internal(builder, ones_like, sig, args)
@@ -3886,8 +3893,9 @@ def numpy_ones_like_dtype_nd(context, builder, sig, args):
 
     def ones_like(arr, dtype):
         arr = np.empty_like(arr, dtype)
-        for idx in np.ndindex(arr.shape):
-            arr[idx] = 1
+        arr_flat = arr.flat
+        for idx, _ in enumerate(arr_flat):
+            arr_flat[idx] = 1
         return arr
 
     res = context.compile_internal(builder, ones_like, sig, args)


Code review comments:
diff_hunk: @@ -3803,8 +3804,9 @@ def numpy_full_dtype_nd(context, builder, sig, args):
 
     def full(shape, value, dtype):
         arr = np.empty(shape, dtype)
-        for idx in np.ndindex(arr.shape):
-            arr[idx] = value
+        arr_flat = arr.flat
+        for idx, _ in enumerate(arr_flat):
path: numba/np/arrayobj.py
text: @user1:
Suggest using `for idx in range(len(arr_flat))` (or similar)? There's no need to set up the more complex loop structure?

@author:
Sure, done. Performance is identical. 

@user1:
Thanks. In the isolated case I'd expect/hope the performance gain would be the same, but in a more involved case, the more simple the loop structure, in general, the more chance LLVM has of doing optimisations.