Как правильно отправить массив CuPy в модуль расширения C Python - PullRequest
Новая
       55

Как правильно отправить массив CuPy в модуль расширения C Python

0 голосов
/ 08 ноября 2019

Я пытаюсь написать собственное расширение CXX для Python, которое принимает в качестве входных данных массивы CuPy. Каков рекомендуемый способ передачи и возврата массивов CuPy из расширения CXX?

Мой текущий подход - передать указатель памяти на расширение как целое число и изменить массив 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);
}
...