Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: support indexing parray by ndarray #89

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

feat: support indexing parray by ndarray #89

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
472405c
fix: bug in reset coherence
yinengy Apr 14, 2023
d9d16eb
feat: support index by array
yinengy Apr 14, 2023
92617c8
fix: minor bug in __getitem__
yinengy Apr 14, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/python/parla/common/parray/coherence.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,7 @@ def reset(self, new_owner: int):
"""
for device_id in self._local_states.keys():
self._local_states[device_id] = self.INVALID
self._versions[device_id] = None
self._versions[device_id] = -1
self._is_complete[device_id] = None

self._local_states[new_owner] = self.MODIFIED
Expand Down
26 changes: 22 additions & 4 deletions src/python/parla/common/parray/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -305,7 +305,22 @@ def print_overview(self) -> None:

# slicing/indexing

def __getitem__(self, slices: SlicesType) -> PArray | Any:
def __getitem__(self, slices: SlicesType | numpy.ndarray | cupy.ndarray) -> PArray | Any:
"""
Acceptable Slices: Slice, Int, tuple of (Slice, Int, List of Int), ndarray of numpy/cupy
Example:
A[0] # int
A[:] # slice
A[0,:,10] # tuple of int slice int
A[2:10:2, 0, [1, 3, 5]] # tuple of slice int list of Int

Note: `:` equals to slice(None, None, None)
Note: `None` or tuple of `None` is not acceptable (even if `numpy.ndarray` accept `None`)
# TODO: support `None`
"""
if isinstance(slices, numpy.ndarray) or isinstance(slices, cupy.ndarray):
slices = slices.tolist()

if self._slices: # resolve saved slices first
ret = self.array[slices]
else:
Expand All @@ -329,9 +344,9 @@ def __getitem__(self, slices: SlicesType) -> PArray | Any:
else:
return ret

def __setitem__(self, slices: SlicesType, value: PArray | ndarray | Any) -> None:
def __setitem__(self, slices: SlicesType | numpy.ndarray | cupy.ndarray, value: PArray | ndarray | Any) -> None:
"""
Acceptable Slices: Slice, Int, tuple of (Slice, Int, List of Int)
Acceptable Slices: Slice, Int, tuple of (Slice, Int, List of Int), ndarray of numpy/cupy
Example:
A[0] # int
A[:] # slice
Expand All @@ -340,11 +355,14 @@ def __setitem__(self, slices: SlicesType, value: PArray | ndarray | Any) -> None

Note: `:` equals to slice(None, None, None)
Note: `None` or tuple of `None` is not acceptable (even if `numpy.ndarray` accept `None`)
# TODO: support `None` and `ndarray` as slices
# TODO: support `None`
"""
if isinstance(value, PArray):
value = value.array

if isinstance(slices, numpy.ndarray) or isinstance(slices, cupy.ndarray):
slices = slices.tolist()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Isn't this inefficient as:

  • numpy supports indexing by both list and np.ndarray of int32/int64/uint32/uint64 type
  • cupy support indexing by list, and both cp.ndarray/np.ndarray ofint32/int64/uint32/uint64 type

This could lead to a lot of device -> host copies from the crosspy calls?

Copy link
Contributor Author

@yinengy yinengy Apr 14, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This could only avoid by caller since this information has to be saved in cpu side that is becase it will later be converted and hashed so it could be recorded for further use within parray. (so indexing by cupy array is not efficient in any senario since parray's internal data structure is still running in cpu)

And I think crosspy should always use numpy array as index, what is the benefit of using cupy array as index? @bozhiyou

Copy link
Contributor

@bozhiyou bozhiyou Apr 14, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

And I think crosspy should always use numpy array as index

No, indices could be as large as the array itself - think of the permutation example arr[shuffle(arange(len(arr))] - and even larger. If you are indexing a cupy array, the indices array has to be copied to GPU as a cupy array.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree with Will that list will be inefficient. Storing all the information on CPU seems to be inefficient as well.

Copy link
Contributor Author

@yinengy yinengy Apr 14, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That makes sense, but currently there is no efficient way to deal with this issue like simply storing the info in GPU or as ndarray (actually all lists will be converted to hash maps in the later step since there is no gurantee the slice mapping is in sequential when you deal with local indices). And each access to array need to query these information (e.g. comparing slices with given slices). Consider the following scenario:

  1. Data initialy stored in cpu then moved to gpu, but user don't know which device it is, should the user index with cp.ndarray or by np.ndarray?
  2. Same as 1 but now in a multidevice task like where crosspy is in, which slices should user use to index into it? (and which device should it be if the cupy.ndarray is used?)
  3. Which device should the cupy ndarray be stored inside parray?
    (a) If we chose to put it on the same device as where data is, when there are multiple copy on different device, each device now has part of the slicing saved, every access regardless where the data is has to access all devices and compare slice information that saved locally, that will be really slow.
    (b) if we chose to put it on the first gpu, operation happened for copy on other devices has to query this gpu to, that is the same as query data in cpu since nvlink is not guranted to be exist between all device and this request will be routed via cpu, which is slower than put it on cpu. This also raise memory issues when the slices is large (you mentioned it might be as large as array), GPU memory are more limited than cpu and runtime has to track the slice size to overwise new data moved there will be OOM, and this also cause imbalanced workload and memory size between devices.
  4. when data is moved, should we move the slices mapping together? Will this increase the data movement overhead?

until we found a good solution to the above issues, CPU is still the best choice to store the slices. (but I could make it numpy array instead of python list)


if self._slices: # resolve saved slices first
self.array.__setitem__(slices, value)
else:
Expand Down
8 changes: 8 additions & 0 deletions testing/python/test_parray.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,14 @@ def check_array_update():
assert a[-1] == 4
assert a._coherence.owner == 1


@spawn(ts[5], dependencies=[ts[4]], placement=gpu(0), inout=[(a,0)])
def check_array_indexing():
a[np.array([0,2])] = 0
assert len(a) == 4
assert a[0] == 0
assert a._coherence.owner == 0

if __name__=="__main__":
test_parray_creation()
test_parray_task()