Skip to content

How numba resolves function calls?

Jump to bottom
Pearu Peterson edited this page Mar 5, 2020 · 6 revisions

Introduction

Numba jit-decorated functions can call other functions provided that the types of these functions can be determined.

Prerequisities

A Python function, say foo, can be Numba jit-compiled in three ways:

  1. Using "nopython"-mode that performs type-inference to compile the function foo to use low-level types so that the execution of foo is performed directly on the low-level values (rather than on Python values that holding these). The "nopython"-mode can be used only when type-inference is possible, that is, the low-level types of all variables can be determined uniquely. The Numba jit-compiled is always faster than the corresponding pure Python program because of the eliminated Python layer but when interpreting benchmark results one needs to take into account the overheads of Numba jit-compilation, converting input Python values to low-level values, and converting the low-level returned values to Python return values.

    The "nopython"-mode compilation can be triggered using the expression

    numba.jit(foo, nopython=True).compile(<a tuple of foo argument Numba types>)

  2. Using "object"-mode that compiles the function foo to use CPython calls so that the execution of foo is performed on Python values using high-level Python operations. The "object"-mode can be used for any valid Python program. The Numba jit-compiled code is not generally faster than the corresponding pure Python program.

    The "object"-mode compilation can be triggered using the expression

    numba.jit(foo, forceobj=True).compile(<a tuple of foo argument Numba types>)

  3. Using "default"-mode that tries "nopython"-mode first. When type-inference is unsuccessful then Numba falls back to using "object"-mode.

    The "default"-mode compilation can be triggered using the expression

    numba.jit(foo).compile(<a tuple of foo argument Numba types>)

Numba types are defined in subpackage numba.core.types that provides various low-level types such as int32, int64, float32, etc. To represent function signatures, Numba defines Signature class (defined in numba/core/typing/templates.py) to hold the return Numba type (attribute return_type), the arguments Numba types (attribute args), as well as inspect.Signature instance (attribute .pysig).

Calling a Python function from Numba jit-decorated functions

Consider the following example

def f(x):
    return x + 1

def foo(x):
    return f(x)

The Numba jit-compilation of foo in "nopython"-mode (say, using numba.jit(foo, nopython=True).compile((numba.int32,))) raises:

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Untyped global name 'f': cannot determine Numba type of <class 'function'>

while in "python"-mode the compilation will succeesful.

Notice that the function f is pure Python object that Numba (<=0.48) cannot resolve to the corresponding low-level function type and hence the failure in "nopython"-mode. However, there exists a unique low-level function type for the given use case that matches with the low-level function type of foo (provided that the Numba argument type numba.int32 supports low-level addition operation with integers).

Moreover, decorating the function f with numba.jit leads to successful type-inference of foo without giving any additional argument type information to the Numba type-inference algorithm. This brings up a question of why Numba type-inference algorithm would not do that implicitly? Are there use cases where resolving the types of pure Python functions automatically would not be desired? See also callback into Python form njit'ed code that illustrates the case where the argument type information can be given by other means than jit-decoration.

Next, we analyze how Numba type-inference works to fix resolving pure Python functions in "nopython"-mode. This is important for the First-class functions support PR 4967.

Compilation of a jit-decorated function

A jit-decorated function is a Dispatcher class instance with the following data:

  • .py_func attribute holding the pure Python function being jit-decorated,
  • .overloads attribute holding a mapping of signature argument types and compilation results,
  • ._compiler attribute holding a _FunctionCompiler instance.

(The actual implementation of Dispatcher class is technically more involved as it is inherited from a CPython extension type that implements a fast signature matching algorithm in C).

Numba jit-compilation is realized by the call

Dispatcher.compile(<argument Numba types or Signature instance>) -> <compilation result>.entry_point

that implements the following algorithm:

  1. Set args = tuple(<argument Numba types>), return_type = <None or function return Numba type>
  2. If args in dispatcher.overloads, return the result.
  3. Otherwise compute cres = dispatcher._compiler.compile(args, return_type), store it in dispatcher.overloads, and return the result.

Calling _FunctionCompile.compile(args, return_type) method boils down to

  1. creating a Compiler instance (in numba.core.compiler.compile_extra function)
  2. getting the bytecode of the dispatcher.py_func as a numba.core.bytecode.FunctionIdentity instance (stored in compiler.state.func_id)
  3. decoding compiler.state.func_id bytecode representation to numba.core.bytecode.ByteCode instance (stored in compiler.state.bc, use NUMBA_DUMP_BYTECODE=1 to see the result)
  4. construct a compilation pipeline to numba.core.compiler_machinery.PassManager instance that will store a sequence of passes that are described below.
  5. apply pipeline passes to compiler.state dictionary (in numba.core.compiler.Compiler._compile_core method). The compile result is stored in compiler.state.rc as numba.core.compiler.CompileResult instance that is also returned by _FunctionCompile.compile method.

The numba.core.compiler.Compiler.define_pipelines constructs a sequence of passes that are applied to the compiler.state dictionary as follows:

  1. TranslateByteCode translates compiler.state.bc to compiler.state.func_ir containing the so-called Numba IR as a numba.core.ir.FunctionIR instance. Use NUMBA_DUMB_CFG=1 to see the Control Flow Graph.
  2. FixupArgs checks that the function and signature have matching number of arguments.
  3. IRProcessing post-processes Numba IR (canonicalizes the CFG, emits del for variables, computes lifetime of variables, etc). Use NUMBA_DUMP_IR=1 to see the Numba IR of the function.
  4. WithLifting extracts with constructs, compile the results, and return compilation result via _EarlyPipelineCompletion exception to _FunctionCompile.compile. If no with statements, continue.
  5. RewriteSemanticConstants rewrites values known to be constants. Set DEBUG=2 in numba.core.analysis.rewrite_semantic_constants to see the resulting Numba IR.
  6. DeadBranchPrune prunes dead branches. Set DEBUG=2 in numba.core.analysis.dead_branch_prune to see the resulting Numba IR.
  7. GenericRewrites applies before-inference rewrites (constant get/set items, constant raises, constant binop arguments, expand Macro lookups and calls)
  8. InlineClosureLikes inlines locally defined closures. Set NUMBA_DEBUG_INLINE_CLOSURE=1 to see debug messages and NUMBA_DUMP_IR=1 to see the resulting Numba IR.
  9. MakeFunctionToJitFunction converts make_function to njit-decorated function
  10. InlineInlinables inlines jit-decorated functions into the site of its call. Set numba.core.untyped_passes.InlineInlinables._DEBUG = True and NUMBA_DEBUG_INLINE_CLOSURE=1 to see debug messages.
  11. FindLiterallyCalls handles numba.literally calls.
  12. LiteralUnroll handles numba.literal_unroll calls.
  13. NopythonTypeInference allows input array returns and rejects pure Python function returns in "nopython"-mode.
  14. AnnotateTypes sets compiler.state.type_annotation as numba.core.annotations.type_annotations.TypeAnnotation instances. Set NUMBA_DUMP_ANNOTATION=1 or NUMBA_DUMP_HTML=<path to html file> to see annotations.
  15. InlineOverloads inline numba.extending.overload decorated functions. Set numba.core.types_passes.InlineOverloads._DEBUG = True to see debug messages.
  16. PreParforPass preprocesses for parfors.
  17. NopythonRewrites applies after-inference rewrites (array creation from list without creating a list, rewrite array expressions as ufunc-like calls)
  18. ParforPass convert data-parallel computations to parfors.
  19. IRLegalization check and legalize invalid IR
  20. NoPythonBackend generates LLVM IR, creates LLVM IR library
  21. DumpParforDiagnostics.

The LLVM IR module of the compiled jit-decorated function contains two functions:

  1. a low-level implementation of the jit-decorated function with a signature 
    <mangled function name>(<return value type>* retptr, <excinfo struct>** excinfo, <argument 0 type> arg0, ...) -> int32 status
  2. a cpython wrapper of the low-level implementation function that takes Python objects as input, converts the Python input values to low-level values, calls the low-level function, converts the low-level return value to Python object, and returns it. The cpython wrapper has a signature: 
    <cpython prefix><mangled function name>(closure, args, kws) -> <return value>
    The cpython wrapper is built in numba.core.cpu.CPUContext.create_cpython_wrapper method.

Once a Dispatcher instance is compiled for a given set of signatures, one can disable any future compilations by calling dispatcher.disable_compile().

Calling jit-decorated functions

... when compilation is enabled

When a jit-decorated function is applied to specific argument values (the call is implemented a Dispatcher_call in numba/_dispatcher.c), Numba finds a compilation result that matches exactly with the given arguments types (see dispatcher_resolve in numba/_dispatcherimpl.cpp). If exactly one match is found, it's entry point is called. In the case of zero or multiple matches, a new compilation is triggered for the given specialization (see compile_and_invoke in numba/_dispatcher.c) by calling Dispatcher._compile_for_args(*args) -> <compilation result>.entry_point method.

