Is there is jitclass equivalent to generated_jit? I’d like Numba to compile a class with a given type before I instantiate it for the first time (alternative approaches certainly welcome).
A minimally reproducible example would be:
from functools import lru_cache
import numba as nb
class _DequeNode:
"""A wrapper around some item type.
This class will be jitclassed in a factory function (so we can define
multiple types to wrap in a DequeNode).
"""
def __init__(self, item, left=None, right=None):
self.item = item
self.left = left
self.right = right
@lru_cache(maxsize=32)
def _make_deque_node_class(item_type):
"""Returns a specialization of DequeNode for the given item type.
The return value from this function is cached so that we can access
the same type across multiple calls (plus we avoid compiling the same
thing over and over).
"""
node_type = nb.deferred_type()
spec = [
("item", item_type),
("left", nb.types.optional(node_type)),
("right", nb.types.optional(node_type)),
]
DequeNode = nb.experimental.jitclass(spec)(_DequeNode)
node_type.define(DequeNode.class_type.instance_type)
return DequeNode
@lru_cache(maxsize=32)
def make_deque(item_type):
DequeNode = _make_deque_node_class(item_type)
deque_spec = [
("first_node", nb.types.optional(DequeNode.class_type.instance_type)),
("last_node", nb.types.optional(DequeNode.class_type.instance_type)),
("len", nb.types.uint64),
]
@nb.experimental.jitclass(deque_spec)
class Deque:
def __init__(self):
self.first_node = None
self.last_node = None
self.len = 0
# ... rest of implementation of Deque, using the DequeNode above to
# wrap any items added to the container
return Deque
Running make_deque almost works:
>>> Deque = make_deque(nb.types.int64)
>>> d = Deque()
RuntimeError: Failed in nopythyon mode pipeline (step: nopython mode backend)
LLVM IR parsing error
<string>:35:122: error: base element of getelementptr must be sized
%".18" = getelementptr inbounds {...etc etc}
I don’t know anything really about LLVM IR parsing but “must be sized” sounds to me like LLVM has no idea how big the specialized DequeNode class will be while creating the Deque type. My suspicion is that Numba is only compiling the DequeNode class the first time it’s instantiated and (in make_deque) it has never been instantiated.
Sure enough, manually instantiating a DequeNode before running the code above works mint (since the results of _make_deque_node_class are cached, we access the same type as the Deque)
>>> DequeNode = _make_deque_node_class(nb.types.int64)
>>> DequeNode(1, None, None)
>>> Deque = make_deque(nb.types.int64)
>>> d = Deque()
( no errors, d works)
