How do I nest dict-of-jitclass

I don’t think this is a dupe because it’s a different error message. I’m trying to do something that’s conceptually similar, though.
If I uncomment the last line below I get an error like “_pickle.PicklingError: Can’t pickle <class ‘numba.experimental.jitclass.base.Level1’>: it’s not found as numba.experimental.jitclass.base.Level1”

from numba import types, typed
from numba.experimental import jitclass

@jitclass([
    ("x", types.int64),
])
class Level1:
    def __init__(self):
        self.x = 12

level1 = Level1()  # Okay, I can create a Level1 instance
level1_instance_type = Level1.class_type.instance_type
level1_kv = (types.int64, level1_instance_type)

@jitclass([
    ("level1s", types.DictType(*level1_kv)),
])
class Level2(object):
    def __init__(self):
        self.level1s = typed.Dict.empty(*level1_kv)
        self.level1s[0] = Level1()


level2 = Level2()  # Okay, I can create a Level2 instance
level2_instance_type = Level2.class_type.instance_type
level2_kv = (types.int64, level2_instance_type)

@jitclass([
    ("level2s", types.DictType(*level2_kv)),
])
class Level3(object):
    def __init__(self):
        self.level2s = typed.Dict.empty(*level2_kv)

# _pickle.PicklingError: Can't pickle <class 'numba.experimental.jitclass.base.Level1'>: it's not found as numba.experimental.jitclass.base.Level1
# level3 = Level3()

Hi @nelson2005

I suspect that this is a bug, Numba’s trying to serialize something for the purposes of box/unboxing the data from the CPython interpreter and it’s got stuck trying to pickle a decorated class (I don’t recall all the details of how jitclass does this from memory). I suggest opening a bug report as it probably ought to work.

This similar code does work:

from numba import types, typed, njit
from numba.experimental import jitclass

@jitclass([
    ("x", types.int64),
])
class Level1:
    def __init__(self):
        self.x = 12

level1 = Level1()  # Okay, I can create a Level1 instance
level1_instance_type = Level1.class_type.instance_type
level1_kv = (types.int64, level1_instance_type)

@jitclass([
    ("level1s", types.DictType(*level1_kv)),
])
class Level2(object):
    def __init__(self):
        self.level1s = typed.Dict.empty(*level1_kv)
        self.level1s[0] = Level1()


level2 = Level2()  # Okay, I can create a Level2 instance
level2_instance_type = Level2.class_type.instance_type
level2_kv = (types.int64, level2_instance_type)

@jitclass([
    ("level2s", types.DictType(*level2_kv)),
])
class Level3(object):
    def __init__(self):
        self.level2s = typed.Dict.empty(*level2_kv)


@njit
def foo():
    Level3() # instantiate in `njit` function

foo()

because the boxing/unboxing part is avoided by virtue of instantiating the Level3 in a jit context, however, if you try and return it, it’ll hit the same pickle issue.

Many Thanks.

That’s outstanding! I can effectively hold it in non-jit space by keeping it in a container, like below. I think this effectively works around the issue… just make sure the object never sees regular-python space. In my particular use case, all I need is a pointer to the jitclass which will be accessed by cfunc callback functions.

@njit
def add_one(lis):
    lis.append(Level3())

level3_instance_type = Level3.class_type.instance_type
lis = typed.List.empty_list(level3_instance_type)
add_one(lis)
print(len(lis))

Bug report submitted
Any suggestion of the best way to get the pointer of the instance from inside the jit func?

Many thanks!

You’d need to use an @intrinsic, this is starting to sound a bit “scary”, what are you trying to do? :slight_smile:

I’ll confess to ‘scary’. :slight_smile: I’m trying to do a sqlite virtual table entirely in numba, currently using this inside the jit-func in order to get the address. It’s effectively a re-ask of a gitter thread (thanks @stuartarchibald for the tip on how to get the link)(hover over the post time to show it)

It seems be going fairly well… this question is also in support of my quest.

The devil is always in the details, but this feels like a perfect application for numba- native performance virtual tables presenting numpy structures to the sql interface. If it works. :slight_smile:

Hey @nelson2005 following up on the post you linked here’s what all my ‘scary’ intrinsics ended up looking like:

code: numbert/utils.py at experimental · DannyWeitekamp/numbert · GitHub
tests: numbert/test_utils.py at experimental · DannyWeitekamp/numbert · GitHub

