CODEWITHSUNDEEP
pythonnumpynumba
Fnord
Fnord

Reputation: 5895

Error when compiling a Numba module with function using other functions inline

Numba documentation specifies that other compiled functions can be inlined and called from other compiled functions. This does not seem to be true when compiling ahead of time.

For example: here are two functions that compute the inner dot product between 2 vector arrays, one of them does the actual product, the other makes the inline call within a loop:

# Module test.py
import numpy as np
from numba import njit, float64

@njit(float64(float64[:], float64[:]))
def product(a, b):
    prod = 0
    for i in range(a.size):
        prod += a[i] * b[i]
    return prod

@njit(float64[:](float64[:,:], float64[:,:]))
def n_inner1d(a, b):
    prod = np.empty(a.shape[0])    
    for i in range(a.shape[0]):
        prod[i] = product(a[i], b[i])

    return prod

As is, I can do import test and use test.n_inner1d perfectly fine. Now lets do some modifications so this can be compiled to a .pyd

# Module test.py
import numpy as np
from numba import float64
from numba.pycc import CC

cc = CC('test')
cc.verbose = True

@cc.export('product','float64(float64[:], float64[:])')
def product(a, b):
    prod = 0
    for i in range(a.size):
        prod += a[i] * b[i]
    return prod

@cc.export('n_inner1d','float64[:](float64[:,:], float64[:,:])')
def n_inner1d(a, b):
    prod = np.empty(a.shape[0])    
    for i in range(a.shape[0]):
        prod[i] = product(a[i], b[i])

    return prod

if __name__ == "__main__":
    cc.compile()

When trying to compile, i get the following error:

# python test.py
Failed at nopython (nopython frontend)
Untyped global name 'product': cannot determine Numba type of <type 'function'>
File "test.py", line 20

QUESTION

For a module compiled ahead of time, is it possible for functions defined within to call one another and be used inline?

Upvotes: 3

Views: 1532

Answers (1)

Fnord
Fnord

Reputation: 5895

I reached out to the numba devs and they kindly answered that adding the @njit decorator after @cc.export will make the function call type resolution work and resolve.

So for example:

@cc.export('product','float64(float64[:], float64[:])')
@njit
def product(a, b):
    prod = 0
    for i in range(a.size):
        prod += a[i] * b[i]
    return prod

Will make the product function available to others. The caveat being that it is entirely possible in some cases that the inlined function ends up with a different type signature to that of the one declared AOT.

Upvotes: 5

Related Questions