[mlir][transform] Add `apply_cse` option to `transform.apply_patterns` op

[springerm]

Summary

Applying the canonicalizer and CSE in an interleaved fashion is useful after bufferization (and maybe other transforms) to fold away self copies.

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

let description = [{
     targeted op until a fixpoint was reached. Patterns are not applied to the
     targeted op itself.
 
     The patterns that should be applied are specified in the graph region of
     this op. They must implement the `PatternDescriptorOpInterface`. The order
     in which patterns are applied is unspecified; i.e., the ordering of ops in
     the region of this op is irrelevant.
 
+    If `apple_cse` is set, the greedy pattern rewrite is interleaved with
+    common subexpression elimination (CSE): both are repeated until a fixpoint
+    is reached.
+
     This transform only reads the target handle and modifies the payload. If a
     pattern erases or replaces a tracked op, the mapping is updated accordingly.
 
     Only replacements via `RewriterBase::replaceOp` or `replaceOpWithNewOp` are
     considered "payload op replacements". Furthermore, only if the replacement
     values are defined by the same op and that op has the same type as the
     original op, the mapping is updated. Otherwise, this transform fails
     silently. More details can be found at the documentation site of
     `TrackingListener`.
 
     This transform also fails silently if the pattern application did not
     converge within the default number of iterations/rewrites of the greedy
     pattern rewrite driver.
   }];
 
   let arguments = (ins
-    TransformHandleTypeInterface:$target);
+    TransformHandleTypeInterface:$target, UnitAttr:$apply_cse);
   let results = (outs);
   let regions = (region MaxSizedRegion<1>:$patterns);
 
   let assemblyFormat = "`to` $target $patterns attr-dict `:` type($target)";
   let hasVerifier = 1;
 
   let skipDefaultBuilders = 1;
   let builders = [

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

for (Operation &op : getRegion().front()) {
           patterns);
     }
   }
 
   // Configure the GreedyPatternRewriteDriver.
   GreedyRewriteConfig config;
   config.listener =
       static_cast<RewriterBase::Listener *>(rewriter.getListener());
+  FrozenRewritePatternSet frozenPatterns(std::move(patterns));
 
+  // Apply patterns and CSE repetitively until a fixpoint is reached. If no CSE
+  // was requested, apply the greedy pattern rewrite only once. (The greedy
+  // pattern rewrite driver already iterates to a fixpoint internally.)
+  bool cseChanged = false;
+  // One or two iterations should be sufficient. Stop iterating after a certain
+  // threshold to make debugging easier.
+  static const int64_t kNumMaxIterations = 50;
+  int64_t iteration = 0;
+  do {
     LogicalResult result = failure();
     if (target->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
       // Op is isolated from above. Apply patterns and also perform region
       // simplification.
-    result = applyPatternsAndFoldGreedily(target, std::move(patterns), config);
+      result = applyPatternsAndFoldGreedily(target, frozenPatterns, config);
     } else {
-    // Manually gather list of ops because the other GreedyPatternRewriteDriver
-    // overloads only accepts ops that are isolated from above. This way,
-    // patterns can be applied to ops that are not isolated from above. Regions
-    // are not being simplified. Furthermore, only a single greedy rewrite
-    // iteration is performed.
+      // Manually gather list of ops because the other
+      // GreedyPatternRewriteDriver overloads only accepts ops that are isolated
+      // from above. This way, patterns can be applied to ops that are not
+      // isolated from above. Regions are not being simplified. Furthermore,
+      // only a single greedy rewrite iteration is performed.
       SmallVector<Operation *> ops;
       target->walk([&](Operation *nestedOp) {
         if (target != nestedOp)
           ops.push_back(nestedOp);
       });
-    result = applyOpPatternsAndFold(ops, std::move(patterns), config);
+      result = applyOpPatternsAndFold(ops, frozenPatterns, config);
     }
 
     // A failure typically indicates that the pattern application did not
     // converge.
     if (failed(result)) {
       return emitSilenceableFailure(target)
              << "greedy pattern application failed";
     }
 
+    if (getApplyCse()) {
+      DominanceInfo domInfo;
+      mlir::eliminateCommonSubExpressions(rewriter, domInfo, target,
+                                          &cseChanged);
+    }
+  } while (cseChanged && ++iteration < kNumMaxIterations);
+
+  if (iteration == kNumMaxIterations)
+    return emitDefiniteFailure() << "fixpoint iteration did not converge";
+
   return DiagnosedSilenceableFailure::success();
 }
 
 LogicalResult transform::ApplyPatternsOp::verify() {
   if (!getRegion().empty()) {
     for (Operation &op : getRegion().front()) {
       if (!isa<transform::PatternDescriptorOpInterface>(&op)) {
         InFlightDiagnostic diag = emitOpError()

mlir/test/Dialect/Transform/test-pattern-application.mlir

module {
     ^bb1(%arg1: !transform.any_op):
       // expected-error @below {{cannot apply transform to itself (or one of its ancestors)}}
       transform.apply_patterns to %arg1 {
         transform.apply_patterns.canonicalization
       } : !transform.any_op
     }
   }
 }
+
+// -----
+
+// CHECK-LABEL: func @canonicalization_and_cse(
+//   CHECK-NOT:   memref.subview
+//   CHECK-NOT:   memref.copy
+func.func @canonicalization_and_cse(%m: memref<5xf32>) {
+  %c2 = arith.constant 2 : index
+  %s0 = memref.subview %m[1] [2] [1] : memref<5xf32> to memref<2xf32, strided<[1], offset: 1>>
+  %s1 = memref.subview %m[1] [%c2] [1] : memref<5xf32> to memref<?xf32, strided<[1], offset: 1>>
+  memref.copy %s0, %s1 : memref<2xf32, strided<[1], offset: 1>> to memref<?xf32, strided<[1], offset: 1>>
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %1 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  transform.apply_patterns to %1 {
+    transform.apply_patterns.canonicalization
+  } {apply_cse} : !transform.any_op
+}