How to pass jit_options to jitclasses?

Is this feature supported? Any workaround if not…

More specifically, I am interested in cache=True option to be passed as jitclass flag

To my knowledge most compilations flags are not supported by jitclass yet (unless that has changed recently).

I had a similar issue witch caching, and in my case the workaround was relatively easy as the jitclass instances I had were actually constant. So I just turned it into a named tuple and made the methods normal jit-functions which can be cached (if you assure that the type of the named tuple is deterministic).

Maybe you can use the new mutuable StructRef that was released yesterday in 0.51, if you require a mutable data container

@Hannes I can’t use StructRef until https://github.com/numba/numba/issues/6077 is fixed. My use case is a little bit more complex. I need both optionals and deferred types to continue with structs. I was able to compile with jitclasses and is working great so far but the compilation is taking way to long everytime you start the app

To my knowledge most compilations flags are not supported by jitclass yet (unless that has changed recently).

This is correct as of Numba 0.51. jitclass does not support compilation flags directly.

@stuartarchibald
This might be the workaround I was looking for:

from numba import typed, typeof, njit, int64, types
from numba.experimental import jitclass
from numba.extending import overload_method, overload

spec = [
    ('value', int64),
]


@jitclass(spec)
class Bag(object):
    def __init__(self, value):
        self.value = value

mybag = Bag(21)

x = typed.List.empty_list(Bag.class_type.instance_type)
x.append(mybag)

@overload_method(types.misc.ClassInstanceType, 'clone', jit_options={'cache': True, 'parallel': True, 'nogil': True, 'boundscheck': False, 'inline': 'always'})
def ol_bag_clone(inst,):
    if inst is Bag.class_type.instance_type:
        def impl(inst,):
            return Bag(inst.value)
        return impl

@njit
def foo(z):
    t = z[0]
    c = t.clone()
    c.value = -21
    z.append(c)
    return z

l = foo(x)

@stuartarchibald, under such an approach, is it possible to separate the implementations for two different jitclasses?, ie

@overload_method(types.misc.ClassInstanceType, 'clone')
def ol_bag_clone(inst,):
    if inst is Bag.class_type.instance_type:
        def impl(inst,):
            return Bag(inst.value)
        return impl

@overload_method(types.misc.ClassInstanceType, 'clone')
def ol_bag_clone2(inst,):
    if inst is Bag2.class_type.instance_type:
        def impl(inst,):
            return 2
        return impl

I’m getting

No implementation of function Function(<function ol_bag_clone at 0x7f13d082d790>) found for signature:

 >>> ol_bag_clone(instance.jitclass.Bag2#7f13d055c640<value:int64>)

which suggests that the second implementation is being ignored.

@luk-f-a you should try this because you are overriding the same type

@overload_method(types.misc.ClassInstanceType, 'clone')
def ol_bag_clone(inst,):
    if inst is Bag.class_type.instance_type:
        def impl(inst,):
            return Bag(inst.value)
        return impl

    if inst is Bag2.class_type.instance_type:
        def impl(inst,):
            return 2
        return impl

Neat! That may come in handy at some point :slight_smile:

I’ve not thought about all the ins-and-outs of this yet, it’s another one for the “workarounds that may be ok” pile.

xref: source of this pattern so I don’t forget :)!

Have you been able to confirm if the caching part is working?

I never really understood the philosophy of jitclass types and instance types which contain some (random?) numbers like in luk’s post

instance.jitclass.Bag2#7f13d055c640<value:int64>

Opening two interpreter sessions easily shows that these numbers are not constant between sessions, so I wonder if the caching machinery can handle this situation, or is this id simply ignored by numba’s type resolution?

@Hannes I don’t think is working… also deferred types suffer from the same problem from what I am seeing

@stuartarchibald not sure how stable this solution is, but it would be good if it is documented… also it would be great if the current/old jitclass implementation would be replaced with structrefs and overloaded methods.

yes, that would work, but I’m creating jitclasses dynamically (via factory function), so it’s impossible to apply that pattern.

@luk-f-a have you tried to dynamically create the function you want to overload, using python’s eval/exec/compile built-in functionality?

The function is already dynamically created, the problem is how to assign it to be the overload of a method of a jitclass (or in my case, a named tuple).

@luk-f-a Is this doing what you want?

dyn_func_str = """
def ol_bag_clone(inst,):
    if inst is Bag.class_type.instance_type:
        def impl(inst,):
            return Bag(-inst.value)
        return impl

    if inst is Bag2.class_type.instance_type:
        def impl(inst,):
            return 2.5
        return impl
"""

exec(dyn_func_str)
dyn_func_call = locals()['ol_bag_clone']
overload_method(types.misc.ClassInstanceType, 'clone')(dyn_func_call)

unfortunately no. Bag and Bag2 never exist together in the same namespace. They are created one by one as needed (think of a jitclass as a generic, Bag, Bag), and after they are created it’s necessary to overload a method one them. There’s never the opportunity to create only one method, they need to overload sequentially.

@luk-f-a you can use global scope and I think you can overload dynamically the same method as many times as you want. only the last one will be use

exec(dyn_func_str, globals())
dyn_func_call = globals()['ol_bag_clone']
overload_method(types.misc.ClassInstanceType, 'clone')(dyn_func_call)
# second overload later on..
overload_method(types.misc.ClassInstanceType, 'clone')(new_ol_bag_clone)