Use manager per queue in tensor implementation

Move implementation from dpctl/utils/__init__.py into
dedicated files.

Author: oleksandr-pavlyk

dpctl/tensor/_accumulation.py

@@ -126,7 +126,7 @@ def _accumulate_common(
             out = dpt.permute_dims(out, perm)
 
     final_ev = dpctl.SyclEvent()
-    _manager = SequentialOrderManager
+    _manager = SequentialOrderManager[q]
     depends = _manager.submitted_events
     if implemented_types:
         if not include_initial:

dpctl/tensor/_clip.py

@@ -299,7 +299,7 @@ def _clip_none(x, val, out, order, _binary_fn):
             x = dpt.broadcast_to(x, res_shape)
         if val_ary.shape != res_shape:
             val_ary = dpt.broadcast_to(val_ary, res_shape)
-        _manager = SequentialOrderManager
+        _manager = SequentialOrderManager[exec_q]
         dep_evs = _manager.submitted_events
         ht_binary_ev, binary_ev = _binary_fn(
             src1=x, src2=val_ary, dst=out, sycl_queue=exec_q, depends=dep_evs
@@ -322,7 +322,7 @@ def _clip_none(x, val, out, order, _binary_fn):
             buf = _empty_like_orderK(val_ary, res_dt)
         else:
             buf = dpt.empty_like(val_ary, dtype=res_dt, order=order)
-        _manager = SequentialOrderManager
+        _manager = SequentialOrderManager[exec_q]
         dep_evs = _manager.submitted_events
         ht_copy_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
             src=val_ary, dst=buf, sycl_queue=exec_q, depends=dep_evs
@@ -449,7 +449,7 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
             else:
                 out = dpt.empty_like(x, order=order)
 
-        _manager = SequentialOrderManager
+        _manager = SequentialOrderManager[exec_q]
         dep_evs = _manager.submitted_events
         ht_copy_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
             src=x, dst=out, sycl_queue=exec_q, depends=dep_evs
@@ -672,7 +672,7 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
                 a_min = dpt.broadcast_to(a_min, res_shape)
             if a_max.shape != res_shape:
                 a_max = dpt.broadcast_to(a_max, res_shape)
-            _manager = SequentialOrderManager
+            _manager = SequentialOrderManager[exec_q]
             dep_ev = _manager.submitted_events
             ht_binary_ev, binary_ev = ti._clip(
                 src=x,
@@ -700,7 +700,7 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
                 buf2 = _empty_like_orderK(a_max, buf2_dt)
             else:
                 buf2 = dpt.empty_like(a_max, dtype=buf2_dt, order=order)
-            _manager = SequentialOrderManager
+            _manager = SequentialOrderManager[exec_q]
             dep_ev = _manager.submitted_events
             ht_copy_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
                 src=a_max, dst=buf2, sycl_queue=exec_q, depends=dep_ev
@@ -756,7 +756,7 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
                 buf1 = _empty_like_orderK(a_min, buf1_dt)
             else:
                 buf1 = dpt.empty_like(a_min, dtype=buf1_dt, order=order)
-            _manager = SequentialOrderManager
+            _manager = SequentialOrderManager[exec_q]
             dep_ev = _manager.submitted_events
             ht_copy_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
                 src=a_min, dst=buf1, sycl_queue=exec_q, depends=dep_ev
@@ -825,7 +825,7 @@ def clip(x, /, min=None, max=None, out=None, order="K"):
         else:
             buf1 = dpt.empty_like(a_min, dtype=buf1_dt, order=order)
 
-        _manager = SequentialOrderManager
+        _manager = SequentialOrderManager[exec_q]
         dep_evs = _manager.submitted_events
         ht_copy1_ev, copy1_ev = ti._copy_usm_ndarray_into_usm_ndarray(
             src=a_min, dst=buf1, sycl_queue=exec_q, depends=dep_evs

dpctl/tensor/_copy_utils.py

@@ -40,13 +40,14 @@ def _copy_to_numpy(ary):
     if not isinstance(ary, dpt.usm_ndarray):
         raise TypeError(f"Expected dpctl.tensor.usm_ndarray, got {type(ary)}")
     nb = ary.usm_data.nbytes
-    hh = dpm.MemoryUSMHost(nb, queue=ary.sycl_queue)
+    q = ary.sycl_queue
+    hh = dpm.MemoryUSMHost(nb, queue=q)
     h = np.ndarray(nb, dtype="u1", buffer=hh).view(ary.dtype)
     itsz = ary.itemsize
     strides_bytes = tuple(si * itsz for si in ary.strides)
     offset = ary.__sycl_usm_array_interface__.get("offset", 0) * itsz
     # ensure that content of ary.usm_data is final
-    dpctl.utils.SequentialOrderManager.wait()
+    q.wait()
     hh.copy_from_device(ary.usm_data)
     return np.ndarray(
         ary.shape,
@@ -105,7 +106,7 @@ def _copy_from_numpy_into(dst, np_ary):
             src_ary = src_ary.astype(np.float32)
         elif src_ary_dt_c == "D":
             src_ary = src_ary.astype(np.complex64)
-    _manager = dpctl.utils.SequentialOrderManager
+    _manager = dpctl.utils.SequentialOrderManager[copy_q]
     dep_ev = _manager.submitted_events
     # synchronizing call
     ti._copy_numpy_ndarray_into_usm_ndarray(
@@ -208,7 +209,7 @@ def _copy_overlapping(dst, src):
         order="C",
         buffer_ctor_kwargs={"queue": q},
     )
-    _manager = dpctl.utils.SequentialOrderManager
+    _manager = dpctl.utils.SequentialOrderManager[q]
     dep_evs = _manager.submitted_events
     hcp1, cp1 = ti._copy_usm_ndarray_into_usm_ndarray(
         src=src, dst=tmp, sycl_queue=q, depends=dep_evs
@@ -232,10 +233,11 @@ def _copy_same_shape(dst, src):
         _copy_overlapping(src=src, dst=dst)
         return
 
-    _manager = dpctl.utils.SequentialOrderManager
+    copy_q = dst.sycl_queue
+    _manager = dpctl.utils.SequentialOrderManager[copy_q]
     dep_evs = _manager.submitted_events
     hev, cpy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
-        src=src, dst=dst, sycl_queue=dst.sycl_queue, depends=dep_evs
+        src=src, dst=dst, sycl_queue=copy_q, depends=dep_evs
     )
     _manager.add_event_pair(hev, cpy_ev)
 
@@ -724,7 +726,7 @@ def _extract_impl(ary, ary_mask, axis=0):
     cumsum_dt = dpt.int32 if mask_nelems < int32_t_max else dpt.int64
     cumsum = dpt.empty(mask_nelems, dtype=cumsum_dt, device=ary_mask.device)
     exec_q = cumsum.sycl_queue
-    _manager = dpctl.utils.SequentialOrderManager
+    _manager = dpctl.utils.SequentialOrderManager[exec_q]
     dep_evs = _manager.submitted_events
     mask_count = ti.mask_positions(
         ary_mask, cumsum, sycl_queue=exec_q, depends=dep_evs
@@ -760,7 +762,7 @@ def _nonzero_impl(ary):
     cumsum = dpt.empty(
         mask_nelems, dtype=cumsum_dt, sycl_queue=exec_q, order="C"
     )
-    _manager = dpctl.utils.SequentialOrderManager
+    _manager = dpctl.utils.SequentialOrderManager[exec_q]
     dep_evs = _manager.submitted_events
     mask_count = ti.mask_positions(
         ary, cumsum, sycl_queue=exec_q, depends=dep_evs
@@ -837,7 +839,7 @@ def _take_multi_index(ary, inds, p):
     res = dpt.empty(
         res_shape, dtype=ary.dtype, usm_type=res_usm_type, sycl_queue=exec_q
     )
-    _manager = dpctl.utils.SequentialOrderManager
+    _manager = dpctl.utils.SequentialOrderManager[exec_q]
     dep_ev = _manager.submitted_events
     hev, take_ev = ti._take(
         src=ary,
@@ -890,7 +892,7 @@ def _place_impl(ary, ary_mask, vals, axis=0):
     cumsum_dt = dpt.int32 if mask_nelems < int32_t_max else dpt.int64
     cumsum = dpt.empty(mask_nelems, dtype=cumsum_dt, device=ary_mask.device)
     exec_q = cumsum.sycl_queue
-    _manager = dpctl.utils.SequentialOrderManager
+    _manager = dpctl.utils.SequentialOrderManager[exec_q]
     dep_ev = _manager.submitted_events
     mask_count = ti.mask_positions(
         ary_mask, cumsum, sycl_queue=exec_q, depends=dep_ev
@@ -990,7 +992,7 @@ def _put_multi_index(ary, inds, p, vals):
     else:
         rhs = dpt.astype(vals, ary.dtype)
     rhs = dpt.broadcast_to(rhs, expected_vals_shape)
-    _manager = dpctl.utils.SequentialOrderManager
+    _manager = dpctl.utils.SequentialOrderManager[exec_q]
     dep_ev = _manager.submitted_events
     hev, put_ev = ti._put(
         dst=ary,

dpctl/tensor/_ctors.py

@@ -159,7 +159,7 @@ def _asarray_from_usm_ndarray(
         )
     eq = dpctl.utils.get_execution_queue([usm_ndary.sycl_queue, copy_q])
     if eq is not None:
-        _manager = dpctl.utils.SequentialOrderManager
+        _manager = dpctl.utils.SequentialOrderManager[eq]
         dep_evs = _manager.submitted_events
         hev, cpy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
             src=usm_ndary, dst=res, sycl_queue=eq, depends=dep_evs
@@ -415,7 +415,7 @@ def _asarray_from_seq(
             sycl_queue=alloc_q,
             order=order,
         )
-        _manager = dpctl.utils.SequentialOrderManager
+        _manager = dpctl.utils.SequentialOrderManager[exec_q]
         _device_copy_walker(seq_obj, res, _manager)
         return res
     else:
@@ -854,7 +854,7 @@ def arange(
     else:
         _step = sc_ty(1)
     _start = _first
-    _manager = dpctl.utils.SequentialOrderManager
+    _manager = dpctl.utils.SequentialOrderManager[sycl_queue]
     # populating newly allocated array, no task dependencies
     hev, lin_ev = ti._linspace_step(_start, _step, res, sycl_queue)
     _manager.add_event_pair(hev, lin_ev)
@@ -1001,7 +1001,7 @@ def ones(
         order=order,
         buffer_ctor_kwargs={"queue": sycl_queue},
     )
-    _manager = dpctl.utils.SequentialOrderManager
+    _manager = dpctl.utils.SequentialOrderManager[sycl_queue]
     # populating new allocation, no dependent events
     hev, full_ev = ti._full_usm_ndarray(1, res, sycl_queue)
     _manager.add_event_pair(hev, full_ev)
@@ -1100,7 +1100,7 @@ def full(
     elif fill_value_type is int and np.issubdtype(dtype, np.integer):
         fill_value = _to_scalar(fill_value, dtype)
 
-    _manager = dpctl.utils.SequentialOrderManager
+    _manager = dpctl.utils.SequentialOrderManager[sycl_queue]
     # populating new allocation, no dependent events
     hev, full_ev = ti._full_usm_ndarray(fill_value, res, sycl_queue)
     _manager.add_event_pair(hev, full_ev)
@@ -1480,7 +1480,7 @@ def linspace(
         start = float(start)
         stop = float(stop)
     res = dpt.empty(num, dtype=dt, usm_type=usm_type, sycl_queue=sycl_queue)
-    _manager = dpctl.utils.SequentialOrderManager
+    _manager = dpctl.utils.SequentialOrderManager[sycl_queue]
     hev, la_ev = ti._linspace_affine(
         start, stop, dst=res, include_endpoint=endpoint, sycl_queue=sycl_queue
     )
@@ -1578,7 +1578,7 @@ def eye(
         buffer_ctor_kwargs={"queue": sycl_queue},
     )
     if n_rows != 0 and n_cols != 0:
-        _manager = dpctl.utils.SequentialOrderManager
+        _manager = dpctl.utils.SequentialOrderManager[sycl_queue]
         hev, eye_ev = ti._eye(k, dst=res, sycl_queue=sycl_queue)
         _manager.add_event_pair(hev, eye_ev)
     return res
@@ -1630,7 +1630,7 @@ def tril(x, /, *, k=0):
             usm_type=x.usm_type,
             sycl_queue=q,
         )
-        _manager = dpctl.utils.SequentialOrderManager
+        _manager = dpctl.utils.SequentialOrderManager[q]
         dep_evs = _manager.submitted_events
         hev, cpy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
             src=x, dst=res, sycl_queue=q, depends=dep_evs
@@ -1652,7 +1652,7 @@ def tril(x, /, *, k=0):
             usm_type=x.usm_type,
             sycl_queue=q,
         )
-        _manager = dpctl.utils.SequentialOrderManager
+        _manager = dpctl.utils.SequentialOrderManager[q]
         dep_evs = _manager.submitted_events
         hev, tril_ev = ti._tril(
             src=x, dst=res, k=k, sycl_queue=q, depends=dep_evs
@@ -1716,7 +1716,7 @@ def triu(x, /, *, k=0):
             usm_type=x.usm_type,
             sycl_queue=q,
         )
-        _manager = dpctl.utils.SequentialOrderManager
+        _manager = dpctl.utils.SequentialOrderManager[q]
         dep_evs = _manager.submitted_events
         hev, cpy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
             src=x, dst=res, sycl_queue=q, depends=dep_evs
@@ -1730,7 +1730,7 @@ def triu(x, /, *, k=0):
             usm_type=x.usm_type,
             sycl_queue=q,
         )
-        _manager = dpctl.utils.SequentialOrderManager
+        _manager = dpctl.utils.SequentialOrderManager[q]
         dep_evs = _manager.submitted_events
         hev, triu_ev = ti._triu(
             src=x, dst=res, k=k, sycl_queue=q, depends=dep_evs

dpctl/tensor/_elementwise_common.py

@@ -236,7 +236,7 @@ def __call__(self, x, /, *, out=None, order="K"):
                 )
 
         exec_q = x.sycl_queue
-        _manager = SequentialOrderManager
+        _manager = SequentialOrderManager[exec_q]
         if buf_dt is None:
             if out is None:
                 if order == "K":
@@ -632,7 +632,7 @@ def __call__(self, o1, o2, /, *, out=None, order="K"):
             )
 
         orig_out = out
-        _manager = SequentialOrderManager
+        _manager = SequentialOrderManager[exec_q]
         if out is not None:
             if not isinstance(out, dpt.usm_ndarray):
                 raise TypeError(

dpctl/tensor/_indexing_functions.py

@@ -121,7 +121,7 @@ def take(x, indices, /, *, axis=None, mode="wrap"):
         res_shape, dtype=x.dtype, usm_type=res_usm_type, sycl_queue=exec_q
     )
 
-    _manager = dpctl.utils.SequentialOrderManager
+    _manager = dpctl.utils.SequentialOrderManager[exec_q]
     deps_ev = _manager.submitted_events
     hev, take_ev = ti._take(
         x, (indices,), res, axis, mode, sycl_queue=exec_q, depends=deps_ev
@@ -278,7 +278,7 @@ def put_vec_duplicates(vec, ind, vals):
         rhs = dpt.astype(vals, x.dtype)
     rhs = dpt.broadcast_to(rhs, val_shape)
 
-    _manager = dpctl.utils.SequentialOrderManager
+    _manager = dpctl.utils.SequentialOrderManager[exec_q]
     deps_ev = _manager.submitted_events
     hev, put_ev = ti._put(
         x, (indices,), rhs, axis, mode, sycl_queue=exec_q, depends=deps_ev
@@ -375,7 +375,7 @@ def place(arr, mask, vals):
     if arr.shape != mask.shape or vals.ndim != 1:
         raise ValueError("Array sizes are not as required")
     cumsum = dpt.empty(mask.size, dtype="i8", sycl_queue=exec_q)
-    _manager = dpctl.utils.SequentialOrderManager
+    _manager = dpctl.utils.SequentialOrderManager[exec_q]
     deps_ev = _manager.submitted_events
     nz_count = ti.mask_positions(
         mask, cumsum, sycl_queue=exec_q, depends=deps_ev

dpctl/tensor/_linear_algebra_functions.py

@@ -189,7 +189,7 @@ def tensordot(x1, x2, axes=2):
             "supported types according to the casting rule ''safe''."
         )
 
-    _manager = SequentialOrderManager
+    _manager = SequentialOrderManager[exec_q]
     if buf1_dt is None and buf2_dt is None:
         out = dpt.empty(
             res_shape,
@@ -408,7 +408,7 @@ def vecdot(x1, x2, axis=-1):
             "supported types according to the casting rule ''safe''."
         )
 
-    _manager = SequentialOrderManager
+    _manager = SequentialOrderManager[exec_q]
     if buf1_dt is None and buf2_dt is None:
         if x1.dtype.kind == "c":
             x1_tmp = _empty_like_orderK(x1, x1.dtype)
@@ -791,7 +791,7 @@ def matmul(x1, x2, out=None, dtype=None, order="K"):
             else "C"
         )
 
-    _manager = SequentialOrderManager
+    _manager = SequentialOrderManager[exec_q]
     if buf1_dt is None and buf2_dt is None:
         if out is None:
             if order == "K":