[mlir][tosa] Fix segmentation fault in case of folding unranked tensor

[Lewuathe]

Summary

Trying to fold the unranked tensor for "tosa.equal" crashes due to null reference. We need to check the dynamic cast result beforehand. This is reported in https://github.com/llvm/llvm-project/issues/60192.

[Lewuathe]

@nicolasvasilache @NatashaKnk @rsuderman Sorry for pinging you them all. Could you review this diff if you get a chance?

[eric-k256]

While waiting for the others to review, you might want to remove the extra include.

mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp

796

This looks like a leftover temporary include.

[rsuderman]

mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp

796

Ditto. Remove the debug code and I'll approve and land.

[Lewuathe]

mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp

796

@eric-k256 @rsuderman Sorry for keeping unnecessary code. My bad. I'll remove it.

[Lewuathe]

Delete unused debug code.

[Lewuathe]

@rsuderman Sorry for bothering you again. I deleted the debug code accordingly. Could you review this change when you get a chance?

[rsuderman]

Yup, I can try to land this week.

mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp

OpFoldResult GreaterEqualOp::fold(FoldAdaptor adaptor) {
   if (!lhsAttr || !rhsAttr)
     return {};
 
   return binaryFolder<APIntFoldGreaterEqual,
                       ComparisonFold<std::greater_equal<APFloat>>>(
       lhsAttr, rhsAttr, resultTy);
 }
 
 OpFoldResult EqualOp::fold(FoldAdaptor adaptor) {
   auto resultTy = getType().dyn_cast<RankedTensorType>();
   auto lhsAttr = adaptor.getInput1().dyn_cast_or_null<DenseElementsAttr>();
   auto rhsAttr = adaptor.getInput2().dyn_cast_or_null<DenseElementsAttr>();
   Value lhs = getInput1();
   Value rhs = getInput2();
   auto lhsTy = lhs.getType().cast<ShapedType>();
 
   // If we are comparing an integer value to itself it is always true. We can
   // not do this with float due to float values.
-  if (lhsTy.getElementType().isa<IntegerType>() && resultTy.hasStaticShape() &&
-      lhs == rhs) {
+  if (lhsTy.getElementType().isa<IntegerType>() && resultTy &&
+      resultTy.hasStaticShape() && lhs == rhs) {
     return DenseElementsAttr::get(resultTy, true);
   }
 
   if (!lhsAttr || !rhsAttr)
     return {};
 
   return binaryFolder<ComparisonFold<std::equal_to<APInt>>,
                       ComparisonFold<std::equal_to<APFloat>>>(lhsAttr, rhsAttr,

mlir/test/Dialect/Tosa/constant_folding.mlir

 // RUN: mlir-opt --test-constant-fold %s | FileCheck %s
 
 // CHECK-LABEL: func @test_const
 func.func @test_const(%arg0 : index) -> tensor<4xi32> {
   // CHECK: "tosa.const"
   %0 = "tosa.const"() {value = dense<[3, 0, 1, 2]> : tensor<4xi32>} : () -> tensor<4xi32>
   return %0 : tensor<4xi32>
 }
+
+// CHECK-LABEL: func @try_fold_equal_with_unranked_tensor
+func.func @try_fold_equal_with_unranked_tensor(%arg0: tensor<4xi32>, %arg1: tensor<i32>) {
+  // CHECK: "tosa.equal"
+  // CHECK-NEXT: return
+  %0 = "tosa.equal"(%arg0, %arg1) : (tensor<4xi32>, tensor<i32>) -> tensor<*xi1>
+  return
+}