The compilation of an array of conditions with numba.jit takes a long time

Question

If I try to compile a function, containing an array of conditions, with numba's jit-compiler, it takes very long. The program looks essentially like

from numba import jit
import numpy as np

@jit(nopython=True)
def foo(a, b):
    valid = [
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0),
        (a - 1 >= 0) and (b - 1 >= 0)
    ]

foo(1, 1)

where I have excluded everything that will not alter the compilation time significantly. The problem arises if I use more than 20 elements.

| elements | time |
-------------------
|    21    | 2.7s |
|    22    | 5.1s |
|    23    |  10s |
|   ...    |  ... |
-------------------

Despite that, the function workes well. Does anybody know, why it takes so long, to compile such function with numba? Creating arrays in a similar way with combinations of integers or floats causes no problem.

Answer 1

1. You may want to report this on the numba issue tracker, feels like something is going awry in the compiler for it to scale so poorly.

2. You may also consider if you really need a large number of array statements like this and if problem could be more clear refactored. E.g. instead of array of booleans, could valid be a function called as needed?

3. That all said, a workaround in the current version of numba is unrolling the conditions.

For your example:

# "codegen"
for i in range(23):
    print(f'    valid[{i}] = (a - 1 >= 0) and (b - 1 >= 0)')

@jit(nopython=True)
def foo(a, b):
    valid = np.empty(23, dtype=np.bool_)
    valid[0] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[1] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[2] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[3] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[4] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[5] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[6] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[7] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[8] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[9] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[10] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[11] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[12] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[13] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[14] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[15] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[16] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[17] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[18] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[19] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[20] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[21] = (a - 1 >= 0) and (b - 1 >= 0)
    valid[22] = (a - 1 >= 0) and (b - 1 >= 0)

%time foo(1,1)
Wall time: 274 ms