[Types] Represent boolean type as uint1

allo/ir/types.py

@@ -195,7 +195,9 @@ def __repr__(self):
         return f"Stream({self.dtype}[{shape}])"
 
 
-bool = Int(1)
+# boolean type should not be used as i1!
+bool = UInt(1)
+# signed integer types
 int1 = Int(1)
 int8 = Int(8)
 int16 = Int(16)
@@ -217,6 +219,7 @@ def __repr__(self):
 int13 = Int(13)
 int14 = Int(14)
 int15 = Int(15)
+# unsigned integer types
 uint1 = UInt(1)
 uint8 = UInt(8)
 uint16 = UInt(16)
@@ -238,7 +241,9 @@ def __repr__(self):
 uint13 = UInt(13)
 uint14 = UInt(14)
 uint15 = UInt(15)
+# index type
 index = Index()
+# floating point types
 float16 = Float(16)
 float32 = Float(32)
 float64 = Float(64)

allo/utils.py

@@ -24,6 +24,7 @@
     "i16": np.int16,
     "i32": np.int32,
     "i64": np.int64,
+    "ui1": np.bool_,
     "ui8": np.uint8,
     "ui16": np.uint16,
     "ui32": np.uint32,
@@ -38,6 +39,7 @@
     "i16": ctypes.c_int16,
     "i32": ctypes.c_int32,
     "i64": ctypes.c_int64,
+    "ui1": ctypes.c_bool,
     "ui8": ctypes.c_uint8,
     "ui16": ctypes.c_uint16,
     "ui32": ctypes.c_uint32,

tests/test_types.py

@@ -10,6 +10,7 @@
     Float,
     Fixed,
     UFixed,
+    bool,
     uint1,
     int32,
     float32,
@@ -407,6 +408,24 @@ def vadd[
     np.testing.assert_allclose(np_A + np_B, allo_C, rtol=1e-5)
 
 
+def test_boolean():
+    def kernel(A: bool[16]) -> bool[16]:
+        B: bool[16] = 0
+        for i in range(16):
+            if A[i]:
+                B[i] = 1
+            else:
+                B[i] = 0
+        return B
+
+    s = allo.customize(kernel)
+    print(s.module)
+    mod = s.build()
+    np_A = np.random.randint(0, 2, size=(16)).astype(np.bool)
+    np_B = mod(np_A)
+    np.testing.assert_array_equal(np_A, np_B)
+
+
 ######################################################################
 # Legacy tests
 ######################################################################