def test_c_structure(self):
spec = OrderedDict()
spec['a'] = int32
spec['b'] = int16
spec['c'] = float64
@jitclass(spec)
class Struct(object):
def __init__(self, a, b, c):
self.a = a
self.b = b
self.c = c
st = Struct(0xabcd, 0xef, 3.1415)
class CStruct(ctypes.Structure):
_fields_ = [
('a', ctypes.c_int32),
('b', ctypes.c_int16),
('c', ctypes.c_double),
]
ptr = ctypes.c_void_p(st._dataptr)
cstruct = ctypes.cast(ptr, ctypes.POINTER(CStruct))[0]
self.assertEqual(cstruct.a, st.a)
self.assertEqual(cstruct.b, st.b)
self.assertEqual(cstruct.c, st.c)
def __init__(self, func, sig, identity=None, targetoptions={}):
assert not targetoptions
self.py_func = func
self.identity = self.parse_identity(identity)
self.signature = sig
self.inputsig, self.outputsig = parse_signature(self.signature)
assert len(self.outputsig) == 1, "only support 1 output"
# { arg_dtype: (return_dtype), cudakernel }
self.kernelmap = OrderedDict()
def _make_Vector2(self):
spec = OrderedDict()
spec['x'] = int32
spec['y'] = int32
@jitclass(spec)
class Vector2(object):
def __init__(self, x, y):
self.x = x
self.y = y
return Vector2
def test_jitclass_datalayout(self):
spec = OrderedDict()
# Boolean has different layout as value vs data
spec['val'] = boolean
@jitclass(spec)
class Foo(object):
def __init__(self, val):
self.val = val
self.assertTrue(Foo(True).val)
self.assertFalse(Foo(False).val)
def __init__(self, func, sig, identity=None, targetoptions={}):
# Allow nopython flag to be set.
if not targetoptions.pop('nopython', True):
raise TypeError("nopython flag must be True")
# Are there any more target options?
if targetoptions:
opts = ', '.join([repr(k) for k in targetoptions.keys()])
fmt = "The following target options are not supported: {0}"
raise TypeError(fmt.format(opts))
self.py_func = func
self.identity = parse_identity(identity)
self.signature = sig
self.inputsig, self.outputsig = parse_signature(self.signature)
assert len(self.outputsig) == 1, "only support 1 output"
# { arg_dtype: (return_dtype), cudakernel }
self.kernelmap = OrderedDict()
def _make_Float2AndArray(self):
spec = OrderedDict()
spec['x'] = float32
spec['y'] = float32
spec['arr'] = float32[:]
@jitclass(spec)
class Float2AndArray(object):
def __init__(self, x, y, arr):
self.x = x
self.y = y
self.arr = arr
def add(self, val):
self.x += val
self.y += val
return val
return Float2AndArray
def test_deferred_type(self):
node_type = deferred_type()
spec = OrderedDict()
spec['data'] = float32
spec['next'] = optional(node_type)
@njit
def get_data(node):
return node.data
@jitclass(spec)
class LinkedNode(object):
def __init__(self, data, next):
self.data = data
self.next = next
def get_next_data(self):
# use deferred type as argument
return get_data(self.next)
node_type.define(LinkedNode.class_type.instance_type)
first = LinkedNode(123, None)
self.assertEqual(first.data, 123)
self.assertIsNone(first.next)
second = LinkedNode(321, first)
self.assertEqual(first._meminfo.refcount, 2)
self.assertEqual(second.next.data, first.data)
self.assertEqual(first._meminfo.refcount, 2)
self.assertEqual(second._meminfo.refcount, 1)
# Test using deferred type as argument
first_val = second.get_next_data()
self.assertEqual(first_val, first.data)
# Check ownership
self.assertEqual(first._meminfo.refcount, 2)
del second
self.assertEqual(first._meminfo.refcount, 1)
def __init__(self, func, identity=None, targetoptions={}):
assert not targetoptions
self.py_func = func
self.identity = parse_identity(identity)
# { arg_dtype: (return_dtype), cudakernel }
self.kernelmap = OrderedDict()
def test_ordereddict_spec(self):
spec = OrderedDict()
spec['x'] = int32
spec['y'] = float32
self._check_spec(spec)
"""
This is a more complicated jitclasses example.
Here, we implement a binarytree and iterative preorder and inorder traversal
function using a handwritten stack.
"""
from __future__ import print_function, absolute_import
import random
from numba.utils import OrderedDict
from numba import njit
from numba import jitclass
from numba import int32, deferred_type, optional
from numba.runtime import rtsys
node_type = deferred_type()
spec = OrderedDict()
spec['data'] = int32
spec['left'] = optional(node_type)
spec['right'] = optional(node_type)
@jitclass(spec)
class TreeNode(object):
def __init__(self, data):
self.data = data
self.left = None
self.right = None
node_type.define(TreeNode.class_type.instance_type)
"""
This example demonstrates jitclasses and deferred type.
This is an extension to the simpler singly-linked-list example in
``linkedlist.py``.
Here, we make a better interface in the Stack class that encapsuate the
underlying linked-list.
"""
from __future__ import print_function, absolute_import
from numba.utils import OrderedDict
from numba import njit
from numba import jitclass
from numba import deferred_type, intp, optional
from numba.runtime import rtsys
linkednode_spec = OrderedDict()
linkednode_type = deferred_type()
linkednode_spec['data'] = data_type = deferred_type()
linkednode_spec['next'] = optional(linkednode_type)
@jitclass(linkednode_spec)
class LinkedNode(object):
def __init__(self, data):
self.data = data
self.next = None
linkednode_type.define(LinkedNode.class_type.instance_type)
stack_spec = OrderedDict()
def register_class_type(cls, spec, class_ctor, builder):
"""
Internal function to create a jitclass.
Args
----
cls: the original class object (used as the prototype)
spec: the structural specification contains the field types.
class_ctor: the numba type to represent the jitclass
builder: the internal jitclass builder
"""
# Normalize spec
if isinstance(spec, Sequence):
spec = OrderedDict(spec)
# Copy methods from base classes
clsdct = {}
for basecls in reversed(inspect.getmro(cls)):
clsdct.update(basecls.__dict__)
methods = dict((k, v) for k, v in clsdct.items()
if isinstance(v, pytypes.FunctionType))
props = dict((k, v) for k, v in clsdct.items() if isinstance(v, property))
others = dict((k, v) for k, v in clsdct.items()
if k not in methods and k not in props)
# Check for name shadowing
shadowed = (set(methods) | set(props)) & set(spec)
if shadowed:
raise NameError("name shadowing: {0}".format(', '.join(shadowed)))
docstring = others.pop('__doc__', "")
_drop_ignored_attrs(others)
if others:
msg = "class members are not yet supported: {0}"
members = ', '.join(others.keys())
raise TypeError(msg.format(members))
for k, v in props.items():
if v.fdel is not None:
raise TypeError("deleter is not supported: {0}".format(k))
jitmethods = {}
for k, v in methods.items():
jitmethods[k] = njit(v)
jitprops = {}
for k, v in props.items():
dct = {}
if v.fget:
dct['get'] = njit(v.fget)
if v.fset:
dct['set'] = njit(v.fset)
jitprops[k] = dct
# Instantiate class type
class_type = class_ctor(cls, ConstructorTemplate, spec, jitmethods,
jitprops)
cls = JitClassType(cls.__name__, (cls, ),
dict(class_type=class_type, __doc__=docstring))
# Register resolution of the class object
typingctx = CPUTarget.typing_context
typingctx.insert_global(cls, class_type)
# Register class
targetctx = CPUTarget.target_context
builder(class_type, methods, typingctx, targetctx).register()
return cls