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

Fix dpctl.tensor.put race condition #1384

Closed
wants to merge 4 commits into from
Closed

Fix dpctl.tensor.put race condition #1384

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a20ef8e
Fix tensor.put possible race condition
PiotrekB416 Sep 1, 2023
b4eb89b
Add check if indices need to be cleaned
PiotrekB416 Sep 1, 2023
994b703
Add option to disable uniqness check
PiotrekB416 Sep 1, 2023
33db448
Add check if unique_only is bool
PiotrekB416 Sep 1, 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
27 changes: 25 additions & 2 deletions dpctl/tensor/_indexing_functions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,7 @@ def take(x, indices, /, *, axis=None, mode="wrap"):
return res


def put(x, indices, vals, /, *, axis=None, mode="wrap"):
def put(x, indices, vals, /, *, axis=None, mode="wrap", unique_only=False):
"""put(x, indices, vals, axis=None, mode="wrap")

Puts values of an array into another array
Expand All @@ -144,6 +144,10 @@ def put(x, indices, vals, /, *, axis=None, mode="wrap"):
negative indices.
"clip" - clips indices to (0 <= i < n)
Default: `"wrap"`.
unique_only:
If True, only unique indices will be used.
This will prevent a undefined behavior when indices repeat,
but may negatively impact performance.
"""
if not isinstance(x, dpt.usm_ndarray):
raise TypeError(
Expand Down Expand Up @@ -175,6 +179,22 @@ def put(x, indices, vals, /, *, axis=None, mode="wrap"):
indices.dtype
)
)
if not isinstance(unique_only, bool):
raise TypeError("`unique_only` expected `bool`, got `{}`".format(
type(unique_only)
))

not_unique = False
if unique_only:
indices = dpt.flip(indices)
_, local_indices = np.unique(indices, return_index=True)
not_unique = len(indices) != len(local_indices)

if not_unique:
local_indices = dpt.asarray(local_indices)
indices = dpt.take(indices, local_indices)
indices = dpt.flip(indices)

queues_.append(indices.sycl_queue)
usm_types_.append(indices.usm_type)
exec_q = dpctl.utils.get_execution_queue(queues_)
Expand Down Expand Up @@ -207,7 +227,10 @@ def put(x, indices, vals, /, *, axis=None, mode="wrap"):
vals = dpt.asarray(
vals, dtype=x.dtype, usm_type=vals_usm_type, sycl_queue=exec_q
)

if not_unique:
vals = dpt.flip(vals)
vals = dpt.take(vals, local_indices)
vals = dpt.flip(vals)
vals = dpt.broadcast_to(vals, val_shape)

hev, _ = ti._put(x, (indices,), vals, axis, mode, sycl_queue=exec_q)
Expand Down