Support for passing a Kokkos::View to python

[bartgol]

I would like to be able to pass my existing C++ Kokkos views to python, so that they can later be consumed by numpy/cupy/pytorch/somethingelse framework. Moreover, I would like the python objects to adopt view semantic (so NOT creating a clone of the input array), as I would like the py code to update the C++ arrays content.

My use case is a large C++ library trying to replace small parts with ML emulators implemented in python. Ideally, I in C++ would do something like py_module.attr("my_func")(view1,view2,...,viewN);, where py_module here is a pybind11 module, but a raw PyObject call is still ok. In this case, view1,...,viewN have to be py-interoperable objects, so I cannot pass a Kokkos::View directly.

My current approach is to create pybind11::array objects (setting up the data pointer as well as shape, stride, data type, etc), and pass that as my viewN args. On the py side, the fcn my_func has to take care of converting the numpy array to the proper framework. E.g., if the Kokkos views were on Cuda space, and the py module plans to use cupy, the py module needs to perform somethign like this:

#########################################################
def get_cu_array(np_arr):
#########################################################
    shape   = np_arr.shape
    dtype   = np_arr.dtype
    ptr     = np_arr.__array_interface__['data'][0]
    strides = np_arr.strides

    # The exact size here does not really matter, as we are just creating an
    # unmanaged mem block, of which we then simply grab the start address.
    # Still, use the correct size for code clarity
    size = shape[0]*strides[0]
    mem = cp.cuda.UnownedMemory(ptr=ptr,owner=None,size=size)
    memptr = cp.cuda.MemoryPointer(mem, 0)

    return cp.ndarray(shape=shape,dtype=dtype,memptr=memptr,strides=strides)

That is, the pybind11::array is used as a "data container" (at least for the non-CPU case), and the py module has to extract data ptr and create an Unmanaged array in the proper framework.

This seems a bit too complicated. Ideally, I'd like to do this:

auto py_view1 = Kokkos::to_pykokkos(view1);
...
auto py_viewN = Kokkos::to_pykokkos(viewN);
my_module.attr("my_func")(py_view1,...,py_viewN);

and in my py module

import pykokkos as pk
from something import my_cool_fcn

# say, 2 args
def my_func(pk_view1, pk_view2): 
   # on_cuda set in the module init before calling this
   if (on_cuda):
     arr1 = pk.to_cupy(pk_view1)
     arr2 = pk.to_cupy(pk_view2)
  else:
     arr1 = pk.to_numpy(pk_view1)
     arr2 = pk.to_numpy(pk_view2)

  my_cool_fcn(arr1,arr2)

and my_cool_fcn may be calling other stuff (e.g., pytorch).

Another option would be to make pykokkos implement from_dlpack(x), to align with the lightweight design pattern of numpy/cupy/torch when it comes to sharing nd arrays among different frameworks.

[JBludau]

Having thought about this for a bit longer, I think this is really something that should have a higher priority, especially since I expect a lot of doe codes would benefit from the possibility to integrate python based ML models into their applications.