[PYTHON] Make IntImm more like an integer

python/tvm/tir/expr.py

@@ -439,6 +439,7 @@ def __init__(self, dtype, value):
         self.__init_handle_by_constructor__(
             tvm.ir._ffi_api.FloatImm, dtype, value)
 
+
 @tvm._ffi.register_object
 class IntImm(ConstExpr):
     """Int constant.
@@ -455,9 +456,24 @@ def __init__(self, dtype, value):
         self.__init_handle_by_constructor__(
             tvm.ir._ffi_api.IntImm, dtype, value)
 
+    def __hash__(self):
+        return self.value
+
     def __int__(self):
         return self.value
 
+    def __nonzero__(self):
+        return self.value != 0
+
+    def __eq__(self, other):
+        return _ffi_api._OpEQ(self, other)
+
+    def __ne__(self, other):
+        return _ffi_api._OpNE(self, other)
+
+    def __bool__(self):
+        return self.__nonzero__()
+
 
 @tvm._ffi.register_object
 class StringImm(ConstExpr):

tests/python/unittest/test_tir_nodes.py

@@ -302,7 +302,21 @@ def test_buffer_load_store():
     assert isinstance(s, tvm.tir.BufferStore)
 
 
+def test_intimm_cond():
+    x = tvm.runtime.convert(1)
+    y = tvm.runtime.convert(1)
+    s = {x}
+    assert y in s
+    assert x == y
+    assert x < 20
+    assert not (x >= 20)
+    assert x < 10 and y < 10
+    assert not tvm.runtime.convert(x != 1)
+    assert x == 1
+
+
 if __name__ == "__main__":
+    test_intimm_cond()
     test_buffer_load_store()
     test_vars()
     test_prim_func()