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
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:
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
:
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.