Skip to content

Commit

Permalink
Interface: use xp_array in Views when View is unmanaged
Browse files Browse the repository at this point in the history
  • Loading branch information
NaderAlAwar committed Mar 10, 2024
1 parent 1862ab5 commit f06cc01
Showing 1 changed file with 13 additions and 2 deletions.
15 changes: 13 additions & 2 deletions pykokkos/interface/views.py
Original file line number Diff line number Diff line change
Expand Up @@ -105,7 +105,10 @@ def fill(self, value: Union[int, float]) -> None:
:param value: the scalar value
"""

self.data.fill(value)
if self.trait is Trait.Unmanaged:
self.xp_array.fill(value)
else:
self.data.fill(value)

def __getitem__(self, key: Union[int, TeamMember, slice, Tuple]) -> Union[int, float, Subview]:
"""
Expand Down Expand Up @@ -145,7 +148,10 @@ def __setitem__(self, key: Union[int, TeamMember], value: Union[int, float]) ->
if "PK_FUSION" in os.environ:
runtime_singleton.runtime.flush_data(self)

self.data[key] = value
if self.trait is Trait.Unmanaged:
self.xp_array[key] = value
else:
self.data[key] = value

def __bool__(self):
# TODO: more complete implementation
Expand Down Expand Up @@ -198,6 +204,9 @@ def __str__(self) -> str:
if "PK_FUSION" in os.environ:
runtime_singleton.runtime.flush_data(self)

if self.trait is Trait.Unmanaged:
return str(self.xp_array)

return str(self.data)

def __deepcopy__(self, memo):
Expand Down Expand Up @@ -502,6 +511,8 @@ def __init__(self, parent_view: Union[Subview, View], data_slice: Union[slice, T

self.data: np.ndarray = parent_view.data[data_slice]
self.dtype = parent_view.dtype
if parent_view.trait is Trait.Unmanaged:
self.xp_array = parent_view.xp_array[data_slice]

is_cpu: bool = self.parent_view.space is MemorySpace.HostSpace
kokkos_lib: ModuleType = km.get_kokkos_module(is_cpu)
Expand Down

0 comments on commit f06cc01

Please sign in to comment.