Repository: numba/numba

Issue:
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.  



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!

Resolving patch:
diff --git a/numba/np/arrayobj.py b/numba/np/arrayobj.py
--- 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 range(len(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 range(len(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 range(len(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 range(len(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 range(len(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 range(len(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 range(len(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 range(len(arr_flat)):
+            arr_flat[idx] = 1
         return arr
 
     res = context.compile_internal(builder, ones_like, sig, args)


Test patch: