refactor(TosaOps): FXML-1981 use hasFolder = 1 for concat folding

Author: gargaroff

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

@@ -1383,6 +1383,7 @@ def Tosa_ConcatOp : Tosa_Op<"concat", [
   );
 
   let hasCanonicalizer = 1;
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//

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

@@ -57,53 +57,9 @@ struct ConcatOptimization : public OpRewritePattern<tosa::ConcatOp> {
   }
 };
 
-struct ConcatFolding : public OpRewritePattern<tosa::ConcatOp> {
-  using OpRewritePattern<tosa::ConcatOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(tosa::ConcatOp op,
-                                PatternRewriter &rewriter) const override {
-    // Fold consecutive concats on the same axis into a single op.
-    uint64_t axis = op.getAxis();
-
-    // Keep track of the operands so we are able to construct a new concat
-    // later. Conservatively assume that we double the number of operands when
-    // folding
-    SmallVector<Value, 8> concatOperands;
-    concatOperands.reserve(2 * op->getNumOperands());
-
-    // Find all operands that are foldable concats
-    bool canFold = false;
-    for (Value operand : op->getOperands()) {
-      concatOperands.emplace_back(operand);
-
-      auto producer = dyn_cast_or_null<ConcatOp>(operand.getDefiningOp());
-      if (!producer)
-        continue;
-
-      // Foldable if axis is the same
-      if (axis != producer.getAxis())
-        continue;
-
-      // Replace the original operand with all incoming operands
-      canFold = true;
-      concatOperands.pop_back();
-      llvm::append_range(concatOperands, producer->getOperands());
-    }
-
-    if (!canFold)
-      return rewriter.notifyMatchFailure(op, "No foldable concats found");
-
-    // Replace the original concat with a new one that contains the original and
-    // folded operands
-    rewriter.replaceOpWithNewOp<tosa::ConcatOp>(op, op->getResultTypes(),
-                                                concatOperands, axis);
-    return success();
-  }
-};
-
 void ConcatOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                            MLIRContext *context) {
-  results.add<ConcatOptimization, ConcatFolding>(context);
+  results.add<ConcatOptimization>(context);
 }
 
 struct ReshapeReshapeOptimization : public OpRewritePattern<tosa::ReshapeOp> {
@@ -1039,3 +995,37 @@ OpFoldResult TransposeOp::fold(ArrayRef<Attribute> operands) {
     return getInput1();
   return {};
 }
+
+OpFoldResult ConcatOp::fold(ArrayRef<Attribute> operands) {
+  // Fold consecutive concats on the same axis into a single op.
+  // Keep track of the operands so we are able to construct a new concat
+  // later. Conservatively assume that we double the number of operands when
+  // folding
+  SmallVector<Value, 8> concatOperands;
+  concatOperands.reserve(2 * getNumOperands());
+
+  // Find all operands that are foldable concats
+  bool canFold = false;
+  for (Value operand : getOperands()) {
+    concatOperands.emplace_back(operand);
+
+    auto producer = dyn_cast_or_null<ConcatOp>(operand.getDefiningOp());
+    if (!producer)
+      continue;
+
+    // Foldable if axis is the same
+    if (getAxis() != producer.getAxis())
+      continue;
+
+    // Replace the original operand with all incoming operands
+    canFold = true;
+    concatOperands.pop_back();
+    llvm::append_range(concatOperands, producer->getOperands());
+  }
+
+  if (!canFold)
+    return {};
+
+  getOperation()->setOperands(concatOperands);
+  return getResult();
+}

mlir/test/Dialect/Tosa/fold_concats.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt --split-input-file --tosa-layerwise-constant-fold %s | FileCheck %s
+// RUN: mlir-opt --split-input-file --canonicalize %s | FileCheck %s
 
 func.func @single_concat(%arg0: tensor<1x1x7x7xf32>) -> tensor<1x2x7x7xf32> {
   %0 = "tosa.concat"(%arg0, %arg0) {axis = 1} : (tensor<1x1x7x7xf32>, tensor<1x1x7x7xf32>) -> tensor<1x2x7x7xf32>