[MLIR][TOSA] Folding cast to bool gives wrong value

[GleasonK]

I'm fixing what I believe to be a bug in StableHLO with our type convert folder, and I believe it may exist in other dialects as well?

Looking at https://www.mlplatform.org/tosa/tosa_spec.html#_cast:

if (is_same<out_t,bool_t>()) {
    out = (in != 0) ? true : false;
}

However with a lack of special casing bools Convert / Cast-like ops these are folded incorrectly:

$ cat /tmp/t.mlir 
func.func @const_f32_bool() -> tensor<i1> {
  %cst = arith.constant dense<10> : tensor<i32>
  %0 = tosa.cast %cst : (tensor<i32>) -> tensor<i1>
  func.return %0 : tensor<i1>
}
$ mlir-opt /tmp/t.mlir --canonicalize
module {
  func.func @const_f32_bool() -> tensor<i1> {
    %0 = "tosa.const"() <{values = dense<false> : tensor<i1>}> : () -> tensor<i1>
    return %0 : tensor<i1>
  }
}

Input was 10 != 0 so probably should be folded to true(?) Not sure if there are other cast-like ops in upstream dialects that may have this issue as well?

Example of my fix for StableHLO in: openxla/xla#29237

[llvmbot]

@llvm/issue-subscribers-mlir

Author: Kevin Gleason (GleasonK)

I'm fixing what I believe to be a bug in StableHLO with our type convert folder, and I believe it may exist in other dialects as well?

Looking at https://www.mlplatform.org/tosa/tosa_spec.html#_cast:

if (is_same<out_t,bool_t>()) {
    out = (in != 0) ? true : false;
}

However with a lack of special casing bools Convert / Cast-like ops these are folded incorrectly:

$ cat /tmp/t.mlir 
func.func @<!-- -->const_f32_bool() -&gt; tensor&lt;i1&gt; {
  %cst = arith.constant dense&lt;10&gt; : tensor&lt;i32&gt;
  %0 = tosa.cast %cst : (tensor&lt;i32&gt;) -&gt; tensor&lt;i1&gt;
  func.return %0 : tensor&lt;i1&gt;
}
$ mlir-opt /tmp/t.mlir --canonicalize
module {
  func.func @<!-- -->const_f32_bool() -&gt; tensor&lt;i1&gt; {
    %0 = "tosa.const"() &lt;{values = dense&lt;false&gt; : tensor&lt;i1&gt;}&gt; : () -&gt; tensor&lt;i1&gt;
    return %0 : tensor&lt;i1&gt;
  }
}

Input was 10 != 0 so probably should be folded to true(?) Not sure if there are other cast-like ops in upstream dialects that may have this issue as well?

[GleasonK]

cc @sjarus maybe?