dpctl.tensor.where output preserves memory order of inputs (#1342)

* Where result now keeps order of operands
- Now when operands are cast, stride simplification can still be performed on non-C contiguous inputs
- Implements _empty_like_triple_orderK to allocate output of where

* Adds test for correct order="K" behavior in where

* Adjusted logic in _empty_like_triple_orderK
- Now calls _empty_like_pair_orderK when two arrays are of equal shape and larger than the third

* Changes to order "K" stride sorting
- Dimensions of size 1 are effectively disregarded in sorting

* Fixed typo in _empty_like_orderK

Author: ndgrigorian

dpctl/tensor/_copy_utils.py

@@ -351,7 +351,9 @@ def _empty_like_orderK(X, dt, usm_type=None, dev=None):
         )
     st = list(X.strides)
     perm = sorted(
-        range(X.ndim), key=lambda d: builtins.abs(st[d]), reverse=True
+        range(X.ndim),
+        key=lambda d: builtins.abs(st[d]) if X.shape[d] > 1 else 0,
+        reverse=True,
     )
     inv_perm = sorted(range(X.ndim), key=lambda i: perm[i])
     sh = X.shape
@@ -395,9 +397,14 @@ def _empty_like_pair_orderK(X1, X2, dt, res_shape, usm_type, dev):
     max_ndim = max(nd1, nd2)
     st1 += [0] * (max_ndim - len(st1))
     st2 += [0] * (max_ndim - len(st2))
+    sh1 = list(X1.shape) + [0] * (max_ndim - nd1)
+    sh2 = list(X2.shape) + [0] * (max_ndim - nd2)
     perm = sorted(
         range(max_ndim),
-        key=lambda d: (builtins.abs(st1[d]), builtins.abs(st2[d])),
+        key=lambda d: (
+            builtins.abs(st1[d]) if sh1[d] > 1 else 0,
+            builtins.abs(st2[d]) if sh2[d] > 1 else 0,
+        ),
         reverse=True,
     )
     inv_perm = sorted(range(max_ndim), key=lambda i: perm[i])
@@ -417,6 +424,74 @@ def _empty_like_pair_orderK(X1, X2, dt, res_shape, usm_type, dev):
     return dpt.permute_dims(R, inv_perm)
 
 
