使用 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 来验证这一点。

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

Passing a cupy pointer to a CUDA kernel using ctypes

python

ctypes

cuda

cupy