[mlir][tosa] Fix integer-to-boolean cast folder (#150370)
According to the TOSA spec, casting to boolean should produce true if
the input is non-zero, and false otherwise — i.e., `out = (in != 0) ?
true : false`. Previous behavior incorrectly relied on truncation, which
could yield incorrect results for non-zero values whose least
significant bit is zero. Fixes #150302.
diff --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 606626d..34e7e42 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -1302,9 +1302,11 @@
auto intVal = operand.getSplatValue<APInt>();
auto bitwidth = outETy.getIntOrFloatBitWidth();
- if (trunc) {
+ // i1 types are boolean in TOSA
+ if (outETy.isInteger(1)) {
+ intVal = APInt(bitwidth, intVal.isZero() ? 0 : 1);
+ } else if (trunc) {
intVal = intVal.trunc(bitwidth);
- // i1 types are boolean in TOSA
} else if (unsignIn || inIntType.isInteger(1)) {
intVal = intVal.zext(bitwidth);
} else {
diff --git a/mlir/test/Dialect/Tosa/canonicalize.mlir b/mlir/test/Dialect/Tosa/canonicalize.mlir
index 11c8d54..6b55442 100644
--- a/mlir/test/Dialect/Tosa/canonicalize.mlir
+++ b/mlir/test/Dialect/Tosa/canonicalize.mlir
@@ -1349,3 +1349,14 @@
%1 = "tosa.cast"(%0) : (tensor<i1>) -> tensor<i32>
return %1 : tensor<i32>
}
+
+// -----
+
+// CHECK-LABEL: @test_fold_i32_to_i1_cast
+// CHECK: %[[OUT:.*]] = "tosa.const"() <{values = dense<true> : tensor<i1>}> : () -> tensor<i1>
+// CHECK: return %[[OUT]] : tensor<i1>
+func.func @test_fold_i32_to_i1_cast() -> tensor<i1> {
+ %0 = "tosa.const"() <{values = dense<10> : tensor<i32>}> : () -> tensor<i32>
+ %1 = "tosa.cast"(%0) : (tensor<i32>) -> tensor<i1>
+ return %1 : tensor<i1>
+}