+def _empty_like_triple_orderK(X1, X2, X3, dt, res_shape, usm_type, dev):
+    if not isinstance(X1, dpt.usm_ndarray):
+        raise TypeError(f"Expected usm_ndarray, got {type(X1)}")
+    if not isinstance(X2, dpt.usm_ndarray):
+        raise TypeError(f"Expected usm_ndarray, got {type(X2)}")
+    if not isinstance(X3, dpt.usm_ndarray):
+        raise TypeError(f"Expected usm_ndarray, got {type(X3)}")
+    nd1 = X1.ndim
+    nd2 = X2.ndim
+    nd3 = X3.ndim
+    if X1.shape == res_shape and X2.shape == res_shape and len(res_shape) > nd3:
+        return _empty_like_pair_orderK(X1, X2, dt, res_shape, usm_type, dev)
+    elif (
+        X2.shape == res_shape and X3.shape == res_shape and len(res_shape) > nd1
+    ):
+        return _empty_like_pair_orderK(X2, X3, dt, res_shape, usm_type, dev)
+    elif (
+        X1.shape == res_shape and X3.shape == res_shape and len(res_shape) > nd2
+    ):
+        return _empty_like_pair_orderK(X1, X3, dt, res_shape, usm_type, dev)
+    fl1 = X1.flags
+    fl2 = X2.flags
+    fl3 = X3.flags
+    if fl1["C"] or fl2["C"] or fl3["C"]:
+        return dpt.empty(
+            res_shape, dtype=dt, usm_type=usm_type, device=dev, order="C"
+        )
+    if fl1["F"] and fl2["F"] and fl3["F"]:
+        return dpt.empty(
+            res_shape, dtype=dt, usm_type=usm_type, device=dev, order="F"
+        )
+    st1 = list(X1.strides)
+    st2 = list(X2.strides)
+    st3 = list(X3.strides)
+    max_ndim = max(nd1, nd2, nd3)
+    st1 += [0] * (max_ndim - len(st1))
+    st2 += [0] * (max_ndim - len(st2))
+    st3 += [0] * (max_ndim - len(st3))
+    sh1 = list(X1.shape) + [0] * (max_ndim - nd1)
+    sh2 = list(X2.shape) + [0] * (max_ndim - nd2)
+    sh3 = list(X3.shape) + [0] * (max_ndim - nd3)
+    perm = sorted(
+        range(max_ndim),
+        key=lambda d: (
+            builtins.abs(st1[d]) if sh1[d] > 1 else 0,
+            builtins.abs(st2[d]) if sh2[d] > 1 else 0,
+            builtins.abs(st3[d]) if sh3[d] > 1 else 0,
+        ),
+        reverse=True,
+    )
+    inv_perm = sorted(range(max_ndim), key=lambda i: perm[i])
+    st1_sorted = [st1[i] for i in perm]
+    st2_sorted = [st2[i] for i in perm]
+    st3_sorted = [st3[i] for i in perm]
+    sh = res_shape
+    sh_sorted = tuple(sh[i] for i in perm)
+    R = dpt.empty(sh_sorted, dtype=dt, usm_type=usm_type, device=dev, order="C")
+    if max(min(st1_sorted), min(st2_sorted), min(st3_sorted)) < 0:
+        sl = tuple(
+            slice(None, None, -1)
+            if (st1_sorted[i] < 0 and st2_sorted[i] < 0 and st3_sorted[i] < 0)
+            else slice(None, None, None)
+            for i in range(nd1)
+        )
+        R = R[sl]
+    return dpt.permute_dims(R, inv_perm)
+
+
 def copy(usm_ary, order="K"):
     """copy(ary, order="K")
 

dpctl/tensor/_search_functions.py

@@ -19,6 +19,7 @@
 import dpctl.tensor._tensor_impl as ti
 from dpctl.tensor._manipulation_functions import _broadcast_shapes
 
+from ._copy_utils import _empty_like_orderK, _empty_like_triple_orderK
 from ._type_utils import _all_data_types, _can_cast
 
 
@@ -121,7 +122,7 @@ def where(condition, x1, x2):
     deps = []
     wait_list = []
     if x1_dtype != dst_dtype:
-        _x1 = dpt.empty_like(x1, dtype=dst_dtype)
+        _x1 = _empty_like_orderK(x1, dst_dtype)
         ht_copy1_ev, copy1_ev = ti._copy_usm_ndarray_into_usm_ndarray(
             src=x1, dst=_x1, sycl_queue=exec_q
         )
@@ -130,39 +131,22 @@ def where(condition, x1, x2):
         wait_list.append(ht_copy1_ev)
 
     if x2_dtype != dst_dtype:
-        _x2 = dpt.empty_like(x2, dtype=dst_dtype)
+        _x2 = _empty_like_orderK(x2, dst_dtype)
         ht_copy2_ev, copy2_ev = ti._copy_usm_ndarray_into_usm_ndarray(
             src=x2, dst=_x2, sycl_queue=exec_q
         )
         x2 = _x2
         deps.append(copy2_ev)
         wait_list.append(ht_copy2_ev)
 
+    dst = _empty_like_triple_orderK(
+        condition, x1, x2, dst_dtype, res_shape, dst_usm_type, exec_q
+    )
+
     condition = dpt.broadcast_to(condition, res_shape)
     x1 = dpt.broadcast_to(x1, res_shape)
     x2 = dpt.broadcast_to(x2, res_shape)
 
-    # dst is F-contiguous when all inputs are F contiguous
-    # otherwise, defaults to C-contiguous
-    if all(
-        (
-            condition.flags.fnc,
-            x1.flags.fnc,
-            x2.flags.fnc,
-        )
-    ):
-        order = "F"
-    else:
-        order = "C"
-
-    dst = dpt.empty(
-        res_shape,
-        dtype=dst_dtype,
-        order=order,
-        usm_type=dst_usm_type,
-        sycl_queue=exec_q,
-    )
-
     hev, _ = ti._where(
         condition=condition,
         x1=x1,

dpctl/tests/test_usm_ndarray_search_functions.py

@@ -370,3 +370,45 @@ def test_where_compute_follows_data():
         dpt.where(dpt.empty((1,), dtype="i4", sycl_queue=q3), x1, x2)
     with pytest.raises(ExecutionPlacementError):
         dpt.where(x1, x1, x2)
+
+
+def test_where_order():
+    get_queue_or_skip()
+
+    test_sh = (
+        20,
+        20,
+    )
+    test_sh2 = tuple(2 * dim for dim in test_sh)
+    n = test_sh[-1]
+
+    for dt1, dt2 in zip(["i4", "i4", "f4"], ["i4", "f4", "i4"]):
+        ar1 = dpt.zeros(test_sh, dtype=dt1, order="C")
+        ar2 = dpt.ones(test_sh, dtype=dt2, order="C")
+        condition = dpt.zeros(test_sh, dtype="?", order="C")
+        res = dpt.where(condition, ar1, ar2)
+        assert res.flags.c_contiguous
+
+        ar1 = dpt.ones(test_sh, dtype=dt1, order="F")
+        ar2 = dpt.ones(test_sh, dtype=dt2, order="F")
+        condition = dpt.zeros(test_sh, dtype="?", order="F")
+        res = dpt.where(condition, ar1, ar2)
+        assert res.flags.f_contiguous
+
+        ar1 = dpt.ones(test_sh2, dtype=dt1, order="C")[:20, ::-2]
+        ar2 = dpt.ones(test_sh2, dtype=dt2, order="C")[:20, ::-2]
+        condition = dpt.zeros(test_sh2, dtype="?", order="C")[:20, ::-2]
+        res = dpt.where(condition, ar1, ar2)
+        assert res.strides == (n, -1)
+
+        ar1 = dpt.ones(test_sh2, dtype=dt1, order="C")[:20, ::-2].mT
+        ar2 = dpt.ones(test_sh2, dtype=dt2, order="C")[:20, ::-2].mT
+        condition = dpt.zeros(test_sh2, dtype="?", order="C")[:20, ::-2].mT
+        res = dpt.where(condition, ar1, ar2)
+        assert res.strides == (-1, n)
+
+        ar1 = dpt.ones(n, dtype=dt1, order="C")
+        ar2 = dpt.broadcast_to(dpt.ones(n, dtype=dt2, order="C"), test_sh)
+        condition = dpt.zeros(n, dtype="?", order="C")
+        res = dpt.where(condition, ar1, ar2)
+        assert res.strides == (20, 1)