High-level extension API

This extension API is exposed through the numba.extending module.

To aid debugging extensions to Numba, it’s recommended to set the following environment variable:

NUMBA_CAPTURED_ERRORS="new_style"

this makes it easy to differentiate between errors in implementation and acceptable errors that can take part in e.g. type inference. For more information see NUMBA_CAPTURED_ERRORS.

Implementing functions

The @overload decorator allows you to implement arbitrary functions for use in nopython mode functions. The function decorated with @overload is called at compile-time with the types of the function’s runtime arguments. It should return a callable representing the implementation of the function for the given types. The returned implementation is compiled by Numba as if it were a normal function decorated with @jit. Additional options to @jit can be passed as dictionary using the jit_options argument.

For example, let’s pretend Numba doesn’t support the len() function on tuples yet. Here is how to implement it using @overload:

from numba import types
from numba.extending import overload

@overload(len)
def tuple_len(seq):
   if isinstance(seq, types.BaseTuple):
       n = len(seq)
       def len_impl(seq):
           return n
       return len_impl

You might wonder, what happens if len() is called with something else than a tuple? If a function decorated with @overload doesn’t return anything (i.e. returns None), other definitions are tried until one succeeds. Therefore, multiple libraries may overload len() for different types without conflicting with each other.

Implementing methods

The @overload_method decorator similarly allows implementing a method on a type well-known to Numba.

numba.core.extending.overload_method(typ, attr, **kwargs)

A decorator marking the decorated function as typing and implementing method attr for the given Numba type in nopython mode.

kwargs are passed to the underlying @overload call.

Here is an example implementing .take() for array types:

@overload_method(types.Array, 'take')
def array_take(arr, indices):
    if isinstance(indices, types.Array):
        def take_impl(arr, indices):
            n = indices.shape[0]
            res = np.empty(n, arr.dtype)
            for i in range(n):
                res[i] = arr[indices[i]]
            return res
        return take_impl

Implementing classmethods

The @overload_classmethod decorator similarly allows implementing a classmethod on a type well-known to Numba.

numba.core.extending.overload_classmethod(typ, attr, **kwargs)

A decorator marking the decorated function as typing and implementing classmethod attr for the given Numba type in nopython mode.

Similar to overload_method.

Here is an example implementing a classmethod on the Array type to call np.arange():

@overload_classmethod(types.Array, "make")
def ov_make(cls, nitems):
    def impl(cls, nitems):
        return np.arange(nitems)
    return impl

The above code will allow the following to work in jit-compiled code:

@njit
def foo(n):
    return types.Array.make(n)

Implementing attributes

The @overload_attribute decorator allows implementing a data attribute (or property) on a type. Only reading the attribute is possible; writable attributes are only supported through the low-level API.

The following example implements the nbytes attribute on Numpy arrays:

@overload_attribute(types.Array, 'nbytes')
def array_nbytes(arr):
   def get(arr):
       return arr.size * arr.itemsize
   return get

Importing Cython Functions

The function get_cython_function_address obtains the address of a C function in a Cython extension module. The address can be used to access the C function via a ctypes.CFUNCTYPE() callback, thus allowing use of the C function inside a Numba jitted function. For example, suppose that you have the file foo.pyx:

from libc.math cimport exp

cdef api double myexp(double x):
    return exp(x)

You can access myexp from Numba in the following way:

import ctypes
from numba.extending import get_cython_function_address

addr = get_cython_function_address("foo", "myexp")
functype = ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double)
myexp = functype(addr)

The function myexp can now be used inside jitted functions, for example:

@njit
def double_myexp(x):
    return 2*myexp(x)

One caveat is that if your function uses Cython’s fused types, then the function’s name will be mangled. To find out the mangled name of your function you can check the extension module’s __pyx_capi__ attribute.

Implementing intrinsics

The @intrinsic decorator is used for marking a function func as typing and implementing the function in nopython mode using the llvmlite IRBuilder API. This is an escape hatch for expert users to build custom LLVM IR that will be inlined into the caller, there is no safety net!

The first argument to func is the typing context. The rest of the arguments corresponds to the type of arguments of the decorated function. These arguments are also used as the formal argument of the decorated function. If func has the signature foo(typing_context, arg0, arg1), the decorated function will have the signature foo(arg0, arg1).

