From a20ef8e2346b6d329c13b58bb1523cbb6197299f Mon Sep 17 00:00:00 2001 From: PiotrekB416 Date: Fri, 1 Sep 2023 14:08:19 +0200 Subject: [PATCH 1/4] Fix tensor.put possible race condition --- dpctl/tensor/_indexing_functions.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/dpctl/tensor/_indexing_functions.py b/dpctl/tensor/_indexing_functions.py index 188f3cd291..de45dc34a5 100644 --- a/dpctl/tensor/_indexing_functions.py +++ b/dpctl/tensor/_indexing_functions.py @@ -175,6 +175,13 @@ def put(x, indices, vals, /, *, axis=None, mode="wrap"): indices.dtype ) ) + + indices = dpt.flip(indices) + _, local_indices = np.unique(indices, return_index=True) + 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_) @@ -207,7 +214,9 @@ 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 ) - + 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) From b4eb89b9a659b48c670f7717e6d0c2fdcf851111 Mon Sep 17 00:00:00 2001 From: PiotrekB416 Date: Fri, 1 Sep 2023 14:14:18 +0200 Subject: [PATCH 2/4] Add check if indices need to be cleaned --- dpctl/tensor/_indexing_functions.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/dpctl/tensor/_indexing_functions.py b/dpctl/tensor/_indexing_functions.py index de45dc34a5..727984ec1d 100644 --- a/dpctl/tensor/_indexing_functions.py +++ b/dpctl/tensor/_indexing_functions.py @@ -178,9 +178,11 @@ def put(x, indices, vals, /, *, axis=None, mode="wrap"): indices = dpt.flip(indices) _, local_indices = np.unique(indices, return_index=True) - local_indices = dpt.asarray(local_indices) - indices = dpt.take(indices, local_indices) - indices = dpt.flip(indices) + 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) @@ -214,9 +216,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 ) - vals = dpt.flip(vals) - vals = dpt.take(vals, local_indices) - vals = dpt.flip(vals) + 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) From 994b70323653fb169eb32ae3e8929484b7666e73 Mon Sep 17 00:00:00 2001 From: PiotrekB416 Date: Fri, 1 Sep 2023 20:54:00 +0200 Subject: [PATCH 3/4] Add option to disable uniqness check --- dpctl/tensor/_indexing_functions.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/dpctl/tensor/_indexing_functions.py b/dpctl/tensor/_indexing_functions.py index 727984ec1d..5e4625a0ea 100644 --- a/dpctl/tensor/_indexing_functions.py +++ b/dpctl/tensor/_indexing_functions.py @@ -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 @@ -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 race condition when indices repeat, + but may negatively impact performance. """ if not isinstance(x, dpt.usm_ndarray): raise TypeError( @@ -176,9 +180,12 @@ def put(x, indices, vals, /, *, axis=None, mode="wrap"): ) ) - indices = dpt.flip(indices) - _, local_indices = np.unique(indices, return_index=True) - not_unique = len(indices) != len(local_indices) + 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) From 33db448056595d5695c7889b610e825ee7e51d39 Mon Sep 17 00:00:00 2001 From: PiotrekB416 Date: Fri, 1 Sep 2023 21:02:57 +0200 Subject: [PATCH 4/4] Add check if unique_only is bool --- dpctl/tensor/_indexing_functions.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/dpctl/tensor/_indexing_functions.py b/dpctl/tensor/_indexing_functions.py index 5e4625a0ea..374e1bb58f 100644 --- a/dpctl/tensor/_indexing_functions.py +++ b/dpctl/tensor/_indexing_functions.py @@ -146,7 +146,7 @@ def put(x, indices, vals, /, *, axis=None, mode="wrap", unique_only=False): Default: `"wrap"`. unique_only: If True, only unique indices will be used. - This will prevent a race condition when indices repeat, + This will prevent a undefined behavior when indices repeat, but may negatively impact performance. """ if not isinstance(x, dpt.usm_ndarray): @@ -179,6 +179,10 @@ def put(x, indices, vals, /, *, axis=None, mode="wrap", unique_only=False): 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: