Support compare operations on tuples

Signed-off-by: Qiqi Xiao <qiqix@nvidia.com>

Author: xiaoqiqi177

src/cuda/tile/_ir/ops.py

@@ -797,6 +797,42 @@ def operator_is_not_impl(x: Var, y: Var):
     return _is_none_compare(x, y, negate=True, op_name="is not")
 
 
+def _tuple_comparison(fn: str, x: Var, y: Var) -> Var:
+    if fn not in ("eq", "ne"):
+        raise TileTypeError(f"Operator '{fn}' is not supported for tuples")
+
+    x_ty = require_tuple_type(x)
+    y_ty = require_tuple_type(y)
+
+    if x.is_constant() and y.is_constant():
+        res = x.get_constant() == y.get_constant()
+        return loosely_typed_const(res if fn == "eq" else not res)
+
+    if len(x_ty) != len(y_ty):
+        return loosely_typed_const(fn == "ne")
+
+    x_items = x.get_aggregate().items
+    y_items = y.get_aggregate().items
+
+    for item in (*x_items, *y_items):
+        item_ty = item.get_type()
+        if isinstance(item_ty, TileTy) and item_ty.ndim > 0:
+            raise TileTypeError("Tuple comparison is not supported for N-D tile elements")
+        if not isinstance(item_ty, (TileTy, TupleTy, LooselyTypedScalar, DTypeSpec, StringTy)):
+            raise TileTypeError(
+                f"Tuple comparison is not supported for elements of type {item_ty}"
+            )
+
+    elem_cmps = [comparison_operator_impl("eq", xi, yi) for xi, yi in zip(x_items, y_items)]
+    result = functools.reduce(lambda a, b: binary_bitwise("and_", a, b), elem_cmps,
+                              loosely_typed_const(True))
+
+    if fn == "ne":
+        result = logical_not_impl(result)
+
+    return result
+
+
 @impl(operator.eq, fixed_args=["eq"])
 @impl(operator.ne, fixed_args=["ne"])
 @impl(operator.lt, fixed_args=["lt"])
@@ -812,6 +848,8 @@ def comparison_operator_impl(fn: str, x: Var, y: Var) -> Var:
             return _binop_propagate_constant(fn, x_ty.dtype, y_ty.dtype, None)
         case StringTy(), StringTy():
             return _binop_propagate_constant(fn, x_ty.value, y_ty.value, None)
+        case TupleTy(), TupleTy():
+            return _tuple_comparison(fn, x, y)
         case _, _:
             return comparison(fn, x, y)
 

test/test_tuple.py

@@ -135,3 +135,133 @@ def kernel():
 
     with pytest.raises(TileTypeError, match=re.escape("Expected a tuple after *")):
         ct.launch(torch.cuda.current_stream(), (1,), kernel, ())
+
+
+def test_tuple_compare_empty_eq():
+    @ct.kernel
+    def kernel(x):
+        if () == ():
+            ct.scatter(x, (), 1)
+        else:
+            ct.scatter(x, (), 0)
+
+    x = torch.zeros((), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (x,))
+    assert x.item() == 1
+
+
+def test_tuple_compare_constants_eq():
+    @ct.kernel
+    def kernel(x):
+        if (1, 2, 3) == (1, 2, 3):
+            ct.scatter(x, (), 1)
+        else:
+            ct.scatter(x, (), 0)
+
+    x = torch.zeros((), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (x,))
+    assert x.item() == 1
+
+
+def test_tuple_compare_constants_ne():
+    @ct.kernel
+    def kernel(x):
+        if (1, 2) != (1, 3):
+            ct.scatter(x, (), 1)
+        else:
+            ct.scatter(x, (), 0)
+
+    x = torch.zeros((), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (x,))
+    assert x.item() == 1
+
+
+def test_tuple_compare_different_lengths():
+    @ct.kernel
+    def kernel(x):
+        a = ct.bid(0)
+        if (a, 1) != (a, 1, 2):
+            ct.scatter(x, (), 1)
+        else:
+            ct.scatter(x, (), 0)
+
+    x = torch.zeros((), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (x,))
+    assert x.item() == 1
+
+
+def test_tuple_compare_0d_tiles_eq():
+    @ct.kernel
+    def kernel(x):
+        a = ct.bid(0)
+        b = ct.bid(1)
+        if (a, b) == (0, 0):
+            ct.scatter(x, (a, b), 1)
+        else:
+            ct.scatter(x, (a, b), -1)
+
+    x = torch.zeros((2, 2), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (2, 2), kernel, (x,))
+    assert x.tolist() == [[1, -1], [-1, -1]]
+
+
+def test_tuple_compare_nd_tile_error():
+    @ct.kernel
+    def kernel():
+        t = ct.ones((4,), dtype=ct.int32)
+        if (t,) == (t,):
+            pass
+
+    with pytest.raises(TileTypeError, match="not supported for N-D tile elements"):
+        ct.launch(torch.cuda.current_stream(), (1,), kernel, ())
+
+
+def test_tuple_compare_unsupported_op():
+    @ct.kernel
+    def kernel():
+        if (1, 2) < (3, 4):
+            pass
+
+    with pytest.raises(TileTypeError, match="not supported for tuples"):
+        ct.launch(torch.cuda.current_stream(), (1,), kernel, ())
+
+
+def test_tuple_compare_nested():
+    @ct.kernel
+    def kernel(x):
+        a = ct.bid(0)
+        if ((a, 1), 2) == ((0, 1), 2):
+            ct.scatter(x, (a, ), 1)
+        else:
+            ct.scatter(x, (a, ), -1)
+
+    x = torch.zeros((2, ), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (2, ), kernel, (x,))
+    assert x.tolist() == [1, -1]
+
+
+def test_tuple_compare_array_element_error():
+    @ct.kernel
+    def kernel(x, y):
+        if (x,) == (y,):
+            pass
+
+    with pytest.raises(TileTypeError, match="not supported for elements of type"):
+        ct.launch(torch.cuda.current_stream(), (1,), kernel,
+                  (torch.zeros(4, dtype=torch.int32, device="cuda"),
+                   torch.zeros(4, dtype=torch.int32, device="cuda")))
+
+
+def test_tuple_compare_constant_args():
+    @ct.kernel
+    def kernel(x, M: ct.Constant[int], N: ct.Constant[int]):
+        if (M, N) == (4, 8):
+            ct.scatter(x, (), 1)
+        else:
+            ct.scatter(x, (), -1)
+
+    x = torch.zeros((), dtype=torch.int32, device="cuda")
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (x, 4, 8))
+    assert x.item() == 1
+    ct.launch(torch.cuda.current_stream(), (1,), kernel, (x, 4, 9))
+    assert x.item() == -1