Python/Numba - Custom class object as input type

Question

I'm starting with numba and my first goal is to try and accelerate a not so complicated function with a nested loop.

Given the following class:

class TestA:
    def __init__(self, a, b):
        self.a = a
        self.b = b

    def get_mult(self):
        return self.a * self.b

and a numpy ndarray that contains class TestA objects. Dimension (N,) where N is usually ~3 million in length.

Now given the following function:

def test_no_jit(custom_class_obj_container):
    container_length = len(custom_class_obj_container)
    sum = 0
    for i in range(container_length):
        for j in range(i + 1, container_length):
            obj_i = custom_class_obj_container[i]
            obj_j = custom_class_obj_container[j]
            sum += (obj_i.get_mult() + obj_j.get_mult())

    return sum

I've tried to play around numba to get it to work with the function above however I cannot seem to get it to work with nopython=True flag, and if it's set to false, then the runtime is higher than the no-jit function.

Here is my latest try in trying to jit the function (also using nb.prange):

@nb.jit(nopython=False, parallel=True)
def test_jit(custom_class_obj_container):
    container_length = len(custom_class_obj_container)
    sum = 0
    for i in nb.prange(container_length):
        for j in nb.prange(i + 1, container_length):
            obj_i = custom_class_obj_container[i]
            obj_j = custom_class_obj_container[j]
            sum += (obj_i.get_mult() + obj_j.get_mult())

    return sum

I've tried to search around but I cannot seem to find a tutorial of how to define a custom class in the signature, and how would I go in order to accelerate a function of that sort and get it to run on GPU and possibly (any info regarding that matter would be highly appreciated) to get it to run with cuda libraries - which are installed and ready to use (previously used with tensorflow)

Answer 1

The numba docs give an example of creating a custom type, even for nopython mode: https://numba.pydata.org/numba-doc/latest/extending/interval-example.html

In your case though, unless this is a really slimmed down version of what you actually want to do, it seems like the easiest approach would be to re-use existing types. Additionally, the construction of a 3M length object array is going to be slow, and produce fragmented memory (as the objects are not being stored in contiguous blocks).

An example of how using record arrays might be used to solve the problem:

x_dt = np.dtype([('a', np.float64),
                 ('b', np.float64)])
n = 30000
buf = np.arange(n*2).reshape((n, 2)).astype(np.float64)
vec3 = np.recarray(n, dtype=x_dt, buf=buf) 

@numba.njit
def mult(a):
    return a.a * a.b

@numba.jit(nopython=True, parallel=True)
def sum_of_prod(vector):
    sum = 0
    vector_len = len(vector)
    for i in numba.prange(vector_len):
        for j in numba.prange(i + 1, vector_len):
            sum += mult(vector[i]) + mult(vector[j])
    return sum

sum_of_prod(vec3)

FWIW, I'm no numba expert. I found this question when searching for how to implement a custom type in numba for non-numerical stuff. In your case, because this is highly numerical, I think a custom type is probably overkill.