Support is not

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

Author: xiaoqiqi177

src/cuda/tile/_ast2ir.py

@@ -186,7 +186,7 @@ def _binop_expr(binop: ast.BinOp, block: ir.Block, ctx: _Context) -> ir.Var:
 
 _cmp_map = {
     ast.Eq: operator.eq, ast.NotEq: operator.ne, ast.Lt: operator.lt, ast.LtE: operator.le,
-    ast.Gt: operator.gt, ast.GtE: operator.ge, ast.Is: operator.is_
+    ast.Gt: operator.gt, ast.GtE: operator.ge, ast.Is: operator.is_, ast.IsNot: operator.is_not,
 }
 
 

src/cuda/tile/_ir/ops.py

@@ -726,16 +726,22 @@ def comparison(fn: str, x: Var, y: Var) -> Var:
     return raw_comparison(fn, x, y)
 
 
+def _is_none_compare(x: Var, y: Var, *, negate: bool, op_name: str) -> Var:
+    x_is_none = x.get_type() is NONE
+    y_is_none = y.get_type() is NONE
+    if not (x_is_none or y_is_none):
+        raise TileTypeError(f"Operator '{op_name}' expects one of the operands to be None")
+    return loosely_typed_const((x_is_none == y_is_none) ^ negate)
+
+
 @impl(operator.is_)
 def operator_is_impl(x: Var, y: Var):
-    x_ty = x.get_type()
-    y_ty = y.get_type()
-    if x_ty is NONE:
-        return loosely_typed_const(y_ty is NONE)
-    elif y_ty is NONE:
-        return loosely_typed_const(x_ty is NONE)
-    else:
-        raise TileTypeError("Operator 'is' expects one of the operands to be None")
+    return _is_none_compare(x, y, negate=False, op_name="is")
+
+
+@impl(operator.is_not)
+def operator_is_not_impl(x: Var, y: Var):
+    return _is_none_compare(x, y, negate=True, op_name="is not")
 
 
 @impl(operator.eq, fixed_args=["eq"])

src/cuda/tile/_ir/typing_support.py

@@ -103,6 +103,7 @@ def wrapped(handler: TypeHandler):
     operator.gt: lambda x, y, /: None,
     operator.ge: lambda x, y, /: None,
     operator.is_: lambda x, y, /: None,
+    operator.is_not: lambda x, y, /: None,
     operator.invert: lambda x, /: None,
     operator.not_: lambda x, /: None,
     operator.pos: lambda x, /: None,

test/test_binary_elementwise.py

@@ -278,6 +278,34 @@ def test_array_compare(shape, tile, dtype, op_symbol, op_func, tmp_path):
         assert_equal(z, ref)
 
 
+def make_is_operator_kernel(cmp):
+    @ct.kernel
+    def is_operator(x):
+        bid = ct.bid(0)
+        a = 1 if cmp is None else -1
+        ct.store(x, index=(bid,), tile=a)
+    return is_operator
+
+
+def make_is_not_operator_kernel(cmp):
+    @ct.kernel
+    def is_not_operator(x):
+        bid = ct.bid(0)
+        a = -1 if cmp is not None else 1
+        ct.store(x, index=(bid,), tile=a)
+    return is_not_operator
+
+
+@pytest.mark.parametrize("make_kernel", [make_is_operator_kernel, make_is_not_operator_kernel])
+@pytest.mark.parametrize("cmp", [None, 1])
+def test_is_or_not_operator(make_kernel, cmp):
+    x = torch.zeros((1,), dtype=torch.int32, device='cuda')
+    kernel = make_kernel(cmp)
+    ct.launch(torch.cuda.current_stream(), (1, 1, 1), kernel, (x, ))
+    ref = 1 if cmp is None else -1
+    assert_equal(x, torch.tensor([ref], dtype=torch.int32, device='cuda'))
+
+
 @pytest.mark.parametrize("max_func", ["max", "ct.maximum"])
 @pytest.mark.parametrize("dtype", int_dtypes + float_dtypes, ids=dtype_id)
 def test_array_max(shape, tile, dtype, tmp_path, max_func):

test/test_constfold.py

@@ -95,24 +95,19 @@ def kernel():
     compile_tile(kernel, (), CompilerOptions())
 
 
-def test_is_op_on_constant():
-
-    def kernel():
-        None is None
-        None is ct.bid(0)
-        ct.bid(0) is None
-
-    compile_tile(kernel, (), CompilerOptions())
-
-
-def test_is_op_on_none_constant():
+@pytest.mark.parametrize("negate", [False, True])
+def test_is_or_not_op_on_none_constant(negate):
 
     def kernel():
         tx = ct.full((1,), 0, ct.float32)
         ty = ct.full((1,), 0, ct.float32)
-        tx is ty
+        if negate:
+            tx is not ty
+        else:
+            tx is ty
 
-    msg = re.escape("Operator 'is' expects one of the operands to be None")
+    op_name = 'is not' if negate else 'is'
+    msg = re.escape(f"Operator '{op_name}' expects one of the operands to be None")
     with pytest.raises(TileTypeError, match=msg):
         compile_tile(kernel, (), CompilerOptions())