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

Posted 2023-02-16

【中文标题】使用 ctypes 将 cupy 指针传递给 CUDA 内核

【英文标题】：Passing a cupy pointer to a CUDA kernel using ctypes

【发布时间】：2020-08-18 18:12:12

【问题描述】：

我有一个 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 内核时犯了一个错误。

【问题讨论】：

我认为可能的原因之一是fromstarts和fromstopscupy数组的dtype是int64。但是您正在创建 int32 指针以传递给库。您可以使用cupy.array([..]), dtype=cupy.int32)

创建数组

@emcastillo 我试过了（并用 dtypes 更新了问题）但我在tonum 中得到了相同的错误值。

【参考方案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 来验证这一点。