The return values of func should be a 2-tuple of expected type signature, and a code-generation function that will passed to lower_builtin(). For an unsupported operation, return None.

Here is an example that cast any integer to a byte pointer:

from numba import types
from numba.extending import intrinsic

@intrinsic
def cast_int_to_byte_ptr(typingctx, src):
    # check for accepted types
    if isinstance(src, types.Integer):
        # create the expected type signature
        result_type = types.CPointer(types.uint8)
        sig = result_type(types.uintp)
        # defines the custom code generation
        def codegen(context, builder, signature, args):
            # llvm IRBuilder code here
            [src] = args
            rtype = signature.return_type
            llrtype = context.get_value_type(rtype)
            return builder.inttoptr(src, llrtype)
        return sig, codegen

it may be used as follows:

from numba import njit

@njit('void(int64)')
def foo(x):
    y = cast_int_to_byte_ptr(x)

foo.inspect_types()

and the output of .inspect_types() demonstrates the cast (note the uint8*):

def foo(x):

    #   x = arg(0, name=x)  :: int64
    #   $0.1 = global(cast_int_to_byte_ptr: <intrinsic cast_int_to_byte_ptr>)  :: Function(<intrinsic cast_int_to_byte_ptr>)
    #   $0.3 = call $0.1(x, func=$0.1, args=[Var(x, check_intrin.py (24))], kws=(), vararg=None)  :: (uint64,) -> uint8*
    #   del x
    #   del $0.1
    #   y = $0.3  :: uint8*
    #   del y
    #   del $0.3
    #   $const0.4 = const(NoneType, None)  :: none
    #   $0.5 = cast(value=$const0.4)  :: none
    #   del $const0.4
    #   return $0.5

    y = cast_int_to_byte_ptr(x)

Implementing mutable structures

Warning

This is an experimental feature, the API may change without warning.

The numba.experimental.structref module provides utilities for defining mutable pass-by-reference structures, a StructRef. The following example demonstrates how to define a basic mutable structure:

Defining a StructRef

from numba/tests/doc_examples/test_structref_usage.py
 1import numpy as np
 2
 3from numba import njit
 4from numba.core import types
 5from numba.experimental import structref
 6
 7from numba.tests.support import skip_unless_scipy
 8
 9
10# Define a StructRef.
11# `structref.register` associates the type with the default data model.
12# This will also install getters and setters to the fields of
13# the StructRef.
14@structref.register
15class MyStructType(types.StructRef):
16    def preprocess_fields(self, fields):
17        # This method is called by the type constructor for additional
18        # preprocessing on the fields.
19        # Here, we don't want the struct to take Literal types.
20        return tuple((name, types.unliteral(typ)) for name, typ in fields)
21
22
23# Define a Python type that can be use as a proxy to the StructRef
24# allocated inside Numba. Users can construct the StructRef via
25# the constructor for this type in python code and jit-code.
26class MyStruct(structref.StructRefProxy):
27    def __new__(cls, name, vector):
28        # Overriding the __new__ method is optional, doing so
29        # allows Python code to use keyword arguments,
30        # or add other customized behavior.
31        # The default __new__ takes `*args`.
32        # IMPORTANT: Users should not override __init__.
33        return structref.StructRefProxy.__new__(cls, name, vector)
34
35    # By default, the proxy type does not reflect the attributes or
36    # methods to the Python side. It is up to users to define
37    # these. (This may be automated in the future.)
38
39    @property
40    def name(self):
41        # To access a field, we can define a function that simply
42        # return the field in jit-code.
43        # The definition of MyStruct_get_name is shown later.
44        return MyStruct_get_name(self)
45
46    @property
47    def vector(self):
48        # The definition of MyStruct_get_vector is shown later.
49        return MyStruct_get_vector(self)
50
51
52@njit
53def MyStruct_get_name(self):
54    # In jit-code, the StructRef's attribute is exposed via
55    # structref.register
56    return self.name
57
58
59@njit
60def MyStruct_get_vector(self):
61    return self.vector
62
63
64# This associates the proxy with MyStructType for the given set of
65# fields. Notice how we are not constraining the type of each field.
66# Field types remain generic.
67structref.define_proxy(MyStruct, MyStructType, ["name", "vector"])

The following demonstrates using the above mutable struct definition:

