When does a numba function compile?

江枫思渺然 提交于 2019-12-06 02:15:16

问题


I'm working off this example:

http://numba.pydata.org/numba-doc/0.15.1/examples.html#multi-threading

and it states that:

You should make sure inner_func is compiled at this point, because the compilation must happen on the main thread. This is the case in this example because we use jit().

It seems in the example that calling jit on a function ensures compilation at that time.

Would the multithreaded example work if instead of calling jit on the function we had used jit with argument types specified as a decorator? I think this is equivalent to asking if the function would be compiled upon definition if jitted with a decorator.

import numba as nb
import numpy as np
def inner_func(result, a, b):
    threadstate = savethread()
    for i in range(len(result)):
        result[i] = np.exp(2.1 * a[i] + 3.2 * b[i])
    restorethread(threadstate)
signature = nb.void(nb.double[:], nb.double[:], nb.double[:])
inner_func_nb = nb.jit(signature, nopython=True)(inner_func)

vs

import numba as nb
import numpy as np
signature = nb.void(nb.double[:], nb.double[:], nb.double[:])
@nb.jit(signature, nopython=True)
def inner_func(result, a, b):
    threadstate = savethread()
    for i in range(len(result)):
        result[i] = np.exp(2.1 * a[i] + 3.2 * b[i])
    restorethread(threadstate)

回答1:


Unless I'm missing something, your two examples are exactly equivalent, not because of anything numba does (or doesn't do), but because that's how decorators work. I'll explain the general principle and you can double-check whether that's the transformation you are asking about. This code:

@d(arg)
def f(x): ...

is by definition equivalent to:

_decorator = d(arg)
def f(x): ...
f = _decorator(f)

which, if we make the reasonable assumptions that d(arg) has no side effects, can be rewritten as:

def f(x): ...
f = d(arg)(f)

The decorator can't even tell the difference (without deliberately using fragile black magic).

The sole difference is that, in the first example, you called the decorated function inner_func_nb instead of replacing inner_func with it. This would cause different behavior if you called inner_func, since in the first example you'd be calling the un-decorated un-jitted function. That doesn't mean the jitting doesn't happen, only that its result is stored under a different name.



来源:https://stackoverflow.com/questions/27510838/when-does-a-numba-function-compile

易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!