TOSA: tosa-layerwise-constant-fold: Folder for reshapes with multiple uses

Author: mgehre-amd

mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp

@@ -573,6 +573,38 @@ static APFloat computeReciprocal(const APFloat &floatVal, FloatType floatTy) {
   return recip;
 }
 
+/// Fold reshapes. This is similar to ReshapeOp::fold, but also allows
+/// to fold with multiple users.
+struct TosaFoldConstantReshape
+    : public TosaFoldConstantUnaryElementwise<TosaFoldConstantReshape,
+                                              ReshapeOp> {
+  using TosaFoldConstantUnaryElementwise::TosaFoldConstantUnaryElementwise;
+
+  LogicalResult matchAndRewrite(ReshapeOp op,
+                                PatternRewriter &rewriter) const override {
+    auto inputTensor = op.getOperand();
+    // Check that we can apply folding
+    auto preCondCheck =
+        notifyIfNoTosaDenseConstantTensor(inputTensor, op, rewriter);
+    if (failed(preCondCheck)) {
+      return preCondCheck;
+    }
+
+    // Extract the tensor values
+    DenseElementsAttr inputValues;
+    matchPattern(inputTensor, m_Constant(&inputValues));
+
+    // Check whether this should be folded.
+    if (!constantUnaryOpShouldBeFolded(op, inputValues)) {
+      return rewriter.notifyMatchFailure(
+          op, "expected reshape op to have a single user");
+    }
+    DenseElementsAttr newTensor = inputValues.reshape(op.getType());
+    rewriter.replaceOpWithNewOp<ConstOp>(op, newTensor.getType(), newTensor);
+    return success();
+  }
+};
+
 struct TosaFoldConstantReciprocal
     : public TosaFoldConstantUnaryElementwise<TosaFoldConstantReciprocal, ReciprocalOp> {
   using TosaFoldConstantUnaryElementwise<TosaFoldConstantReciprocal,
@@ -1723,6 +1755,8 @@ void mlir::tosa::populateTosaFoldConstantPatterns(
     MLIRContext *ctx, RewritePatternSet &patterns,
     bool foldSplatOrSingleUseOnly,
     bool enableIntCastFolding) {
+
+  patterns.add<TosaFoldConstantReshape>(ctx, foldSplatOrSingleUseOnly);
   patterns.add<TosaFoldConstantTranspose>(ctx, foldSplatOrSingleUseOnly);
   patterns.add<TosaFoldConstantReciprocal>(ctx, foldSplatOrSingleUseOnly);
   patterns.add<TosaFoldConstantRSQRT>(ctx, foldSplatOrSingleUseOnly);

mlir/test/Dialect/Tosa/constant-reshape-fold.mlir

@@ -0,0 +1,33 @@
+// RUN: mlir-opt --split-input-file --tosa-layerwise-constant-fold %s | FileCheck %s
+// RUN: mlir-opt --split-input-file --tosa-layerwise-constant-fold="fold-splat-or-single-use-only=0" %s | FileCheck %s --check-prefix CHECK-MULTI
+// CHECK-LABEL: @reshape_single_user
+func.func @reshape_single_user() -> tensor<1x2xf32> {
+  // CHECK: %[[RES:.*]] = "tosa.const"{{.*}}-> tensor<1x2xf32>
+  // CHECK: return %[[RES]]
+  %0 = "tosa.const"() {value = dense<4.0> : tensor<2xf32>} : () -> tensor<2xf32>
+  %1 = tosa.reshape %0 {new_shape = array<i64: 1, 2>}: (tensor<2xf32>) -> tensor<1x2xf32>
+  return %1 : tensor<1x2xf32>
+}
+
+// CHECK-LABEL: @reshape_multi_user
+func.func @reshape_multi_user_splat() -> (tensor<1x2xf32>, tensor<2xf32>) {
+  // CHECK-DAG: %[[RES:.*]] = "tosa.const"{{.*}}-> tensor<2xf32>
+  // CHECK-DAG: %[[RESHAPED:.*]] = "tosa.const"{{.*}}-> tensor<1x2xf32>
+  // CHECK: return %[[RESHAPED]], %[[RES]]
+  %0 = "tosa.const"() {value = dense<4.0> : tensor<2xf32>} : () -> tensor<2xf32>
+  %1 = tosa.reshape %0 {new_shape = array<i64: 1, 2>}: (tensor<2xf32>) -> tensor<1x2xf32>
+  return %1, %0 : tensor<1x2xf32>, tensor<2xf32>
+}
+
+// CHECK-LABEL: @reshape_multi_user_non_splat
+func.func @reshape_multi_user_non_splat() -> (tensor<1x2xf32>, tensor<2xf32>) {
+  // CHECK: %[[CONST:.*]] = "tosa.const"{{.*}}-> tensor<2xf32>
+  // CHECK: %[[RES:.*]] = tosa.reshape
+  // CHECK: return %[[RES]], %[[CONST]]
+  // CHECK-MULTI-DAG: %[[RES:.*]] = "tosa.const"{{.*}}-> tensor<2xf32>
+  // CHECK-MULTI-DAG: %[[RESHAPED:.*]] = "tosa.const"{{.*}}-> tensor<1x2xf32>
+  // CHECK-MULTI: return %[[RESHAPED]], %[[RES]]
+  %0 = "tosa.const"() {value = dense<[4.0, 3.0]> : tensor<2xf32>} : () -> tensor<2xf32>
+  %1 = tosa.reshape %0 {new_shape = array<i64: 1, 2>}: (tensor<2xf32>) -> tensor<1x2xf32>
+  return %1, %0 : tensor<1x2xf32>, tensor<2xf32>
+}