python-c-api - 如何正确地将 CuPy 数组发送到 Python C 扩展模块

Question

我正在尝试为 Python 编写一个自定义 CXX 扩展，它将 CuPy 数组作为输入。从 CXX 扩展传递和返回 CuPy 数组的推荐方法是什么？

我目前的方法是将内存指针作为整数传递给扩展，并就地修改 CuPy 数组。但是，使用这种方法时，当数组被传递给两个扩展函数而不被重新转换为中间的 CuPy 数组时，会出现一些奇怪的行为。

我正在研究一个有效的示例，但与此同时，这是我的方法的示意图，我想要关于如何在 CuPy 和 CXX 之间交换数组的建议。

import cupy as cp

a = cp.zeros((1, 2, 3))

cxx_function_modifying_a_in_place(a.data.ptr)

a = cp.array(a)  # required else something goes wrong

cxx_function_modifying_a_in_place(a.data.ptr)

以下函数使用 pybind11 或 SWIG 包装

void cxx_function_modifying_array_in_place(size_t g_)
{
  // convert pointers to correct type
  g = (float2 *)g_;
  // do some stuff to g on a GPU
  cuda_kernel<<<1, 1>>>(g);
  cufftExecC2C(plan2d, (cufftComplex *)g, (cufftComplex *)g, CUFFT_FORWARD);
}