Issues to folders used for shape refinement

[sdasgup3]

Here are some of the problematic shape-refinement foldings.

$ cat file1.mlir 
func.func @add_op_test_i1() {
  %0 = stablehlo.constant dense<[false, false, true, true]> : tensor<4xi1>
  %1 = stablehlo.constant dense<[false, true, false, true]> : tensor<4xi1>
  %2 = stablehlo.add %0, %1 : tensor<4xi1>
  check.expect_eq_const %2, dense<[false, true, true, true]> : tensor<4xi1>
  func.return
}

$ ./bin/stablehlo-opt --stablehlo-refine-shapes  file.mlir
func.func @add_op_test_i1() {
    %c = stablehlo.constant dense<[false, true, true, false]> : tensor<4xi1>
    check.expect_eq_const %c, dense<[false, true, true, true]> : tensor<4xi1>
    return
  }


$ cat file.mlir
func.func @select_op_test_si64_scalar() {
  %pred = stablehlo.constant dense<false> : tensor<i1>
  %on_true = stablehlo.constant dense<[2, 3, -1]> : tensor<3xi64>
  %on_false = stablehlo.constant dense<[3, 7, -3]> : tensor<3xi64>
  %result = stablehlo.select %pred, %on_true, %on_false : (tensor<i1>, tensor<3xi64>, tensor<3xi64>) -> tensor<3xi64>
  check.expect_eq_const %result, dense<[3, 7, -3]> : tensor<3xi64>
  func.return
}

$ ./bin/stablehlo-opt --stablehlo-refine-shapes  file.mlir
 func.func @select_op_test_si64_scalar() {
    %c = stablehlo.constant dense<3> : tensor<3xi64>
    check.expect_eq_const %c, dense<[3, 7, -3]> : tensor<3xi64>
    return
  }

$ cat file.mlir
func.func @clamp_op_test_si64_min_scalar() {
  %min = stablehlo.constant dense<[0, 0, -2]> : tensor<3xi64>
  %operand = stablehlo.constant dense<[2, 3, -1]> : tensor<3xi64>
  %max = stablehlo.constant dense<1> : tensor<i64>
  %result = stablehlo.clamp %min, %operand, %max : (tensor<3xi64>, tensor<3xi64>, tensor<i64>) -> tensor<3xi64>
  check.expect_eq_const %result, dense<[1, 1, -1]> : tensor<3xi64>
  func.return
}

$ ./bin/stablehlo-opt --stablehlo-refine-shapes  file.mlir
 func.func @clamp_op_test_si64_min_scalar() {
    %c = stablehlo.constant dense<1> : tensor<3xi64>
    check.expect_eq_const %c, dense<[1, 1, -1]> : tensor<3xi64>
    return
  }

The issues can be categorised into the following:

• For Add: semantics mismatch for boolean add
  • Shape Folder:

    stablehlo/stablehlo/transforms/StablehloRefineShapes.cpp

    Line 324 in de6488a

    [&](APSInt lhs, APSInt rhs) { return lhs + rhs; });

  • Interpreter eval:

    stablehlo/stablehlo/reference/Element.cpp

    Line 341 in de6488a

    [](bool lhs, bool rhs) -> bool { return lhs | rhs; },

• For select: The pred has to be bropadcated to match the True/False operand shape.
• For clamp: The max (and even min) has to be broadcasted to operand shape.