Avoid unecessary synchronizations applying sorter, or type promoting

Author: oleksandr-pavlyk

dpctl/tensor/_searchsorted.py

@@ -1,10 +1,13 @@
 from typing import Literal, Union
 
+import dpctl
 import dpctl.utils as du
 
-from ._copy_utils import astype
+from ._copy_utils import _empty_like_orderK
 from ._ctors import empty
 from ._data_types import int32, int64
+from ._tensor_impl import _copy_usm_ndarray_into_usm_ndarray as ti_copy
+from ._tensor_impl import _take as ti_take
 from ._tensor_sorting_impl import _searchsorted_left, _searchsorted_right
 from ._type_utils import iinfo, isdtype, result_type
 from ._usmarray import usm_ndarray
@@ -74,6 +77,14 @@ def searchsorted(
             "inferred from input arguments."
         )
 
+    if x1.ndim != 1:
+        raise ValueError("First argument array must be one-dimensional")
+
+    x1_dt = x1.dtype
+    x2_dt = x2.dtype
+
+    host_evs = []
+    ev = dpctl.SyclEvent()
     if sorter is not None:
         if not isdtype(sorter.dtype, "integral"):
             raise ValueError(
@@ -84,29 +95,78 @@ def searchsorted(
                 "Sorter array must be one-dimension with the same "
                 "shape as the first argument array"
             )
-        x1 = x1[sorter]
-
-    if x1.ndim != 1:
-        raise ValueError("First argument array must be one-dimensional")
+        res = empty(x1.shape, dtype=x1_dt, usm_type=x1.usm_type, sycl_queue=q)
+        ind = (sorter,)
+        axis = 0
+        wrap_out_of_bound_indices_mode = 0
+        ht_ev, ev = ti_take(
+            x1,
+            ind,
+            res,
+            axis,
+            wrap_out_of_bound_indices_mode,
+            sycl_queue=q,
+            depends=[
+                ev,
+            ],
+        )
+        x1 = res
+        host_evs.append(ht_ev)
 
-    if x1.dtype != x2.dtype:
+    if x1_dt != x2_dt:
         dt = result_type(x1, x2)
-        x1 = astype(x1, dt, copy=None)
-        x2 = astype(x2, dt, copy=None)
+        if x1_dt != dt:
+            x1_buf = _empty_like_orderK(x1, dt)
+            ht_ev, ev = ti_copy(
+                src=x1,
+                dst=x1_buf,
+                sycl_queue=q,
+                depends=[
+                    ev,
+                ],
+            )
+            host_evs.append(ht_ev)
+            x1 = x1_buf
+        if x2_dt != dt:
+            x2_buf = _empty_like_orderK(x2, dt)
+            ht_ev, ev = ti_copy(
+                src=x2,
+                dst=x2_buf,
+                sycl_queue=q,
+                depends=[
+                    ev,
+                ],
+            )
+            host_evs.append(ht_ev)
+            x2 = x2_buf
 
     dst_usm_type = du.get_coerced_usm_type([x1.usm_type, x2.usm_type])
     dst_dt = int32 if x2.size <= iinfo(int32).max else int64
 
-    dst = empty(x2.shape, dtype=dst_dt, usm_type=dst_usm_type, sycl_queue=q)
+    dst = _empty_like_orderK(x2, dst_dt, usm_type=dst_usm_type)
 
     if side == "left":
         ht_ev, _ = _searchsorted_left(
-            hay=x1, needles=x2, positions=dst, sycl_queue=q
+            hay=x1,
+            needles=x2,
+            positions=dst,
+            sycl_queue=q,
+            depends=[
+                ev,
+            ],
         )
     else:
         ht_ev, _ = _searchsorted_right(
-            hay=x1, needles=x2, positions=dst, sycl_queue=q
+            hay=x1,
+            needles=x2,
+            positions=dst,
+            sycl_queue=q,
+            depends=[
+                ev,
+            ],
         )
-    ht_ev.wait()
+
+    host_evs.append(ht_ev)
+    dpctl.SyclEvent.wait_for(host_evs)
 
     return dst