... when compilation is disabled

When a jit-decorated function is applied to specific argument values (the call is implemented a Dispatcher_call in numba/_dispatcher.c), Numba finds a compilation result that matches exactly with the given arguments types. If exactly one match is found, it's entry point is called. If not match is found, Dispatcher._search_new_conversions(*args) -> bool method is called that installs new conversions if can_convert of argument types allows; finally, another dispatcher resolve is triggered if a new conversion is found. If still no match is found and "python"-mode is effective, then a call to pure Python function call will be made. Otherwise, Dispatcher._explain_matching_error is called.

Storing compilation results

Numba jit-decorated functions are Dispatcher instances that store compilation results in the attribute overloads. The overloads attribute is a (sorted-)dictionary with tuples of Numba types as keys and CompileResult instances (see numba/core/compiler.py) as values. The CompileResult structure holds the following fields:

  • signature - Signature instance
  • library - contains LLVM IR
  • typing_context, target_context, entry_point, type_annotation, etc.

Note that when a Numba jit-decorated function is used as an argument to a Numba jit-decorated function, then the corresponding argument Numba type is a numba.types.Dispatcher instance that holds the Dispatcher instance (in attribute dispatcher). So, each Numba jit-decorated function defines its own unique Numba type, despite the fact that the Dispatcher instance can contain multiple (argument type specific) realizations of the same Python function.

First-class function type

The first-class function type support aims at representing all realizations of Python functions that have the same Numba signature. The first-class function type is represented as FunctionType (see numba/core/types/function.py) that instances hold the Numba types of a function arguments and return values.

Numba FunctionType and Dispatcher types are dual: the former represents functions that realizations have the same proto-type and the latter represents a single function that may have many realizations with all different proto-types.

When a callable is used as an argument to a Numba jit-decorated function, the corresponding Dispatcher instance needs to decide if the callable argument is interpreted as a dispatcher object or as a first-class function object:

  1. If a callable argument is interpreted as a dispatcher object then the Dispatcher instance performs type-inference which leads to a new realization of the instance with callable realization embedded (that is also a result of the type-inference, inlined or to be called as static function). The new dispatcher's overloads key will contain the Numba dispatcher type.

  2. If a callable argument is interpreted as first-class function object then it's prototype(s) can be used in type-inference or type-inference can determine the callable's prototype. In either case, the realizations of the callable and of the dispatcher will be stored independently and are re-usable. The new dispatcher's overloads key will contain Numba function type.

If a callable is Numba cfunc-decorated function or WAP object, then the callable prototype is known and type-inference can use it as fixed input. If a callable is Numba jit-decorated function or a Python function with no pre-specified prototypes then type-inference should resolve the callable type or leave it as undefined (for instance, when the callable is never used in the dispatcher call). If a callable already contains realizations, the type-inference should try to use their prototypes as possible inputs.

Numba jit-decorated functions as first-class functions

Consider the following example

@numba.jit
def foo(f, x):
    return f(x)

@numba.jit
def a(x):
    return x + 1

When making a call to foo with a as the first argument, the specializations of both functions will depend on the type of foo's second argument: if x is an integer value, then

a:   int64(int64)
foo: int64(int64(int64), int64)

If x is a float value, then

a:   float64(float64)
foo: float64(float64(float64), float64)

Etc.

Note that even if a contains a specialization for int64(int64), it cannot be use if x is a float, because of the cast int64(x) would be inexact. On the other hand, if a contains a specialization for float64(float64), it could be used if x is a int. However, this would mean that the type of the result foo(a, 2) will depend on whether foo(a, 2.5) was called before (leading to float64 return value) or not (leading to int64 value). So, in the following we will require that the types of input arguments must match exactly with the corresponding types of jit-decorated function specializations when resolving the compiled result for given arguments.

The above note means that when a is used as a first-type function argument to foo, its existing specializations can be used only when the types of the rest of foo arguments match exactly with the existing specializations of foo. If not exact match exists, a should be treated as a Numba jit-decorated function with no existing specializations (even when a.overloads is non-empty). To ensure that Numba type-inference will not use the existing specializations of a, it may be necessary to use a_ = numba.jit(a.py_func) in type-inference (and later do a.overloads.update(a_.overloads)).

Clone this wiki locally