Хороший способ раскрыть CUPY MemoryPointer в C / C ++? - PullRequest
Новая
       77

Хороший способ раскрыть CUPY MemoryPointer в C / C ++?

0 голосов
/ 05 августа 2020

NumPy предоставляет четко определенные C API, так что можно легко обрабатывать массив NumPy в пространстве C / C ++. Например, если у меня есть функция C, которая принимает массивы C (указатели) в качестве аргументов, я могу просто #include <numpy/arrayobject.h> и передать ей массив NumPy, обратившись к его члену data (или используя C API PyArray_DATA).

Недавно я хочу добиться того же для CuPy, но я не могу найти файл заголовка, который я мог бы включить. Чтобы быть точным c, моя цель следующая:

  • У меня есть несколько ядер CUDA и их вызывающие программы, написанные на C / C ++. Вызывающие программы запускаются на хосте, но в качестве аргументов принимают дескрипторы буферов памяти на устройстве. Вычисленные результаты вызывающих абонентов также хранятся на устройстве.
  • Я хочу обернуть вызывающих абонентов в Python функции, чтобы я мог контролировать, когда передавать данные с устройства на хост в Python. Это означает, что я должен заключить указатели памяти устройства в объекты Python. ndarray CuPy - лучший выбор, о котором я могу думать.

Я не могу использовать механизм CuPy, определяемый пользователем-kenrel, потому что функции, которые я хочу обернуть, не являются непосредственно ядрами CUDA. Они должны содержать код хоста.

В настоящее время я нашел обходной путь. Я пишу функции Python в cython, которые принимают массивы CuPy в качестве входных данных и возвращают массивы CuPy. А затем я преобразовал атрибут .data.ptr в тип size_t C, а затем преобразовал его в любой нужный мне тип указателя. Ниже приводится пример кода.

Пример кода 

//kernel.cu

#include <math.h>

__global__ void vecSumKernel(float *A, float *B, float *C, int n) {
    int i = threadIdx.x + blockIdx.x * blockDim.x;

    if (i < n)
        C[i] = A[i] + B[i];
}

// This is the C function I want to wrap into Python.
// Notice it does not allocate any memory on device. I want that to be done by cupy.
extern "C" void vecSum(float *A_d, float *B_d, float *C_d, int n) {
    int threadsPerBlock = 512;
    if (threadsPerBlock > n) threadsPerBlock = n;
    int nBlocks = (int)ceilf((float)n / (float)threadsPerBlock);

    vecSumKernel<<<nBlocks, threadsPerBlock>>>(A_d, B_d, C_d, n);
}

//kernel.h

#ifndef KERNEL_H_
#define KERNEL_H_

void vecSum(float *A_d, float *B_d, float *C_d, int n);

#endif

# test_module.pyx

import cupy as cp
import numpy as np


cdef extern from "kernel.h":

    void vecSum(float *A_d, float *B_d, float *C_d, int n)


cdef vecSum_wrapper(size_t aPtr, size_t bPtr, size_t cPtr, int n):
    # here the Python int -- cp.ndarray.data.ptr -- is first cast to size_t,
    # and then cast to (float *).

    vecSum(<float*>aPtr, <float*>bPtr, <float*>cPtr, n)


# This is the Python function I want to use
# a, b are cupy arrays
def vec_sum(a, b):
    a_ptr = a.data.ptr
    b_ptr = b.data.ptr

    n = a.shape[0]

    output = cp.empty(shape=(n,), dtype=a.dtype)
    c_ptr = output.data.ptr

    vecSum_wrapper(a_ptr, b_ptr, c_ptr, n)
    return output

Компиляция и запуск

Для компиляции можно сначала скомпилировать kernel.cu в стат c библиотека, скажем, libVecSum. Затем используйте cython для компиляции test_module.pyx int test_module.c и создайте расширение Python как обычно. 

# setup.py

from setuptools import Extension, setup

ext_module = Extension(
    "cupyExt.test_module",
    sources=["cupyExt/test_module.c"],
        library_dirs=["cupyExt/"],
        libraries=['libVecSum', 'cudart'])

setup(
    name="cupyExt",
    version="0.0.0",
    ext_modules = [ext_module],
)

Кажется, работает. 

>>> import cupy as cp
>>> from cupyExt import test_module
>>> a = cp.ones(5, dtype=cp.float32) * 3
>>> b = cp.arange(5, dtype=cp.float32)
>>> c = test_module.vec_sum(a, b)
>>> print(c.device)
<CUDA Device 0>
>>> print(c)
[3. 4. 5. 6. 7.]

Любые лучшие способы ?

Я не уверен, безопасен ли этот способ для памяти. Я также считаю, что приведение указателей от .data.ptr к C не очень хорошо. Я хочу знать мысли и комментарии людей по этому поводу.

...