How to use prange in jitclass

user_7713315 0
Viewed 245

How to use prange in function that is member of jitclass? For example I want to parallel for in function do_smth using prange.

@jitclass([('x', float64[:])])
class Test:
    def __init__(self):
        self.x = np.arange(10)

    def do_smth(self):
        for i in range(len(self.x)):
            pass
numba python
1 Answers

I haven't found a way to specify parallel=True in a jitclass member. But you can call an external function instead. This small overhead should be minimal if parallelization is really required.

@nb.experimental.jitclass([('x', nb.float64[:, :])])
class Test:

    def __init__(self):
        self.x = np.arange(10000.).reshape((10,1000))

    def do_smth(self):
        rows, cols = self.x.shape
        sums = np.zeros(cols)
        for r in nb.prange(rows):
            print(r)
            sums[r] = np.sum(self.x[r])
        return np.sum(sums)

    def do_other(self):
        return do_other(self.x)

@nb.njit([nb.float64(nb.float64[:,:])], parallel=True)
def do_other(x):
    rows, cols = x.shape
    sums = np.zeros(cols)
    for r in nb.prange(rows):
        print(r)
        sums[r] = np.sum(x[r])
    return np.sum(sums)

t = Test()
print("Class member")
print(t.do_smth())
print("External function")
print(t.do_other())

Even though prange is used explicitly in Test.do_smth(), the loop is not parallelized. The order in which rows are processed is a good indication of when the loop is parallelized:

Class member
0
1
2
3
4
5
6
7
8
9
49995000.0
External function
0
2
6
3
5
9
8
7
4
1
49995000.0
user_21484160
Related