Added tests for where and type utility functions

Author: ndgrigorian

dpctl/tensor/_search_functions.py

@@ -89,7 +89,7 @@ def where(condition, x1, x2):
 
     deps = []
     wait_list = []
-    if x1_dtype is not dst_dtype:
+    if x1_dtype != dst_dtype:
         _x1 = dpt.empty_like(x1, dtype=dst_dtype)
         ht_copy1_ev, copy1_ev = ti._copy_usm_ndarray_into_usm_ndarray(
             src=x1, dst=_x1, sycl_queue=exec_q
@@ -98,7 +98,7 @@ def where(condition, x1, x2):
         deps.append(copy1_ev)
         wait_list.append(ht_copy1_ev)
 
-    if x2_dtype is not dst_dtype:
+    if x2_dtype != dst_dtype:
         _x2 = dpt.empty_like(x2, dtype=dst_dtype)
         ht_copy2_ev, copy2_ev = ti._copy_usm_ndarray_into_usm_ndarray(
             src=x2, dst=_x2, sycl_queue=exec_q

dpctl/tensor/_type_utils.py

@@ -84,7 +84,7 @@ def _all_data_types(_fp16, _fp64):
             ]
 
 
-def is_maximal_inexact_type(dt: dpt.dtype, _fp16: bool, _fp64: bool):
+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
@@ -106,7 +106,7 @@ def _can_cast(from_: dpt.dtype, to_: dpt.dtype, _fp16: bool, _fp64: bool):
         if (
             from_.kind in "biu"
             and to_.kind in "fc"
-            and is_maximal_inexact_type(to_, _fp16, _fp64)
+            and _is_maximal_inexact_type(to_, _fp16, _fp64)
         ):
             return True
 

dpctl/tests/test_type_utils.py

@@ -0,0 +1,68 @@
+#                      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 pytest
+
+import dpctl.tensor as dpt
+from dpctl.tensor._type_utils import (
+    _all_data_types,
+    _can_cast,
+    _is_maximal_inexact_type,
+)
+
+
+def test_all_data_types():
+    fp16_fp64_types = set([dpt.float16, dpt.float64, dpt.complex128])
+    fp64_types = set([dpt.float64, dpt.complex128])
+
+    all_dts = _all_data_types(True, True)
+    assert fp16_fp64_types.issubset(all_dts)
+
+    all_dts = _all_data_types(True, False)
+    assert dpt.float16 in all_dts
+    assert not fp64_types.issubset(all_dts)
+
+    all_dts = _all_data_types(False, True)
+    assert dpt.float16 not in all_dts
+    assert fp64_types.issubset(all_dts)
+
+    all_dts = _all_data_types(False, False)
+    assert not fp16_fp64_types.issubset(all_dts)
+
+
+@pytest.mark.parametrize("fp16", [True, False])
+@pytest.mark.parametrize("fp64", [True, False])
+def test_maximal_inexact_types(fp16, fp64):
+    assert not _is_maximal_inexact_type(dpt.int32, fp16, fp64)
+    assert fp64 == _is_maximal_inexact_type(dpt.float64, fp16, fp64)
+    assert fp64 == _is_maximal_inexact_type(dpt.complex128, fp16, fp64)
+    assert fp64 != _is_maximal_inexact_type(dpt.float32, fp16, fp64)
+    assert fp64 != _is_maximal_inexact_type(dpt.complex64, fp16, fp64)
+
+
+def test_can_cast_device():
+    assert _can_cast(dpt.int64, dpt.float64, True, True)
+    # if f8 is available, can't cast i8 to f4
+    assert not _can_cast(dpt.int64, dpt.float32, True, True)
+    assert not _can_cast(dpt.int64, dpt.float32, False, True)
+    # should be able to cast to f8 when f2 unavailable
+    assert _can_cast(dpt.int64, dpt.float64, False, True)
+    # casting to f4 acceptable when f8 unavailable
+    assert _can_cast(dpt.int64, dpt.float32, True, False)
+    assert _can_cast(dpt.int64, dpt.float32, False, False)
+    # can't safely cast inexact type to inexact type of lesser precision
+    assert not _can_cast(dpt.float32, dpt.float16, True, False)
+    assert not _can_cast(dpt.float64, dpt.float32, False, True)

dpctl/tests/test_usm_ndarray_search_functions.py

@@ -20,6 +20,9 @@
 from numpy.testing import assert_array_equal
 
 import dpctl.tensor as dpt
+from dpctl.tensor._search_functions import _where_result_type
+from dpctl.tensor._type_utils import _all_data_types
+from dpctl.utils import ExecutionPlacementError
 
 _all_dtypes = [
     "u1",
@@ -38,6 +41,12 @@
 ]
 
 
+class mock_device:
+    def __init__(self, fp16, fp64):
+        self.has_aspect_fp16 = fp16
+        self.has_aspect_fp64 = fp64
+
+
 def test_where_basic():
     get_queue_or_skip()
 
@@ -54,7 +63,16 @@ def test_where_basic():
     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()
+
+    out = dpt.where(cond, dpt.ones(cond.shape), dpt.zeros(cond.shape))
+    assert (dpt.asnumpy(out) == dpt.asnumpy(out_expected)).all()
 
+    out = dpt.where(
+        cond,
+        dpt.ones(cond.shape[0])[:, dpt.newaxis],
+        dpt.zeros(cond.shape[0])[:, dpt.newaxis],
+    )
     assert (dpt.asnumpy(out) == dpt.asnumpy(out_expected)).all()
 
 
@@ -70,6 +88,31 @@ def _dtype_all_close(x1, x2):
         return np.allclose(x1, x2)
 
 
+@pytest.mark.parametrize("dt1", _all_dtypes)
+@pytest.mark.parametrize("dt2", _all_dtypes)
+@pytest.mark.parametrize("fp16", [True, False])
+@pytest.mark.parametrize("fp64", [True, False])
+def test_where_result_types(dt1, dt2, fp16, fp64):
+    dev = mock_device(fp16, fp64)
+
+    dt1 = dpt.dtype(dt1)
+    dt2 = dpt.dtype(dt2)
+    res_t = _where_result_type(dt1, dt2, dev)
+
+    if fp16 and fp64:
+        assert res_t == dpt.result_type(dt1, dt2)
+    else:
+        if res_t:
+            assert res_t.kind == dpt.result_type(dt1, dt2).kind
+        else:
+            # some illegal cases are covered above, but
+            # this guarantees that _where_result_type
+            # produces None only when one of the dtypes
+            # is illegal given fp aspects of device
+            all_dts = _all_data_types(fp16, fp64)
+            assert dt1 not in all_dts or dt2 not in all_dts
+
+
 @pytest.mark.parametrize("dt1", _all_dtypes)
 @pytest.mark.parametrize("dt2", _all_dtypes)
 def test_where_all_dtypes(dt1, dt2):
@@ -78,17 +121,39 @@ def test_where_all_dtypes(dt1, dt2):
     skip_if_dtype_not_supported(dt2, q)
 
     cond = dpt.asarray([False, False, False, True, True], sycl_queue=q)
-    x1 = dpt.asarray(2, sycl_queue=q)
-    x2 = dpt.asarray(3, sycl_queue=q)
-
+    x1 = dpt.asarray(2, dtype=dt1, sycl_queue=q)
+    x2 = dpt.asarray(3, dtype=dt2, sycl_queue=q)
     res = dpt.where(cond, x1, x2)
+
     res_check = np.asarray([3, 3, 3, 2, 2], dtype=res.dtype)
+    assert _dtype_all_close(dpt.asnumpy(res), res_check)
 
-    dev = q.sycl_device
+    # contiguous cases
+    x1 = dpt.full(cond.shape, 2, dtype=dt1, sycl_queue=q)
+    x2 = dpt.full(cond.shape, 3, dtype=dt2, sycl_queue=q)
+    res = dpt.where(cond, x1, x2)
+    assert _dtype_all_close(dpt.asnumpy(res), res_check)
 
-    if not dev.has_aspect_fp16 or not dev.has_aspect_fp64:
-        assert res.dtype.kind == dpt.result_type(x1.dtype, x2.dtype).kind
 
+@pytest.mark.parametrize("dt1", _all_dtypes)
+@pytest.mark.parametrize("dt2", _all_dtypes)
+def test_where_mask_dtypes(dt1, dt2):
+    q = get_queue_or_skip()
+    skip_if_dtype_not_supported(dt1, q)
+    skip_if_dtype_not_supported(dt2, q)
+
+    cond = dpt.asarray([0, 1, 3, 0, 10], dtype=dt1, sycl_queue=q)
+    x1 = dpt.asarray(2, dtype=dt2, sycl_queue=q)
+    x2 = dpt.asarray(3, dtype=dt2, sycl_queue=q)
+    res = dpt.where(cond, x1, x2)
+
+    res_check = np.asarray([3, 2, 2, 3, 2], dtype=res.dtype)
+    assert _dtype_all_close(dpt.asnumpy(res), res_check)
+
+    # contiguous cases
+    x1 = dpt.full(cond.shape, 2, dtype=dt2, sycl_queue=q)
+    x2 = dpt.full(cond.shape, 3, dtype=dt2, sycl_queue=q)
+    res = dpt.where(cond, x1, x2)
     assert _dtype_all_close(dpt.asnumpy(res), res_check)
 
 
@@ -116,12 +181,14 @@ def test_where_empty():
 
     assert_array_equal(dpt.asnumpy(res), res_np)
 
+    # check that broadcasting is performed
+    with pytest.raises(ValueError):
+        dpt.where(empty, x1, dpt.empty((1, 2)))
+
 
-@pytest.mark.parametrize("dt", _all_dtypes)
 @pytest.mark.parametrize("order", ["C", "F"])
-def test_where_contiguous(dt, order):
-    q = get_queue_or_skip()
-    skip_if_dtype_not_supported(dt, q)
+def test_where_contiguous(order):
+    get_queue_or_skip()
 
     cond = dpt.asarray(
         [
@@ -131,14 +198,81 @@ def test_where_contiguous(dt, order):
             [[False, False, False], [True, False, True]],
             [[True, True, True], [True, False, True]],
         ],
-        sycl_queue=q,
         order=order,
     )
 
-    x1 = dpt.full(cond.shape, 2, dtype=dt, order=order, sycl_queue=q)
-    x2 = dpt.full(cond.shape, 3, dtype=dt, order=order, sycl_queue=q)
+    x1 = dpt.full(cond.shape, 2, order=order)
+    x2 = dpt.full(cond.shape, 3, order=order)
+    expected = np.where(dpt.asnumpy(cond), dpt.asnumpy(x1), dpt.asnumpy(x2))
+    res = dpt.where(cond, x1, x2)
+
+    assert _dtype_all_close(dpt.asnumpy(res), expected)
+
+
+def test_where_contiguous1D():
+    get_queue_or_skip()
+
+    cond = dpt.asarray([True, False, True, False, False, True])
+
+    x1 = dpt.full(cond.shape, 2)
+    x2 = dpt.full(cond.shape, 3)
+    expected = np.where(dpt.asnumpy(cond), dpt.asnumpy(x1), dpt.asnumpy(x2))
+    res = dpt.where(cond, x1, x2)
+    assert _dtype_all_close(dpt.asnumpy(res), expected)
+
+    # test with complex dtype (branch in kernel)
+    x1 = dpt.astype(x1, dpt.complex64)
+    x2 = dpt.astype(x2, dpt.complex64)
+    expected = np.where(dpt.asnumpy(cond), dpt.asnumpy(x1), dpt.asnumpy(x2))
+    res = dpt.where(cond, x1, x2)
+    assert _dtype_all_close(dpt.asnumpy(res), expected)
+
+
+def test_where_strided():
+    get_queue_or_skip()
+
+    s0, s1 = 4, 9
+    cond = dpt.reshape(
+        dpt.asarray(
+            [True, False, False, False, True, True, False, True, False] * s0
+        ),
+        (s0, s1),
+    )[:, ::3]
 
+    x1 = dpt.ones((cond.shape[0], cond.shape[1] * 2))[:, ::2]
+    x2 = dpt.zeros((cond.shape[0], cond.shape[1] * 3))[:, ::3]
     expected = np.where(dpt.asnumpy(cond), dpt.asnumpy(x1), dpt.asnumpy(x2))
     res = dpt.where(cond, x1, x2)
 
     assert _dtype_all_close(dpt.asnumpy(res), expected)
+
+
+def test_where_arg_validation():
+    get_queue_or_skip()
+
+    check = dict()
+    x1 = dpt.empty((1,))
+    x2 = dpt.empty((1,))
+
+    with pytest.raises(TypeError):
+        dpt.where(check, x1, x2)
+    with pytest.raises(TypeError):
+        dpt.where(x1, check, x2)
+    with pytest.raises(TypeError):
+        dpt.where(x1, x2, check)
+
+
+def test_where_compute_follows_data():
+    q1 = get_queue_or_skip()
+    q2 = get_queue_or_skip()
+    q3 = get_queue_or_skip()
+
+    x1 = dpt.empty((1,), sycl_queue=q1)
+    x2 = dpt.empty((1,), sycl_queue=q2)
+
+    with pytest.raises(ExecutionPlacementError):
+        dpt.where(dpt.empty((1,), sycl_queue=q1), x1, x2)
+    with pytest.raises(ExecutionPlacementError):
+        dpt.where(dpt.empty((1,), sycl_queue=q3), x1, x2)
+    with pytest.raises(ExecutionPlacementError):
+        dpt.where(x1, x1, x2)