Added utility functions, basic where tests

- Utility functions are for finding an output type for universal and binary functions when the device of allocation lacks fp16 or fp64

Author: ndgrigorian

dpctl/tensor/_search_functions.py

@@ -19,6 +19,25 @@
 import dpctl.tensor._tensor_impl as ti
 from dpctl.tensor._manipulation_functions import _broadcast_shapes
 
+from ._type_utils import _all_data_types, _can_cast
+
+
+def _where_result_type(dt1, dt2, dev):
+    res_dtype = dpt.result_type(dt1, dt2)
+    fp16 = dev.has_aspect_fp16
+    fp64 = dev.has_aspect_fp64
+
+    all_dts = _all_data_types(fp16, fp64)
+    if res_dtype in all_dts:
+        return res_dtype
+    else:
+        for res_dtype_ in all_dts:
+            if _can_cast(dt1, res_dtype_, fp16, fp64) and _can_cast(
+                dt2, res_dtype_, fp16, fp64
+            ):
+                return res_dtype_
+        return None
+
 
 def where(condition, x1, x2):
     if not isinstance(condition, dpt.usm_ndarray):
@@ -52,7 +71,7 @@ def where(condition, x1, x2):
 
     x1_dtype = x1.dtype
     x2_dtype = x2.dtype
-    dst_dtype = dpt.result_type(x1.dtype, x2.dtype)
+    dst_dtype = _where_result_type(x1_dtype, x2_dtype, exec_q.sycl_device)
 
     if condition.size == 0:
         return dpt.asarray(

dpctl/tensor/_type_utils.py

@@ -0,0 +1,113 @@
+#                      Data Parallel Control (dpctl)
+#
+# Copyright 2020-2022 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import dpctl.tensor as dpt
+
+
+def _all_data_types(_fp16, _fp64):
+    if _fp64:
+        if _fp16:
+            return [
+                dpt.bool,
+                dpt.int8,
+                dpt.uint8,
+                dpt.int16,
+                dpt.uint16,
+                dpt.int32,
+                dpt.uint32,
+                dpt.int64,
+                dpt.uint64,
+                dpt.float16,
+                dpt.float32,
+                dpt.float64,
+                dpt.complex64,
+                dpt.complex128,
+            ]
+        else:
+            return [
+                dpt.bool,
+                dpt.int8,
+                dpt.uint8,
+                dpt.int16,
+                dpt.uint16,
+                dpt.int32,
+                dpt.uint32,
+                dpt.int64,
+                dpt.uint64,
+                dpt.float32,
+                dpt.float64,
+                dpt.complex64,
+                dpt.complex128,
+            ]
+    else:
+        if _fp16:
+            return [
+                dpt.bool,
+                dpt.int8,
+                dpt.uint8,
+                dpt.int16,
+                dpt.uint16,
+                dpt.int32,
+                dpt.uint32,
+                dpt.int64,
+                dpt.uint64,
+                dpt.float16,
+                dpt.float32,
+                dpt.complex64,
+            ]
+        else:
+            return [
+                dpt.bool,
+                dpt.int8,
+                dpt.uint8,
+                dpt.int16,
+                dpt.uint16,
+                dpt.int32,
+                dpt.uint32,
+                dpt.int64,
+                dpt.uint64,
+                dpt.float32,
+                dpt.complex64,
+            ]
+
+
+def is_maximal_inexact_type(dt: dpt.dtype, _fp16: bool, _fp64: bool):
+    """
+    Return True if data type `dt` is the
+    maximal size inexact data type
+    """
+    if _fp64:
+        return dt in [dpt.float64, dpt.complex128]
+    return dt in [dpt.float32, dpt.complex64]
+
+
+def _can_cast(from_: dpt.dtype, to_: dpt.dtype, _fp16: bool, _fp64: bool):
+    """
+    Can `from_` be cast to `to_` safely on a device with
+    fp16 and fp64 aspects as given?
+    """
+    can_cast_v = dpt.can_cast(from_, to_)  # ask NumPy
+    if _fp16 and _fp64:
+        return can_cast_v
+    if not can_cast_v:
+        if (
+            from_.kind in "biu"
+            and to_.kind in "fc"
+            and is_maximal_inexact_type(to_, _fp16, _fp64)
+        ):
+            return True
+
+    return can_cast_v

dpctl/tests/test_usm_ndarray_search_functions.py

@@ -0,0 +1,84 @@
+#                      Data Parallel Control (dpctl)
+#
+# Copyright 2020-2022 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import pytest
+from helper import get_queue_or_skip, skip_if_dtype_not_supported
+from numpy.testing import assert_array_equal
+
+import dpctl.tensor as dpt
+
+_all_dtypes = [
+    "u1",
+    "i1",
+    "u2",
+    "i2",
+    "u4",
+    "i4",
+    "u8",
+    "i8",
+    "e",
+    "f",
+    "d",
+    "F",
+    "D",
+]
+
+
+def test_where_basic():
+    get_queue_or_skip
+
+    cond = dpt.asarray(
+        [
+            [True, False, False],
+            [False, True, False],
+            [False, False, True],
+            [False, False, False],
+            [True, True, True],
+        ]
+    )
+    out = dpt.where(cond, dpt.asarray(1), dpt.asarray(0))
+    out_expected = dpt.asarray(
+        [[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 0, 0], [1, 1, 1]]
+    )
+
+    assert (dpt.asnumpy(out) == dpt.asnumpy(out_expected)).all()
+
+
+@pytest.mark.parametrize("dt1", _all_dtypes)
+@pytest.mark.parametrize("dt2", _all_dtypes)
+def test_where_all_dtypes(dt1, dt2):
+    q = get_queue_or_skip()
+    skip_if_dtype_not_supported(dt1, q)
+    skip_if_dtype_not_supported(dt2, q)
+
+    cond_np = np.arange(5) > 2
+    x1_np = np.asarray(2, dtype=dt1)
+    x2_np = np.asarray(3, dtype=dt2)
+
+    cond = dpt.asarray(cond_np, sycl_queue=q)
+    x1 = dpt.asarray(x1_np, sycl_queue=q)
+    x2 = dpt.asarray(x2_np, sycl_queue=q)
+
+    res = dpt.where(cond, x1, x2)
+    res_np = np.where(cond_np, x1_np, x2_np)
+
+    if res.dtype != res_np.dtype:
+        assert res.dtype.kind == res_np.dtype.kind
+        assert_array_equal(dpt.asnumpy(res).astype(res_np.dtype), res_np)
+
+    else:
+        assert_array_equal(dpt.asnumpy(res), res_np)