1

I'm trying to write a custom CXX extension for Python which takes CuPy arrays as input. What's the recommended way to hand off and return CuPy arrays from a CXX extension?

My current approach is to pass a memory pointer to the extension as an integer and modify the CuPy array in-place. However, with this approach there is some strange behavior when the array is passed to two extension functions without being re-cast as a CuPy array in-between.

I'm working on a functioning example, but in the mean-time, here is a schematic of my approach, and I'd like suggestions on how to do exchange arrays between CuPy and 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)

The following function is wrapped using pybind11 or 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);
}
Daniel Ching
  • 443
  • 5
  • 9
  • Isn't `a.data.ptr` a pointer to some memory on the CUDA device, and so probably can't be accessed as a simple `float*` in C? – DavidW Nov 09 '19 at 10:59
  • Yes, it is. I'm using CUDA on the CXX side, so all operations to this array are being done on a GPU. – Daniel Ching Nov 09 '19 at 22:35
  • If you want to create an CuPy ndarray from raw CUDA pointer, see here: https://groups.google.com/forum/#!topic/cupy/3yzb9aIbOWI . However be aware that in most cases just passing a CUDA pointer from Python to C++ is insufficient because the array is not always C-contiguous. – kmaehashi Nov 11 '19 at 03:59
  • That's a good thought. I should check that the array is contiguous when it is returned. I'm always creating the array on the Python side, but I am using the CuFFT library, which is a black box. – Daniel Ching Nov 11 '19 at 16:06
  • In my particular case, there is something unexpected happening with the in-place cuFFT that is somehow fixed(?) by asking CuPy to make a copy of the array. I have switched to using the out-of-place cuFFT operator by creating an intermediate array, and this fixes the unexpected behavior I was observing without me needing to pepper my Python code with calls to `cp.array`. – Daniel Ching Nov 11 '19 at 19:46

0 Answers0