[MLIR][Tosa] Turn reshape(const()) from canonicalization into fold; fix dynamic shape case

[mgehre-amd]

Summary

1. Turns the canonicalization into a fold, so it can cleanup IR within other passes.

2. When the output of the reshape is a dynamic shaped tensor,

we cannot apply the fold to the constant, because constants are required to have
static shape.

[mgehre-amd]

mlir/test/Dialect/Tosa/canonicalize.mlir

378–379

Note that I modified this test to be non-splat because a test for a splat constant already exists below (see reshape_canonicalize_const_spat)

[jpienaar]

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

836–841

Unless it has single use ?

[eric-k256]

Looks good to me.

mlir/test/Dialect/Tosa/canonicalize.mlir

378–379

I hadn't noticed the typo in _spat before. I'll try to remember to do a typo fix later.

[mgehre-amd]

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

836–841

If it has a single use, then adding the reshaped variant will not duplicate it because the non-reshaped one will be removed.

[mgehre-amd]

mlir/test/Dialect/Tosa/canonicalize.mlir

378–379

I also didn't notice the typo when copying the function name :-D. I have added a typo fix in my commit.

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

if (tosa::ReshapeOp reshapeOp = dyn_cast<tosa::ReshapeOp>(definingOp)) {
           op, op.getType(), reshapeOp.getInput1(), op.getNewShape());
       return success();
     }
 
     return failure();
   }
 };
 
-struct ReshapeConstOptimization : public OpRewritePattern<tosa::ReshapeOp> {
-  using OpRewritePattern<tosa::ReshapeOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(tosa::ReshapeOp op,
-                                PatternRewriter &rewriter) const override {
-    Value input = op.getInput1();
-    ShapedType inputTy = llvm::cast<ShapedType>(input.getType());
-    ShapedType resultTy = llvm::cast<ShapedType>(op.getType());
-
-    if (inputTy.getElementType() != resultTy.getElementType())
-      return rewriter.notifyMatchFailure(op, "element type does not match.");
-
-    // Check if input is constant
-    DenseElementsAttr inputAttr;
-    if (!matchPattern(input, m_Constant(&inputAttr)))
-      return rewriter.notifyMatchFailure(op, "Non-constant input.");
-
-    // Check if has >1 consumer and is not splat
-    if (!input.hasOneUse() && !inputAttr.isSplat())
-      return rewriter.notifyMatchFailure(op,
-                                         "Used more than once or not-splat");
-
-    // Build new const op with correct output shape
-    DenseElementsAttr outputAttr = inputAttr.reshape(
-        llvm::cast<ShapedType>(inputAttr.getType()).clone(op.getNewShape()));
-    rewriter.replaceOpWithNewOp<tosa::ConstOp>(op, resultTy, outputAttr);
-    return success();
-  }
-};
-
 void ReshapeOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
   results.add<ReshapeReshapeOptimization>(context);
