[Mlir-commits] [mlir] [mlir][tosa] Harden folds/canonicalizations for unranked and dynamic shapes (PR #188188)

[Hocky Yudhiono]

================
@@ -1784,28 +1752,22 @@ OpFoldResult ResizeOp::fold(FoldAdaptor adaptor) {
     return {};
   }
 
-  auto input = getInput();
-  auto inputTy = llvm::cast<RankedTensorType>(input.getType());
-  auto resultTy = llvm::cast<RankedTensorType>(getType());
-  if (inputTy != resultTy)
-    return {};
-
-  return input;
+  return foldToInputIfTypeMatches(*this, getInput());
 }
 
 OpFoldResult ReverseOp::fold(FoldAdaptor adaptor) {
   auto operand = getInput1();
   auto operandTy = llvm::cast<ShapedType>(operand.getType());
   auto axis = getAxis();
-  auto operandAttr =
-      llvm::dyn_cast_if_present<SplatElementsAttr>(adaptor.getInput1());
-  if (operandAttr)
-    return operandAttr;
+  bool noOpReverse =
+      llvm::isa_and_nonnull<SplatElementsAttr>(adaptor.getInput1());
+
+  // If the dim-length is 1, or reversing axis is unit-dim, also a no-op.
+  noOpReverse |= operandTy.hasRank() &&
+                 (operandTy.getRank() == 0 || operandTy.getDimSize(axis) == 1);
 
-  // If the dim-length is 1, tosa.reverse is a no-op.
-  if (operandTy.hasRank() &&
-      (operandTy.getRank() == 0 || operandTy.getDimSize(axis) == 1))
-    return operand;
+  if (noOpReverse)
+    return foldToInputIfTypeMatches(*this, operand);
----------------
hockyy wrote:

```
hocky at hocky:~/llvm-project/build$ mlir-opt huh.mlir 
module {
  func.func @reverse_splat() -> tensor<10xi32> {
    %0 = "tosa.const"() <{values = dense<42> : tensor<10xi32>}> : () -> tensor<10xi32>
    %1 = tosa.reverse %0 {axis = 0 : i32} : (tensor<10xi32>) -> tensor<10xi32>
    return %1 : tensor<10xi32>
  }
}

hocky at hocky:~/llvm-project/build$ mlir-opt huh.mlir  --test-single-fold
module {
  func.func @reverse_splat() -> tensor<10xi32> {
    %0 = "tosa.const"() <{values = dense<42> : tensor<10xi32>}> : () -> tensor<10xi32>
    %1 = tosa.reverse %0 {axis = 0 : i32} : (tensor<10xi32>) -> tensor<10xi32>
    return %1 : tensor<10xi32>
  }
}

hocky at hocky:~/llvm-project/build$ ninja
[10/10] Linking CXX executable bin/mlir-transform-opt
hocky at hocky:~/llvm-project/build$ mlir-opt huh.mlir  --test-single-fold
module {
  func.func @reverse_splat() -> tensor<10xi32> {
    %0 = "tosa.const"() <{values = dense<42> : tensor<10xi32>}> : () -> tensor<10xi32>
    return %0 : tensor<10xi32>
  }
}

```

It seems like changing it to that would fail this testcase 👀

https://github.com/llvm/llvm-project/pull/188188