使用ctypes将cupy指针传递给CUDA内核

Question

我有一个CUDA内核-

template <typename T, typename C>
__global__
void cuda_ListArray_num(
  C *tonum,
  const T *fromstarts,
  const T *fromstops
) {
  int64_t block_id = blockIdx.x + blockIdx.y * gridDim.x + gridDim.x * gridDim.y * blockIdx.z;
  int64_t thread_id = block_id * blockDim.x + threadIdx.x;
  int64_t start = fromstarts[thread_id];
  int64_t stop = fromstops[thread_id];
  tonum[thread_id] = (C) (stop - start);
}

ERROR
awkward_ListArray32_num_64(
  int64_t* tonum,
  const int32_t* fromstarts,
  const int32_t* fromstops,
  int64_t length) {

  dim3 blocks_per_grid;
  dim3 threads_per_block;

  if (length > 1024) {
    blocks_per_grid = dim3(ceil((length) / 1024.0), 1, 1);
    threads_per_block = dim3(1024, 1, 1);
  } else {
    blocks_per_grid = dim3(1, 1, 1);
    threads_per_block = dim3(length, 1, 1);
  }

  cuda_ListArray_num<int32_t, int64_t><<<blocks_per_grid, threads_per_block>>>(
    tonum,
    fromstarts,
    fromstops);

  cudaDeviceSynchronize();

  return success();
}

我可以将它添加到.so文件中，并使用ctypes从Python加载它。之后，我尝试从Python中使用它

这是在上述代码块中返回的ERROR结构的Python等价物-

class Error(ctypes.Structure):
    _fields_ = [
        ("str", ctypes.POINTER(ctypes.c_char)),
        ("identity", ctypes.c_int64),
        ("attempt", ctypes.c_int64),
        ("pass_through", ctypes.c_bool),
    ]

下面是我如何从Python中使用它-

lib = ctypes.CDLL("cuda-kernels.so")

funcC = getattr(lib, 'awkward_ListArray32_num_64')
funcC.restype = Error

tonum = cupy.array([123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123, 123], dtype=cupy.in64)
tonumx = ctypes.cast(tonum.data.ptr, ctypes.POINTER(ctypes.c_int64))
fromstarts = cupy.array([2, 0, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1], dtype=cupy.int32)
fromstarts = ctypes.cast(fromstarts.data.ptr, ctypes.POINTER(ctypes.c_int32))
fromstops = cupy.array([3, 2, 4, 5, 3, 4, 2, 5, 3, 4, 6, 11], dtype=cupy.int32)
fromstops = ctypes.cast(fromstops.data.ptr, ctypes.POINTER(ctypes.c_int32))
length = 3
funcC.argtypes = (ctypes.POINTER(ctypes.c_int64), ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_int32), ctypes.c_int64)
ret_pass = funcC(tonumx, fromstarts, fromstops, length)

但是当我打印tonum

>>> tonum[:3]
array([0, 0, 0])

但是值应该是-[1, 2, 2]（基于cuda_ListArray_num的工作方式）

我可能做错了什么？我想我可能在如何将cupy指针传递到cuda内核上犯了一个错误

Answer 1

您必须将python代码更改为

fromstarts = cupy.array([2, 0, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1], dtype=cupy.int32)
fromstarts_ctypes = ctypes.cast(fromstarts.data.ptr, ctypes.POINTER(ctypes.c_int32))
fromstops = cupy.array([4, 2, 4, 5, 3, 4, 2, 5, 3, 4, 6, 11], dtype=cupy.int32)
fromstops_ctypes = ctypes.cast(fromstops.data.ptr, ctypes.POINTER(ctypes.c_int32))
length = 3
funcC.argtypes = (ctypes.POINTER(ctypes.c_int64), ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_int32), ctypes.c_int64)
ret_pass = funcC(tonumx, fromstarts_ctypes, fromstops_ctypes, length)

原因是CuPy数组是用RAII管理的，因此当您将fromstarts变量重新分配给另一个对象（ctypes指针）时，实际数组将被销毁，其内存块将返回到CuPy的内存池。在此之后，当您创建fromstops数组时，它将使用相同的内存块，覆盖fromstarts数组的内容，因为该数组不再活动，并且共享同一个指针

然后，当您调用c代码时，fromstarts和fromstops实际上是同一个指针。您可以使用调试器或printf验证这一点