Compiling Python classes with @jitclass
¶
Note
This is a early version of jitclass support. Not all compiling features are exposed or implemented, yet.
Numba supports code generation for classes via the numba.jitclass()
decorator. A class can be marked for optimization using this decorator along
with a specification of the types of each field. We call the resulting class
object a jitclass. All methods of a jitclass are compiled into nopython
functions. The data of a jitclass instance is allocated on the heap as a
C-compatible structure so that any compiled functions can have direct access
to the underlying data, bypassing the interpreter.
Basic usage¶
Here’s an example of a jitclass:
import numpy as np
from numba import int32, float32 # import the types
from numba.experimental import jitclass
spec = [
('value', int32), # a simple scalar field
('array', float32[:]), # an array field
]
@jitclass(spec)
class Bag(object):
def __init__(self, value):
self.value = value
self.array = np.zeros(value, dtype=np.float32)
@property
def size(self):
return self.array.size
def increment(self, val):
for i in range(self.size):
self.array[i] += val
return self.array
@staticmethod
def add(x, y):
return x + y
n = 21
mybag = Bag(n)
In the above example, a spec
is provided as a list of 2-tuples. The tuples
contain the name of the field and the Numba type of the field. Alternatively,
user can use a dictionary (an OrderedDict
preferably for stable field
ordering), which maps field names to types.
The definition of the class requires at least a __init__
method for
initializing each defined fields. Uninitialized fields contains garbage data.
Methods and properties (getters and setters only) can be defined. They will be
automatically compiled.
Specifying numba.typed
containers as class members¶
It is often desirable to use a numba.typed.Dict
or a numba.typed.List
as
a class member in a jitclass
. Methods for using these types and various
common patterns are presented in the following:
First, using explicit Numba types and explicit construction.
from numba import jitclass, types, typed
# key and value types
kv_ty = (types.int64, types.unicode_type)
# A container class with:
# * member 'd' holding a typed dictionary of int64 -> unicode string (kv_ty)
# * member 'l' holding a typed list of float64
@jitclass([('d', types.DictType(*kv_ty)),
('l', types.ListType(types.float64))])
class ContainerHolder(object):
def __init__(self):
# initialize the containers
self.d = typed.Dict.empty(*kv_ty)
self.l = typed.List.empty_list(types.float64)
container = ContainerHolder()
container.d[1] = "apple"
container.d[2] = "orange"
container.l.append(123.)
container.l.append(456.)
print(container.d) # {1: apple, 2: orange}
print(container.l) # [123.0, 456.0]
Another useful pattern is to use the numba.typed
container attribute
_numba_type_
to find the type of a container, this can be accessed directly
from an instance of the container in the Python interpreter. The same
information can be obtained by calling numba.typeof()
on the instance. For
example:
from numba import jitclass, typed, typeof
d = typed.Dict()
d[1] = "apple"
d[2] = "orange"
l = typed.List()
l.append(123.)
l.append(456.)
@jitclass([('d', typeof(d)), ('l', typeof(l))])
class ContainerInstHolder(object):
def __init__(self, dict_inst, list_inst):
self.d = dict_inst
self.l = list_inst
container = ContainerInstHolder(d, l)
print(container.d) # {1: apple, 2: orange}
print(container.l) # [123.0, 456.0]
It is worth noting that the instance of the container in a jitclass
must be
initialized before use, for example, this will cause an invalid memory access
as self.d
is written to without d
being initialized as a type.Dict
instance of the type specified.
from numba import jitclass, types
dict_ty = types.DictType(types.int64, types.unicode_type)
@jitclass([('d', dict_ty)])
class NotInitilisingContainer(object):
def __init__(self):
self.d[10] = "apple" # this is invalid, `d` is not initialized
NotInitilisingContainer() # segmentation fault/memory access violation
Support operations¶
The following operations of jitclasses work in both the interpreter and Numba compiled functions:
- calling the jitclass class object to construct a new instance
(e.g.
mybag = Bag(123)
); - read/write access to attributes and properties (e.g.
mybag.value
); - calling methods (e.g.
mybag.increment(3)
); - calling static methods as instance attributes (e.g.
mybag.add(1, 1)
); - calling static methods as class attributes (e.g.
Bag.add(1, 2)
);
Using jitclasses in Numba compiled function is more efficient.
Short methods can be inlined (at the discretion of LLVM inliner).
Attributes access are simply reading from a C structure.
Using jitclasses from the interpreter has the same overhead of calling any
Numba compiled function from the interpreter. Arguments and return values
must be unboxed or boxed between Python objects and native representation.
Values encapsulated by a jitclass does not get boxed into Python object when
the jitclass instance is handed to the interpreter. It is during attribute
access to the field values that they are boxed.
Calling static methods as class attributes is only supported outside of the
class definition (i.e. you can’t call Bag.add()
from within another method
of Bag
).
Limitations¶
- A jitclass class object is treated as a function (the constructor) inside a Numba compiled function.
isinstance()
only works in the interpreter.- Manipulating jitclass instances in the interpreter is not optimized, yet.
- Support for jitclasses are available on CPU only. (Note: Support for GPU devices is planned for a future release.)
The decorator: @jitclass
¶
-
numba.experimental.
jitclass
(spec)¶ A decorator for creating a jitclass.
arguments:
- spec:
- Specifies the types of each field on this class. Must be a dictionary or a sequence. With a dictionary, use collections.OrderedDict for stable ordering. With a sequence, it must contain 2-tuples of (fieldname, fieldtype).
returns:
A callable that takes a class object, which will be compiled.