Move eye() to dpctl_ext/tensor and reuse it in dpnp

Author: vlad-perevezentsev

dpctl_ext/tensor/__init__.py

@@ -35,6 +35,7 @@
     to_numpy,
 )
 from dpctl_ext.tensor._ctors import (
+    eye,
     full,
     tril,
     triu,
@@ -58,6 +59,7 @@
     "astype",
     "copy",
     "extract",
+    "eye",
     "from_numpy",
     "full",
     "nonzero",

dpctl_ext/tensor/_ctors.py

@@ -58,6 +58,38 @@ def _cast_fill_val(fill_val, dt):
         return fill_val
 
 
+def _ensure_native_dtype_device_support(dtype, dev) -> None:
+    """Check that dtype is natively supported by device.
+
+    Arg:
+        dtype:
+            Elemental data-type
+        dev (:class:`dpctl.SyclDevice`):
+            The device about which the query is being made.
+    Returns:
+        None
+    Raise:
+        ValueError:
+            if device does not natively support this `dtype`.
+    """
+    if dtype in [dpt.float64, dpt.complex128] and not dev.has_aspect_fp64:
+        raise ValueError(
+            f"Device {dev.name} does not provide native support "
+            "for double-precision floating point type."
+        )
+    if (
+        dtype
+        in [
+            dpt.float16,
+        ]
+        and not dev.has_aspect_fp16
+    ):
+        raise ValueError(
+            f"Device {dev.name} does not provide native support "
+            "for half-precision floating point type."
+        )
+
+
 def _to_scalar(obj, sc_ty):
     """A way to convert object to NumPy scalar type.
     Raises OverflowError if obj can not be represented
@@ -67,6 +99,102 @@ def _to_scalar(obj, sc_ty):
     return zd_arr[()]
 
 
+def eye(
+    n_rows,
+    n_cols=None,
+    /,
+    *,
+    k=0,
+    dtype=None,
+    order="C",
+    device=None,
+    usm_type="device",
+    sycl_queue=None,
+):
+    """
+    eye(n_rows, n_cols=None, /, *, k=0, dtype=None, \
+        device=None, usm_type="device", sycl_queue=None)
+
+    Creates :class:`dpctl.tensor.usm_ndarray` with ones on the `k`-th
+    diagonal.
+
+    Args:
+        n_rows (int):
+            number of rows in the output array.
+        n_cols (int, optional):
+            number of columns in the output array. If ``None``,
+            ``n_cols = n_rows``. Default: ``None``
+        k (int):
+            index of the diagonal, with ``0`` as the main diagonal.
+            A positive value of ``k`` is a superdiagonal, a negative value
+            is a subdiagonal.
+            Raises :exc:`TypeError` if ``k`` is not an integer.
+            Default: ``0``
+        dtype (optional):
+            data type of the array. Can be typestring,
+            a :class:`numpy.dtype` object, :mod:`numpy` char string, or
+            a NumPy scalar type. Default: ``None``
+        order ("C" or "F"):
+            memory layout for the array. Default: ``"C"``
+        device (optional):
+            array API concept of device where the output array
+            is created. ``device`` can be ``None``, a oneAPI filter selector
+            string, an instance of :class:`dpctl.SyclDevice` corresponding to
+            a non-partitioned SYCL device, an instance of
+            :class:`dpctl.SyclQueue`, or a :class:`dpctl.tensor.Device` object
+            returned by :attr:`dpctl.tensor.usm_ndarray.device`.
+            Default: ``None``
+        usm_type (``"device"``, ``"shared"``, ``"host"``, optional):
+            The type of SYCL USM allocation for the output array.
+            Default: ``"device"``
+        sycl_queue (:class:`dpctl.SyclQueue`, optional):
+            The SYCL queue to use
+            for output array allocation and copying. ``sycl_queue`` and
+            ``device`` are complementary arguments, i.e. use one or another.
+            If both are specified, a :exc:`TypeError` is raised unless both
+            imply the same underlying SYCL queue to be used. If both are
+            ``None``, a cached queue targeting default-selected device is
+            used for allocation and population. Default: ``None``
+
+    Returns:
+        usm_ndarray:
+            A diagonal matrix.
+    """
+    if not isinstance(order, str) or len(order) == 0 or order[0] not in "CcFf":
+        raise ValueError(
+            "Unrecognized order keyword value, expecting 'F' or 'C'."
+        )
+    order = order[0].upper()
+    n_rows = operator.index(n_rows)
+    n_cols = n_rows if n_cols is None else operator.index(n_cols)
+    k = operator.index(k)
+    if k >= n_cols or -k >= n_rows:
+        return dpt.zeros(
+            (n_rows, n_cols),
+            dtype=dtype,
+            order=order,
+            device=device,
+            usm_type=usm_type,
+            sycl_queue=sycl_queue,
+        )
+    dpctl.utils.validate_usm_type(usm_type, allow_none=False)
+    sycl_queue = normalize_queue_device(sycl_queue=sycl_queue, device=device)
+    dtype = _get_dtype(dtype, sycl_queue)
+    _ensure_native_dtype_device_support(dtype, sycl_queue.sycl_device)
+    res = dpt.usm_ndarray(
+        (n_rows, n_cols),
+        dtype=dtype,
+        buffer=usm_type,
+        order=order,
+        buffer_ctor_kwargs={"queue": sycl_queue},
+    )
+    if n_rows != 0 and n_cols != 0:
+        _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
+
+
 def _validate_fill_value(fill_val):
     """Validates that `fill_val` is a numeric or boolean scalar."""
     # TODO: verify if `np.True_` and `np.False_` should be instances of

dpnp/dpnp_container.py

@@ -196,7 +196,7 @@ def eye(
         order = "C"
 
     """Creates `dpnp_array` with ones on the `k`th diagonal."""
-    array_obj = dpt.eye(
+    array_obj = dpt_ext.eye(
         N,
         M,
         k=k,