python - 如何在jitclass中使用prange

Question

如何在属于 jitclass 成员的函数中使用 prange？例如，我想使用 prange 在函数 do_smth 中并行执行。

@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

Answer 1

我还没有找到在 jitclass 成员中指定 parallel=True 的方法。但是您可以改为调用外部函数。如果确实需要并行化，这个小开销应该是最小的。

@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())

即使 prange 在 Test.do_smth() 中显式使用，循环也不会并行化。处理行的顺序很好地指示了循环何时并行化:

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