在 NUMBA 中使用多线程

Question

>我有一个在多边形测试中执行点的函数。它采用两个 2D numpy 数组作为输入(一系列点和一个多边形(。该函数返回一个布尔值作为输出(如果点位于多边形内，则为 True，否则为 False(。代码是从这个 SO 答案借来的。下面是一个示例：

from numba import jit
from numba.pycc import CC
cc = CC('nbspatial')
import numpy as np
@cc.export('array_tracing2', 'b1[:](f8[:,:], f8[:,:])')
@jit(nopython=True, nogil=True)
def array_tracing2(xy, poly):
D = np.empty(len(xy), dtype=numba.boolean)
n = len(poly)
for i in range(1, len(D) - 1):
inside = False
p2x = 0.0
p2y = 0.0
xints = 0.0
p1x,p1y = poly[0]
x = xy[i][0]
y = xy[i][1]
for i in range(n+1):
p2x,p2y = poly[i % n]
if y > min(p1y,p2y):
if y <= max(p1y,p2y):
if x <= max(p1x,p2x):
if p1y != p2y:
xints = (y-p1y)*(p2x-p1x)/(p2y-p1y)+p1x
if p1x == p2x or x <= xints:
inside = not inside
p1x,p1y = p2x,p2y
D[i] = inside
return D

if __name__ == "__main__":
cc.compile()

上面的代码可以通过运行python numba_src.py来编译，并使用以下命令进行测试：

import numpy as np
# regular polygon for testing
lenpoly = 10000
polygon = np.array([[np.sin(x)+0.5,np.cos(x)+0.5] for x in np.linspace(0,2*np.pi,lenpoly)[:-1]])
# random points set of points to test 
N = 100000
# making a list instead of a generator to help debug
pp = np.array([np.random.random(N), np.random.random(N)]).reshape(N,2)

import nbspatial
nbspatial.array_tracing2(pp, polygon) 

我的尝试是并行化上面的代码，以便利用所有可用的CPU。

我尝试按照numba官方文档中的示例使用@njit

import numba
@njit(parallel=True)
def array_tracing3(xy, poly):
D = np.empty(len(xy), dtype=numba.boolean)
n = len(poly)
for i in range(1, len(D) - 1):
inside = False
p2x = 0.0
p2y = 0.0
xints = 0.0
p1x,p1y = poly[0]
x = xy[i][0]
y = xy[i][1]
for i in range(n+1):
p2x,p2y = poly[i % n]
if y > min(p1y,p2y):
if y <= max(p1y,p2y):
if x <= max(p1x,p2x):
if p1y != p2y:
xints = (y-p1y)*(p2x-p1x)/(p2y-p1y)+p1x
if p1x == p2x or x <= xints:
inside = not inside
p1x,p1y = p2x,p2y
D[i] = inside
return D

上面的代码在预编译的串行版本的55''与1' 33''中完成了N=1000000。系统监视器显示只有一个 CPU 以 100% 的速度运行。

如何尝试使用可用的整个 CPU，并在单个布尔数组中返回结果？

Answer 1

Numba 的parallel=True只为某些函数启用自动并行性，而不是为所有循环启用自动并行性。 应将其中一个range()表达式替换为numba.prange以启用多核计算。

请参阅：https://numba.pydata.org/numba-doc/dev/user/parallel.html