Be sure to incref and decref stuff when appropriate and test lots. Although using intrinsics like this has the advantage of allowing you to dereference ‘pointers’, it has a tendency to introduce segfaults and memory leaks if you’re not careful. Personally although I’ve come to rely on this for my own application and have been able to make things work robustly I’m a little bit afraid of it because numba or llvmlite could change something with how they handle increfs/decrefs and then all my code could break. Livin’ dangerously I guess.

Awesome, I’ve been wondering where that ended up. I’ve been using your kit like this:

from llvmlite import ir
from numba import njit, types, i8
from numba.core import cgutils
from numba.experimental import jitclass
from numba.experimental.jitclass import _box
from numba.experimental.structref import _Utils
from numba.extending import intrinsic

int32_t = ir.IntType(32)

def fflush(builder):
    """
    Calls fflush(NULL) which flushes all open streams.
    """
    int8_t = ir.IntType(8)
    fflush_fnty = ir.FunctionType(int32_t, [int8_t.as_pointer()])
    fflush_fn = builder.module.get_or_insert_function(
        fflush_fnty, name='fflush')

    builder.call(fflush_fn, [int8_t.as_pointer()(None)])

# https://github.com/numba/numba/issues/5679
@intrinsic
def printf(typingctx, format_type, *args):
    """printf that can be called from Numba jit-decorated functions.
    """
    if isinstance(format_type, types.StringLiteral):
        sig = types.void(format_type, types.BaseTuple.from_types(args))

        def codegen(context, builder, signature, args):
            cgutils.printf(builder, format_type.literal_value, *args[1:])
            fflush(builder)

        return sig, codegen

@intrinsic
def _struct_get_data_pointer(typingctx, inst_type):
    def codegen(context, builder, sig, args):
        val_ty, = sig.args
        val, = args
        utils = _Utils(context, builder, val_ty)
        dataptr = utils.get_data_pointer(val)
        ret = builder.ptrtoint(dataptr, cgutils.intp_t)
        return ret

    sig = i8(inst_type, )
    return sig, codegen

@jitclass([])
class Demo:
    def __init__(self):
        pass
    def get_addr(self):
        return _struct_get_data_pointer(self)

@njit
def ptr_info():
    demo = Demo()
    printf("demo.get_addr=%p\n", demo.get_addr())
    printf("_struct_get_data_pointer=%p\n", _struct_get_data_pointer(demo))
    return demo

demo = ptr_info()
print(f"{_box.box_get_dataptr(demo):X}")

outputs

demo.get_addr=000001DCF2157D08
_struct_get_data_pointer=000001DCF2157D08
_box.box_get_dataptr=1DCF2157D08

This seems to work fine but I was poking around to see if there are any alternatives. :slight_smile:
Anyway, I appreciate you sharing- that’s great stuff!

Is it possible to get a pointer to a member, suitable for passing to a native library? For example, if there was a native function taking a pointer-to-int64 like

void func(long){}

I think @DannyWeitekamp’s _struct_get_attr_offset will work after jitclass is migrated to StructRef (which I wasn’t aware of until @luk-f-a closed 5111 with such a comment)

Is there a way that works with the current (0.52.0) jitclass? Or, should _struct_get_attr_offset work with with the current jitclass?

Not entirely sure how to make the equivalent function with jitclasses. I haven’t had a lot of luck with jitclasses in general and have been relying mostly on structrefs. There is a helper functions in the linked repo for auto-generating the source code for structrefs from a field list if that helps:

Thanks for sharing that- are structrefs laid out in memory the way a C library would expect? That is, if I get a pointer to the first attribute and pass it to a native library, will the native library be able to access all the attributes in normal “C” fashion?

Core devs will have a better answer. So take this with a grain of salt.
As I understand it, and I really don’t understand it all that well, the numba runtime has it’s own managed memory refcounting (dunno if there’s a delayed garbage collector or it just frees on zero), so structrefs and most other object-like things have a meminfo object with a recount and a pointer to the actual underlying data.
From experimenting with finding the address offsets of structref members, it seems that everything in the underlying data is laid out contiguously (with the first member having offset 0) and with the bit-widths you would expect (at least for numerical types). So in that sense I guess things are laid out in the “C” fashion.
However you should be aware that beyond numerical types things might be a bit different than you are expecting. For example unicode_type (which I assume is pretty important for your use case) looks quite a bit different than a C++ std::string so you might need to poke around to figure out how to reformat that if that’s what you’re expecting.

That’s very helpful, thanks for sharing your observations. I’d done some thinking about the nrt based on the docs but your observations are helpful to me- in particular, I had a mental model of the data pointer as being a parent of the meminfo block but your explanation helped me realize that the reverse is true.

For what it’s worth, by my reading the memory is freed on refcount=0