Any Customizable Container that Doesn't Box/Unbox when passed around?

Numba Community Support
1

I’ve noticed that how things get passed between numba and python can significantly effect the runtime. For example consider the following three ways of passing parameters into an njit function:

from numba import njit, f8
from numba.types import unicode_type, Tuple
from numba.typed import List
import timeit
N=1000
def time_ms(f):
	f() #warm start
	return " %0.6f ms" % (1000.0*(timeit.timeit(f, number=N)/float(N)))

@njit
def separate(f1,f2,f3,f4, s1,s2,s3,s4):
	return f1 + f2 + f4 + f4 + len(s1) + len(s2) + len(s3) + len(s4)

@njit
def together(t):
	f1,f2,f3,f4,s1,s2,s3,s4 = t
	return f1 + f2 + f4 + f4 + len(s1) + len(s2) + len(s3) + len(s4)

@njit
def from_list(l):
	f1,f2,f3,f4,s1,s2,s3,s4 = l[0]
	return f1 + f2 + f4 + f4 + len(s1) + len(s2) + len(s3) + len(s4)

tup = (1,2,3,4,"A","BB", "CC", "DDD")

def s():
	separate(1,2,3,4,"A","BB","CC","DDD")

def t():
	together(tup)

lst = List.empty_list(Tuple((f8,f8,f8,f8,unicode_type,unicode_type,unicode_type,unicode_type)))
lst.append(tup)

def l():
	from_list(lst)

print(time_ms(s)) # 0.006556 ms
print(time_ms(t)) # 0.040448 ms
print(time_ms(l)) # 0.001097 ms

Somewhat counter-intuitively wrapping the arguments in a tuple in a typed list works best. I assume this is because the actual guts of List() doesn’t get boxed/unboxed when you just pass it between numba and python. My use case involves moving a lot of different Dicts between python and numba. Most of the time I don’t need to read into these Dicts on the python side, but I do need to pass them around. It would be nice to have some kind of custom context object that was just like a struct or tuple that I could pass into numba functions and change as needed, but that doesn’t unbox into a native type.

My question is is there any straightforward way to make this sort of custom type that I can use to just pass around my dictionaries and stuff without any substantial boxing/unboxing. Something that can be AOT-compiled is preferable. Haven’t had much luck with jitclasses and structrefs, they don’t seem to vibe well with AOT compilation.

2

Seems my aversion to using structrefs for this purpose came about because of a bug possibly on numba’s end.

import numba as nb
from numba import njit, types, typed, float64, int64, f8
from numba.core import types
from numba.experimental import structref
from numba.core.extending import overload_method
from numba.pycc import CC


####   STRUCT STUFF   #####
@structref.register
class VectorStructType(types.StructRef):
    def preprocess_fields(self, fields):
        return tuple((name, types.unliteral(typ)) for name, typ in fields)

class VectorStruct(structref.StructRefProxy):
    def __new__(cls, x, y, z):
        return structref.StructRefProxy.__new__(cls, x, y, z)

    @property
    def x(self):
        return VectorStruct_get_x(self)

    @property
    def y(self):
        return VectorStruct_get_y(self)

    @property
    def z(self):
        return VectorStruct_get_z(self)

@njit
def VectorStruct_get_x(self):
    return self.x

@njit
def VectorStruct_get_y(self):
    return self.y

@njit
def VectorStruct_get_z(self):
    return self.z

structref.define_proxy(VectorStruct, VectorStructType, ["x", "y", "z"])
vector_struct_type = VectorStructType(fields=(('x', f8), ('y', f8), ('z', f8)))

####   FUNCTIONS    #####
cc = CC('my_module')
@cc.export('add',f8(vector_struct_type,))
@njit
def add(v):
   return v.x + v.y + v.z

cc.compile()

####   NJIT'ed    #####
v_f = VectorStruct(1.,2.,3.)
print(add(v_f))
v_i = VectorStruct(1,2,3)
print(add(v_i))

####   AOT'ed    #####
from my_module import add

v_f = VectorStruct(1.,2.,3.)
print(add(v_f))
v_i = VectorStruct(1,2,3)
print(add(v_i))

####   Sanity Check    #####
cc = CC('my_module2')
@cc.export('add2',f8(f8,f8,f8))
@njit
def add2(x,y,z):
   return x + y + z

cc.compile()

print(add2(1,2,3))

from my_module2 import add2

print(add2(1,2,3))

produces:

6.0
6
6.0
3e-323
6
6.0

In other words the function using an AOT compiled structref wasn’t properly casting the int arguments to floats. I assume this is just a minor bug with numba? Or maybe not because the AOT compiled code has a strict signature and shouldn’t bother casting… I’m not sure.

Anyway structrefs seem to be doing the trick so maybe I answered my own question :man_shrugging:

1 Like