Support for 4 dimensional CuPy arrays.

kennykos commented 3 months ago

When trying to pass a 4 dimensional cupy array into a PyKokkos workunit, as in the script below,

import cupy as cp
import pykokkos as pk

@pk.workunit
def work(wid, a):
    a[wid][wid][wid][wid] = a[wid][wid][wid][wid] + 1

def main():
    N = 10
    a = cp.ones((N, N, N, N))
    pk.set_default_space(pk.Cuda)
    pk.parallel_for("work", 10, work, a=a)
    print(a)

main()

I am met with the error

Traceback (most recent call last):
  File "/work/09661/gkk345/ls6/3dcapsules/python/development/gridding/tmp4.py", line 15, in <module>
    main()
  File "/work/09661/gkk345/ls6/3dcapsules/python/development/gridding/tmp4.py", line 12, in main
    pk.parallel_for("work", 10, work, a=a)
  File "/work/09661/gkk345/ls6/pykokkos/pykokkos/interface/parallel_dispatch.py", line 155, in parallel_for
    convert_arrays(kwargs)
  File "/work/09661/gkk345/ls6/pykokkos/pykokkos/interface/parallel_dispatch.py", line 135, in convert_arrays
    kwargs[k] = array(v)
                ^^^^^^^^
  File "/work/09661/gkk345/ls6/pykokkos/pykokkos/interface/views.py", line 802, in array
    return from_array(array)
           ^^^^^^^^^^^^^^^^^
  File "/work/09661/gkk345/ls6/pykokkos/pykokkos/interface/views.py", line 763, in from_array
    return from_numpy(np_array, memory_space, layout, array)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/09661/gkk345/ls6/pykokkos/pykokkos/interface/views.py", line 706, in from_numpy
    return View(ret_list, dtype, space=space, trait=Trait.Unmanaged, array=array, layout=layout, cp_array=cp_array)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/09661/gkk345/ls6/pykokkos/pykokkos/interface/views.py", line 264, in __init__
    self._init_view(shape, dtype, space, layout, trait, array, cp_array)
  File "/work/09661/gkk345/ls6/pykokkos/pykokkos/interface/views.py", line 370, in _init_view
    self.array = kokkos_lib.unmanaged_array(array, dtype=self.dtype.value, space=self.space.value, layout=self.layout.value)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/09661/gkk345/ls6/miniconda3/envs/pyse/lib/python3.11/site-packages/pykokkos_base-0.0.7-py3.11-linux-x86_64.egg/kokkos/utility.py", line 206, in unmanaged_array
    return array(*_args, **_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/work/09661/gkk345/ls6/miniconda3/envs/pyse/lib/python3.11/site-packages/pykokkos_base-0.0.7-py3.11-linux-x86_64.egg/kokkos/utility.py", line 176, in array
    raise ValueError(
ValueError: pykokkos-base build only supports 3 ranks. Requested 4 ranks

I do not expect this result as it seems that I can create a view in up to 8 dimensions.

NaderAlAwar commented 3 months ago

Hi @kennykos. The issue here seems to be from pykokkos-base being built with only 3 ranks. Most pykokkos applications use a maximum of 3 dimensions, so the default pykokkos-base installation command listed here sets -DENABLE_VIEW_RANKS=3. The way to fix this is to rebuild pykokkos-base but set this option to 4 instead of 3. Note that this might increase the pykokkos-base build time. Let me know if you run into issues where the build takes too much time.

kennykos commented 2 months ago

This build took around 3 hours, but it solved the issue. Thanks a ton!

NaderAlAwar commented 2 months ago

Closing this issue for now but feel free to reopen it if you need more help.

kokkos / pykokkos

Support for 4 dimensional CuPy arrays. #273