from test_type_definition of numba/tests/doc_examples/test_structref_usage.py
 1# Let's test our new StructRef.
 2
 3# Define one in Python
 4alice = MyStruct("Alice", vector=np.random.random(3))
 5
 6# Define one in jit-code
 7@njit
 8def make_bob():
 9    bob = MyStruct("unnamed", vector=np.zeros(3))
10    # Mutate the attributes
11    bob.name = "Bob"
12    bob.vector = np.random.random(3)
13    return bob
14
15bob = make_bob()
16
17# Out: Alice: [0.5488135  0.71518937 0.60276338]
18print(f"{alice.name}: {alice.vector}")
19# Out: Bob: [0.88325739 0.73527629 0.87746707]
20print(f"{bob.name}: {bob.vector}")
21
22# Define a jit function to operate on the structs.
23@njit
24def distance(a, b):
25    return np.linalg.norm(a.vector - b.vector)
26
27# Out: 0.4332647200356598
28print(distance(alice, bob))

Defining a method on StructRef

Methods and attributes can be attached using @overload_* as shown in the previous sections.

The following demonstrates the use of @overload_method to insert a method for instances of MyStructType:

from test_overload_method of numba/tests/doc_examples/test_structref_usage.py
 1from numba.core.extending import overload_method
 2from numba.core.errors import TypingError
 3
 4# Use @overload_method to add a method for
 5# MyStructType.distance(other)
 6# where *other* is an instance of MyStructType.
 7@overload_method(MyStructType, "distance")
 8def ol_distance(self, other):
 9    # Guard that *other* is an instance of MyStructType
10    if not isinstance(other, MyStructType):
11        raise TypingError(
12            f"*other* must be a {MyStructType}; got {other}"
13        )
14
15    def impl(self, other):
16        return np.linalg.norm(self.vector - other.vector)
17
18    return impl
19
20# Test
21@njit
22def test():
23    alice = MyStruct("Alice", vector=np.random.random(3))
24    bob = MyStruct("Bob", vector=np.random.random(3))
25    # Use the method
26    return alice.distance(bob)

numba.experimental.structref API Reference

Utilities for defining a mutable struct.

A mutable struct is passed by reference; hence, structref (a reference to a struct).

class numba.experimental.structref.StructRefProxy(*args)

A PyObject proxy to the Numba allocated structref data structure.

Notes

  • Subclasses should not define __init__.

  • Subclasses can override __new__.

numba.experimental.structref.define_attributes(struct_typeclass)

Define attributes on struct_typeclass.

Defines both setters and getters in jit-code.

This is called directly in register().

numba.experimental.structref.define_boxing(struct_type, obj_class)

Define the boxing & unboxing logic for struct_type to obj_class.

Defines both boxing and unboxing.

  • boxing turns an instance of struct_type into a PyObject of obj_class

  • unboxing turns an instance of obj_class into an instance of struct_type in jit-code.

Use this directly instead of define_proxy() when the user does not want any constructor to be defined.

numba.experimental.structref.define_constructor(py_class, struct_typeclass, fields)

Define the jit-code constructor for struct_typeclass using the Python type py_class and the required fields.

Use this instead of define_proxy() if the user does not want boxing logic defined.

numba.experimental.structref.define_proxy(py_class, struct_typeclass, fields)

Defines a PyObject proxy for a structref.

This makes py_class a valid constructor for creating a instance of struct_typeclass that contains the members as defined by fields.

Parameters
py_classtype

The Python class for constructing an instance of struct_typeclass.

struct_typeclassnumba.core.types.Type

The structref type class to bind to.

fieldsSequence[str]

A sequence of field names.

Returns
None
numba.experimental.structref.register(struct_type)

Register a numba.core.types.StructRef for use in jit-code.

This defines the data-model for lowering an instance of struct_type. This defines attributes accessor and mutator for an instance of struct_type.

Parameters
struct_typetype

A subclass of numba.core.types.StructRef.

Returns
struct_typetype

Returns the input argument so this can act like a decorator.

Examples

class MyStruct(numba.core.types.StructRef):
    ...  # the simplest subclass can be empty

numba.experimental.structref.register(MyStruct)

Determining if a function is already wrapped by a jit family decorator

The following function is provided for this purpose.

extending.is_jitted()

Returns True if a function is wrapped by one of the Numba @jit decorators, for example: numba.jit, numba.njit

The purpose of this function is to provide a means to check if a function is already JIT decorated.