CuVec#

import numpy, mymod
arr = numpy.zeros((1337, 42, 7), "float32")
assert all(numpy.mean(arr, axis=(0, 1)) == 0)
print(mymod.myfunc(arr).sum())

import cuvec.cpython as cuvec, numpy, mymod
arr = cuvec.zeros((1337, 42, 7), "float32")
assert all(numpy.mean(arr, axis=(0, 1)) == 0)
print(cuvec.asarray(mymod.myfunc(arr)).sum())

import cuvec.pybind11 as cuvec, numpy, mymod
arr = cuvec.zeros((1337, 42, 7), "float32")
assert all(numpy.mean(arr, axis=(0, 1)) == 0)
print(cuvec.asarray(mymod.myfunc(arr.cuvec)).sum())

import cuvec.swig as cuvec, numpy, mymod
arr = cuvec.zeros((1337, 42, 7), "float32")
assert all(numpy.mean(arr, axis=(0, 1)) == 0)
print(cuvec.retarray(mymod.myfunc(arr.cuvec)).sum())

#include <numpy/arrayobject.h>
#include "mycudafunction.h"

static PyObject *myfunc(PyObject *self, PyObject *args, PyObject *kwargs) {
  PyObject *o_src = NULL;
  PyObject *o_dst = NULL;
  static const char *kwds[] = {"src", "output", NULL};
  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|O", (char **)kwds,
                                   &o_src, &o_dst))
    return NULL;
  PyArrayObject *p_src = (PyArrayObject *)PyArray_FROM_OTF(
    o_src, NPY_FLOAT32, NPY_ARRAY_IN_ARRAY);
  if (p_src == NULL) return NULL;
  float *src = (float *)PyArray_DATA(p_src);

  // hardcode upsampling factor 2
  npy_intp *src_shape = PyArray_SHAPE(p_src);
  Py_ssize_t dst_shape[3];
  dst_shape[2] = src_shape[2] * 2;
  dst_shape[1] = src_shape[1] * 2;
  dst_shape[0] = src_shape[0] * 2;

  PyArrayObject *p_dst = NULL;
  if (o_dst == NULL)
    p_dst = (PyArrayObject *)PyArray_ZEROS(3, dst_shape, NPY_FLOAT32, 0);
  else
    p_dst = (PyArrayObject *)PyArray_FROM_OTF(
      o_dst, NPY_FLOAT32, NPY_ARRAY_INOUT_ARRAY);
  if (p_dst == NULL) return NULL;
  float *dst = (float *)PyArray_DATA(p_dst);

  mycudafunction(dst, src, dst_shape);

  Py_DECREF(p_src);
  return PyArray_Return(p_dst);
}
...
PyMODINIT_FUNC PyInit_mymod(void) {
  ...
  import_array();
  ...
}

#include "cuvec_cpython.cuh"
#include "mycudafunction.h"

static PyObject *myfunc(PyObject *self, PyObject *args, PyObject *kwargs) {
  PyCuVec<float> *src = NULL;
  PyCuVec<float> *dst = NULL;
  static const char *kwds[] = {"src", "output", NULL};
  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O&|O&", (char **)kwds,
                                  &asPyCuVec_f, &src, &asPyCuVec_f, &dst))
    return NULL;


  if (!src) return NULL;


  // hardcode upsampling factor 2

  std::vector<Py_ssize_t> dst_shape = src->shape();
  for (auto &i : dst_shape) i *= 2;




  if (!dst)
    dst = PyCuVec_zeros<float>(dst_shape);



  if (!dst) return NULL;


  mycudafunction(dst->vec.data(), src->vec.data(), dst_shape.data());


  return dst;
}
...
PyMODINIT_FUNC PyInit_mymod(void) {
  ...

  ...
}

#include "cuvec.cuh"
#include "mycudafunction.h"

NDCuVec<float> *myfunc(NDCuVec<float> &src, NDCuVec<float> *output = nullptr) {











  // hardcode upsampling factor 2

  std::vector<size_t> dst_shape = src.shape;
  for (auto &i : dst_shape) i *= 2;




  if (!dst)
    dst = new NDCuVec<float>(dst_shape);



  if (!dst) throw pybind11::value_error("could not allocate output");


  mycudafunction(dst->vec.data(), src.vec.data(), dst_shape.data());


  return dst;
}
using namespace pybind11::literals;
PYBIND11_MODULE(mymod, m){
  ...
  m.def("myfunc", &myfunc, "src"_a, "output"_a = nullptr);
  ...
}

/// SWIG interface file mymod.i (not mymod.cpp)
%module mymod
%include "cuvec.i"
%{
#include "mycudafunction_swig.h"
%}
NDCuVec<float> *myfunc(
  NDCuVec<float> &src, NDCuVec<float> *output = NULL);
































// that's it :)
// all non-wrapper logic is actually moved to the implementation below

void mycudafunction(float *dst, float *src, Py_ssize_t *shape) {
  float *d_src;
  int dst_size = shape[0] * shape[1] * shape[2] * sizeof(float);
  cudaMalloc(&d_src, dst_size / 8);
  cudaMemcpy(d_src, src, dst_size / 8, cudaMemcpyHostToDevice);
  float *d_dst;
  cudaMalloc(&d_dst, dst_size);
  mykernel<<<...>>>(d_dst, d_src, shape[0], shape[1], shape[2]);
  cudaMemcpy(dst, d_dst, dst_size, cudaMemcpyDeviceToHost);
  cudaFree(d_dst);
  cudaFree(d_src);
}

void mycudafunction(float *dst, float *src, Py_ssize_t *shape) {






  mykernel<<<...>>>(dst, src, shape[0], shape[1], shape[2]);
  cudaDeviceSynchronize();


}

void mycudafunction(float *dst, float *src, size_t *shape) {






  mykernel<<<...>>>(dst, src, shape[0], shape[1], shape[2]);
  cudaDeviceSynchronize();


}

#include "cuvec.cuh"
NDCuVec<float> *myfunc(
    NDCuVec<float> &src, NDCuVec<float> *output = NULL) {
  // hardcode upsampling factor 2
  std::vector<size_t> shape = src.shape;
  for (auto &i : shape) i *= 2;
  if (!output) output = new NDCuVec<float>(shape);
  mykernel<<<...>>>(output->vec.data(), src.vec.data(),
                    shape[0], shape[1], shape[2]);
  cudaDeviceSynchronize();
  return output;
}

cmake_minimum_required(VERSION 3.24 FATAL_ERROR)
project(myproj)
find_package(AMYPADcuvec COMPONENTS cuvec REQUIRED)
add_executable(myexe ...)
set_target_properties(myexe PROPERTIES CXX_STANDARD 11)
target_link_libraries(myexe PRIVATE AMYPAD::cuvec)