[Mlir-commits] [mlir] 9d0b596 - [mlir][tosa] Fix segmentation fault in case of folding unranked tensor

[Rob Suderman]

Author: Kai Sasaki
Date: 2023-02-13T13:12:25-08:00
New Revision: 9d0b596aada6fb2166dd4f6f58e359fbac483154

URL: https://github.com/llvm/llvm-project/commit/9d0b596aada6fb2166dd4f6f58e359fbac483154
DIFF: https://github.com/llvm/llvm-project/commit/9d0b596aada6fb2166dd4f6f58e359fbac483154.diff

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

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.

Reviewed By: rsuderman

Differential Revision: https://reviews.llvm.org/D143034

Added: 
    

Modified: 
    mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
    mlir/test/Dialect/Tosa/constant_folding.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
index 74325c85de2b3..1a8a5782e11f6 100644
--- a/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
+++ b/mlir/lib/Dialect/Tosa/IR/TosaCanonicalizations.cpp
@@ -803,8 +803,8 @@ OpFoldResult EqualOp::fold(FoldAdaptor adaptor) {
 
   // 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);
   }
 

diff  --git a/mlir/test/Dialect/Tosa/constant_folding.mlir b/mlir/test/Dialect/Tosa/constant_folding.mlir
index 3111b12f231d2..259b2eaf0bd13 100644
--- a/mlir/test/Dialect/Tosa/constant_folding.mlir
+++ b/mlir/test/Dialect/Tosa/constant_folding.mlir
@@ -6,3 +6,11 @@ func.func @test_const(%arg0 : index) -> tensor<4xi32> {
   %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
+}