(tvm-tirx) branch main updated: [FIX] Lower LLVM boolean comparisons as unsigned integers (#22)

[tqchen]

This is an automated email from the ASF dual-hosted git repository.

tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git


The following commit(s) were added to refs/heads/main by this push:
     new cce4f6a2 [FIX] Lower LLVM boolean comparisons as unsigned integers 
(#22)
cce4f6a2 is described below

commit cce4f6a27c10b60c424d9da3af072f7a1824ede9
Author: Tianqi Chen <tqc...@users.noreply.github.com>
AuthorDate: Thu Oct 22 07:13:45 2026 -0400

    [FIX] Lower LLVM boolean comparisons as unsigned integers (#22)
---
 src/codegen/llvm/codegen_llvm.cc        |  6 +++---
 tests/python/codegen/llvm/test_basic.py | 20 ++++++++++++++++++++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/src/codegen/llvm/codegen_llvm.cc b/src/codegen/llvm/codegen_llvm.cc
index 0a54fe46..a9448e03 100644
--- a/src/codegen/llvm/codegen_llvm.cc
+++ b/src/codegen/llvm/codegen_llvm.cc
@@ -1595,7 +1595,7 @@ DEFINE_CODEGEN_BINARY_OP(Mul);
   llvm::Value* CodeGenLLVM::Create##Op(PrimType t, llvm::Value* a, 
llvm::Value* b) {    \
     if (t.MatchesCode(DLDataTypeCode::kDLInt)) {                               
         \
       return builder_->CreateICmpS##Op(a, b);                                  
         \
-    } else if (t.MatchesCode(DLDataTypeCode::kDLUInt)) {                       
         \
+    } else if (t.MatchesCode(DLDataTypeCode::kDLUInt, 
DLDataTypeCode::kDLBool)) {       \
       return builder_->CreateICmpU##Op(a, b);                                  
         \
     } else {                                                                   
         \
       TVM_FFI_ICHECK(t.MatchesCode(DLDataTypeCode::kDLFloat));                 
         \
@@ -1655,7 +1655,7 @@ llvm::Value* CodeGenLLVM::Dispatch_(const prim::EQNode* 
op) {
   llvm::Value* a = MakeValue(op->a);
   llvm::Value* b = MakeValue(op->b);
   PrimType dtype(op->a.ty()->dtype);
-  if (dtype.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) {
+  if (dtype.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt, 
DLDataTypeCode::kDLBool)) {
     return builder_->CreateICmpEQ(a, b);
   } else {
     return builder_->CreateFCmpOEQ(a, b);
@@ -1666,7 +1666,7 @@ llvm::Value* CodeGenLLVM::Dispatch_(const prim::NENode* 
op) {
   llvm::Value* a = MakeValue(op->a);
   llvm::Value* b = MakeValue(op->b);
   PrimType dtype(op->a.ty()->dtype);
-  if (dtype.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) {
+  if (dtype.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt, 
DLDataTypeCode::kDLBool)) {
     return builder_->CreateICmpNE(a, b);
   } else {
     return builder_->CreateFCmpONE(a, b);
diff --git a/tests/python/codegen/llvm/test_basic.py 
b/tests/python/codegen/llvm/test_basic.py
index af1d5ef7..895963fb 100644
--- a/tests/python/codegen/llvm/test_basic.py
+++ b/tests/python/codegen/llvm/test_basic.py
@@ -32,6 +32,26 @@ from tvm_tirx.support import utils
 from tvm_tirx.testing import env
 
 
+@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
+def test_bool_comparisons():
+    @T.function
+    def compare(A: T.Tensor((4,), "bool"), B: T.Tensor((4,), "bool"), C: 
T.Tensor((6, 4), "bool")):
+        for i in range(4):
+            C[0, i] = A[i] == B[i]
+            C[1, i] = A[i] != B[i]
+            C[2, i] = A[i] < B[i]
+            C[3, i] = A[i] <= B[i]
+            C[4, i] = A[i] > B[i]
+            C[5, i] = A[i] >= B[i]
+
+    a = np.array([False, False, True, True])
+    b = np.array([False, True, False, True])
+    output = tvm_tirx.runtime.empty((6, 4), "bool")
+    compiled = tvm_tirx.tirx.compile(compare, target="llvm")
+    compiled(tvm_tirx.runtime.tensor(a), tvm_tirx.runtime.tensor(b), output)
+    np.testing.assert_array_equal(output.numpy(), [a == b, a != b, a < b, a <= 
b, a > b, a >= b])
+
+
 @pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
 def test_duplicate_function_global_symbol_diagnostic():
     @I.ir_module