Могу ли я использовать общую память Numba для ускорения с 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]

Вот часть моего кода, v - это [9,2] массив чашки.

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] ])

Предыдущий - это определение постоянного массива 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

где fin - массив [9,420,180] cupy. Я хочу вычислить массив u, [2,420,180]. Обычное вычисление (без GPU):

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

Однако оно вызывает ошибку, как описано ранее. Не знаю почему. Типы данных v и sfin не совпадают? Должен ли я перейти на массив Numpy, чтобы использовать общую память в Numba?

Answer 1

Я столкнулся с той же проблемой, и я использовал массивы Numpy с функцией cuda.to_device () для передачи их на GPU. Я думаю, что на данный момент Cupy не совместим с массивами разделяемой памяти.