diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc index 131c8212c597..8882eb523dcd 100644 --- a/src/target/llvm/codegen_llvm.cc +++ b/src/target/llvm/codegen_llvm.cc @@ -928,7 +928,7 @@ llvm::Value* CodeGenLLVM::CreateCast(DataType from, DataType to, llvm::Value* va } else if (to.is_bool()) { if (from.is_float()) { llvm::Constant* zero = llvm::ConstantFP::get(DTypeToLLVMType(from), 0.); - return builder_->CreateFCmpONE(value, zero); + return builder_->CreateFCmpUNE(value, zero); } else { llvm::Constant* zero = llvm::ConstantInt::get(DTypeToLLVMType(from), 0); return builder_->CreateICmpNE(value, zero); diff --git a/tests/python/codegen/test_target_codegen_llvm.py b/tests/python/codegen/test_target_codegen_llvm.py index 88b791d1aa52..e23ed06211d9 100644 --- a/tests/python/codegen/test_target_codegen_llvm.py +++ b/tests/python/codegen/test_target_codegen_llvm.py @@ -378,6 +378,31 @@ def check_llvm(n): check_llvm(64) +@tvm.testing.requires_llvm +def test_llvm_cast_float_to_bool(): + a_np = np.array([0.0, 1.0, np.nan, np.inf], dtype="float32") + n = a_np.shape[0] + + A = te.placeholder((n,), name="A", dtype="float32") + C = te.compute((n,), lambda i: A[i].astype("bool"), name="C") + + # Convert to TIR and create schedule + mod = te.create_prim_func([A, C]) + sch = tir.Schedule(mod) + + # build and invoke the kernel. + f = tvm.compile(sch.mod, target="llvm") + dev = tvm.cpu(0) + + # launch the kernel. + a = tvm.runtime.tensor(a_np, dev) + c = tvm.runtime.empty((n,), dtype="bool", device=dev) + f(a, c) + c_np = np.array([False, True, True, True], dtype="bool") + + tvm.testing.assert_allclose(c.numpy(), c_np) + + @tvm.testing.requires_llvm def test_rank_zero(): def check_llvm(n):