ENH: smoother deep_copy() experience

[tylerjereddy]

We recently discussed that pykokkos should stop making an extra host allocation of a view when operating in i.e., CUDA memory space so that it can be more competitive with CuPy in benchmarks. PyKokkos currently needs double the memory transfers in both directions for the same ufunc kernels.

To empower this, we need to be able to smoothly copy a NumPy array on the host to a pykokkos view on the device. In CuPy, it is a frictionless one-liner:

x_gpu = cp.asarray(x_cpu) # move the data to the current device.

In kokkos/pykokkos at the moment, it appears that the user is burdened with concepts like bitwise identical "mirrors" and deep_copy() to move the data in either direction.

I think pykokkos is still missing some components to do this, and if I try to do this naively in kokkos it feels like quite a bit more friction than CuPy:

    view_host = kokkos.array("view_host", shape=[array_size], space=kokkos.HostSpace)
    view_gpu = kokkos.array("view_gpu", shape=[array_size], space=kokkos.CudaSpace)
    # I want to put something interesting from NumPy in view_host, then copy it to view_gpu
    # and I'm already getting confused as a pykokkos developer...
    mirror_gpu_on_host = kokkos.create_mirror_view(view_gpu)
    kokkos.deep_copy(view_gpu, mirror_gpu_on_host)

Also, if I try to assign to a host view, or the mirror_gpu_on_host from a NumPy array, I can't even put something interesting in there anyway at the moment:

    view_host = kokkos.array("view_host", shape=[array_size], space=kokkos.HostSpace)
    view_host[:] = data_np

Traceback (most recent call last):
  File "run_pk_single_ufunc.py", line 30, in <module>
    run_single_pykokkos_ufunc_bench()
  File "run_pk_single_ufunc.py", line 18, in run_single_pykokkos_ufunc_bench
    view_host[:] = data_np
TypeError: __setitem__(): incompatible function arguments. The following argument types are supported:
    1. (self: kokkos.libpykokkos.KokkosView_float64_HostSpace_LayoutRight_1, arg0: int, arg1: float) -> None
    2. (self: kokkos.libpykokkos.KokkosView_float64_HostSpace_LayoutRight_1, arg0: Tuple[int], arg1: float) -> None

Invoked with: <kokkos.libpykokkos.KokkosView_float64_HostSpace_LayoutRight_1 object at 0x151830933270>, slice(None, None, None), array([0.22733602, 0.31675834, 0.79736546, ..., 0.79951774, 0.89607902,
       0.92564455])

At the time of writing, I don't seem to be able to pass view_gpu successfully into a pykokkos workunit with an error like:

  File "/vast/home/treddy/github_projects/pykokkos/pykokkos/lib/ufuncs.py", line 103, in log
    return math.log(view)
TypeError: must be real number, not kokkos.libpykokkos.KokkosView_float64_CudaSpace_La

This probably stems from confusion about using pykokkos-base vs. pykokkos View objects. Speaking of the latter:

    view = pk.View([array_size], dtype=pk.double)
    kokkos.create_mirror_view(view)

errors out: AttributeError: 'View' object has no attribute 'create_mirror_view'

This seems like an awful lot to worry about for an operation that is going to be super common in the real world.