-  results.add<ReshapeConstOptimization>(context);
 }
 
 LogicalResult SelectOp::canonicalize(SelectOp op, PatternRewriter &rewriter) {
   auto notOp = op.getPred().getDefiningOp<tosa::LogicalNotOp>();
   if (!notOp)
     return failure();
   rewriter.updateRootInPlace(op, [&]() {
     op.getOperation()->setOperands(
OpFoldResult ReshapeOp::fold(FoldAdaptor adaptor) {
   auto outputTy = llvm::dyn_cast<RankedTensorType>(getType());
 
   if (!inputTy || !outputTy)
     return {};
 
   if (inputTy == outputTy)
     return getInput1();
 
+  // Constants must have static shape.
+  if (!outputTy.hasStaticShape())
+    return {};
+
   auto operand = llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getInput1());
-  if (operand && outputTy.hasStaticShape() && operand.isSplat()) {
+  if (!operand)
+    return {};
+
+  // Okay to duplicate splat constants.
+  if (operand.isSplat()) {
     return SplatElementsAttr::get(outputTy, operand.getSplatValue<Attribute>());
   }
 
+  // Don't duplicate other constants.
+  if (!getInput1().hasOneUse())
     return {};
+
+  return operand.reshape(
+      llvm::cast<ShapedType>(operand.getType()).clone(getNewShape()));
 }
 
 OpFoldResult PadOp::fold(FoldAdaptor adaptor) {
   // If the pad is all zeros we can fold this operation away.
   if (adaptor.getPadding()) {
     auto densePad = llvm::cast<DenseElementsAttr>(adaptor.getPadding());
     if (densePad.isSplat() && densePad.getSplatValue<APInt>().isZero()) {
       return getInput1();

mlir/test/Dialect/Tosa/canonicalize.mlir

func.func @reshape_canonicalize_double(%arg0: tensor<?x10xf32>) -> tensor<?x5xf32> {
   // CHECK: %[[VAR0:.+]] = "tosa.reshape"(%arg0) <{new_shape = array<i64: -1, 5>}
   // CHECK: return %[[VAR0]]
   %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 5, -1>}: (tensor<?x10xf32>) -> tensor<5x?xf32>
   %1 = "tosa.reshape"(%0) {new_shape = array<i64: -1, 5>}: (tensor<5x?xf32>) -> tensor<?x5xf32>
   return %1 : tensor<?x5xf32>
 }
 
 // CHECK-LABEL: @reshape_canonicalize_const
-func.func @reshape_canonicalize_const() -> tensor<1x10xi32> {
-  // CHECK: %[[VAR0:.+]] = "tosa.const"() <{value = dense<0> : tensor<1x10xi32>}
+func.func @reshape_canonicalize_const() -> tensor<1x5xi32> {
+  // CHECK: %[[VAR0:.+]] = "tosa.const"() <{value = dense<{{\[\[}}0, 1, 2, 3, 4]]> : tensor<1x5xi32>}
   // CHECK: return %[[VAR0]]
-  %0 = "tosa.const"() {value = dense<0> : tensor<10xi32>} : () -> tensor<10xi32>
-  %1 = "tosa.reshape"(%0) {new_shape = array<i64: 1, 10>} : (tensor<10xi32>) -> tensor<1x10xi32>
-  return %1 : tensor<1x10xi32>
+  %0 = "tosa.const"() {value = dense<[0, 1, 2, 3, 4]> : tensor<5xi32>} : () -> tensor<5xi32>
+  %1 = "tosa.reshape"(%0) {new_shape = array<i64: 1, 5>} : (tensor<5xi32>) -> tensor<1x5xi32>
+  return %1 : tensor<1x5xi32>
+}
+
+// CHECK-LABEL: @reshape_canonicalize_const_dynamic
+func.func @reshape_canonicalize_const_dynamic() -> tensor<1x?xi32> {
+  // CHECK: tosa.reshape
+  %0 = "tosa.const"() {value = dense<[0, 1, 2, 3, 4]> : tensor<5xi32>} : () -> tensor<5xi32>
+  %1 = "tosa.reshape"(%0) {new_shape = array<i64: 1, 5>} : (tensor<5xi32>) -> tensor<1x?xi32>
+  return %1 : tensor<1x?xi32>
 }
 
-// CHECK-LABEL: @reshape_canonicalize_const_spat
-func.func @reshape_canonicalize_const_spat() -> (tensor<10xi32>, tensor<1x10xi32>) {
+// CHECK-LABEL: @reshape_canonicalize_const_splat
+func.func @reshape_canonicalize_const_splat() -> (tensor<10xi32>, tensor<1x10xi32>) {
   // CHECK-DAG: %[[VAR0:.+]] = "tosa.const"() <{value = dense<0> : tensor<10xi32>}
   // CHECK-DAG: %[[VAR1:.+]] = "tosa.const"() <{value = dense<0> : tensor<1x10xi32>}
   // CHECK: return %[[VAR0]], %[[VAR1]]
   %0 = "tosa.const"() {value = dense<0> : tensor<10xi32>} : () -> tensor<10xi32>
   %1 = "tosa.reshape"(%0) {new_shape = array<i64: 1, 10>} : (tensor<10xi32>) -> tensor<1x10xi32>
   return %0 , %1 : tensor<10xi32>, tensor<1x10xi32>
 }