非常慢的Python麻婆鱼

Question

我在numba cuda中运行这个简单的代码，它发现非常慢。 知道瓶颈吗？在

   @cuda.jit('int32(float64,  float64, int32)', device=True)
   def mandelbrot_numbagpu(creal,cimag,maxiter):
       real = creal
       imag = cimag
       for n in range(maxiter):
           real2 = real*real
           imag2 = imag*imag
           if real2 + imag2 > 4.0:
               return n
           imag = 2* real*imag + cimag
           real = real2 - imag2 + creal
       return 0



   @cuda.jit
   def        mandelbrot_set_numbagpu(xmin,xmax,ymin,ymax,width,height,maxiter,n3,r1,r2):
      for i in range(width):
         for j in range(height):
            n3[i,j] = mandelbrot_numbagpu(r1[i],r2[j],maxiter)


   r1 = np.linspace(-2.0,0.5,1000, dtype=np.float )
   r2 = np.linspace(-1.25,1.25,1000, dtype=np.float)
   n3 = np.zeros((1000,1000),  dtype=np.uint8)

   %timeit mandelbrot_set_numbagpu(-2.0,0.5,-1.25,1.25,1000,1000,80,n3,r1,r2)
   #1 loops, best of 3: 4.84 s per loop

如果我运行JIT，它会快10倍！。。。。在

Answer 1

在Numba/CUDA（我认为是CUDA）中，函数不应该迭代数组。相反，它们应该处理一个数组元素，Numbda/CUDA处理程序将一大堆数组元素分配给一大堆GPU核心，因此所有事情都会快速并行地发生。这是all documented。在

不幸的是，这意味着您不能仅仅将@jit更改为{}，而是必须对其进行调整。在

以下工作：

# mandelbrot_numbagpu as before...

# I've removed some of the useless arguments for simplicity
@cuda.jit
def mandelbrot_set_numbagpu(n3,r1,r2,maxiter):
    # numba provides this function for working out which element you're
    # supposed to be accessing
    i,j = cuda.grid(2)
    if i<n3.shape[0] and j<n3.shape[1]: # check we're in range
        # do work on a single element
        n3[i,j] = mandelbrot_numbagpu(r1[i],r2[j],maxiter)

你就把它叫做

^{pr2}$

在我的电脑上，速度提高了3800倍。我不知道这和等效的CPU程序相比如何。在