我可以使用 Numba 共享的 memory 与 Cupy 一起加速吗？

Question

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
Invalid use of Function(<built-in function getitem>) with argument(s) of type(s): 
(array(float64, 2d, C), Tuple(int64, int32, int64))

 * parameterized

File "<ipython-input-34-637851842bfe>", line 34:

def macroscopic(fin,u,v):

<source elided>

  for k in range(TPB):

    tmp1 = v[i,0] * sfin[i,tx,k]

Here is part of my code, v is a [9,2] cupy array.

import cupy as cp
from numba import cuda
v = cp.array([ [ 1,  1], [ 1,  0], [ 1, -1], [ 0,  1], [ 0,  0],

           [ 0, -1], [-1,  1], [-1,  0], [-1, -1] ])

Previous is the definition of the constant array v

@cuda.jit
def macroscopic(fin,u,v):
  #Shared Memory
  # The computation will be done on blocks of TPBxTPB elements.
  TPB = 16
  sfin = cuda.shared.array(shape=(TPB, TPB), dtype=numba.float64)
  x,y = cuda.grid(2)
  tx = cuda.threadIdx.x
  ty = cuda.threadIdx.y
  bpg = cuda.gridDim.x    # blocks per grid
  if x >= fin.shape[1] and y >= fin.shape[2]:
    return

  tmp1 = 0.
  tmp2 = 0.

  for i in range(9):
    for j in range(bpg):
      #preload data to shared memory
      sfin[tx, ty] = fin[i,x, ty + j * TPB]
      cuda.syncthreads()
      #compute in the shared memory
      for k in range(TPB):
        tmp1 = v[i,0] * sfin[i,tx,k]
        tmp2 = v[i,1] * sfin[i,tx,k]

    cuda.syncthreads()
    u[0,x,y] += tmp1
    u[1,x,y] += tmp2

where, fin is a [9,420,180] cupy array. I want to compute the u, [2,420,180] array. The normal computation (non-GPU) is:

for i in range(9):
    u[0,:,:] += v[i,0] * fin[i,:,:]
    u[1,:,:] += v[i,1] * fin[i,:,:]

However, it throws an error as previously described. I don't know why. Is the data type of v and sfin not match? Shall I change to Numpy array to use the shared memory in Numba?

Answer 1

I run into the same problem, and I used Numpy arrays with cuda.to_device() function to transfer them to the GPU. I think at the moment Cupy is not compatible with shared memory arrays.