Poor np.full performace with jit

numba version: 0.54.1

This is the testing code.

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.

[rishi-kulkarni]

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.

@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:

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.

[szsdk]

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