Add tests & fix call site

Author: jpienaar

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

@@ -170,9 +170,8 @@ applyPatternsGreedily(Operation *op, const FrozenRewritePatternSet &patterns,
   bool failed = false;
   for (Region &region : op->getRegions()) {
     bool regionChanged;
-    failed |=
-        applyPatternsAndFoldGreedily(region, patterns, config, &regionChanged)
-            .failed();
+    failed |= applyPatternsGreedily(region, patterns, config, &regionChanged)
+                  .failed();
     anyRegionChanged |= regionChanged;
   }
   if (changed)

mlir/test/Transforms/test-operation-folder.mlir

@@ -1,5 +1,7 @@
 // RUN: mlir-opt -test-greedy-patterns='top-down=false' %s | FileCheck %s
 // RUN: mlir-opt -test-greedy-patterns='top-down=true' %s | FileCheck %s
+// RUN: mlir-opt -test-greedy-patterns='cse-constants=false' %s | FileCheck %s --check-prefix=NOCSE
+// RUN: mlir-opt -test-greedy-patterns='fold=false' %s | FileCheck %s --check-prefix=NOFOLD
 
 func.func @foo() -> i32 {
   %c42 = arith.constant 42 : i32
@@ -25,7 +27,8 @@ func.func @test_fold_before_previously_folded_op() -> (i32, i32) {
 }
 
 func.func @test_dont_reorder_constants() -> (i32, i32, i32) {
-  // Test that we don't reorder existing constants during folding if it isn't necessary.
+  // Test that we don't reorder existing constants during folding if it isn't
+  // necessary.
   // CHECK: %[[CST:.+]] = arith.constant 1
   // CHECK-NEXT: %[[CST:.+]] = arith.constant 2
   // CHECK-NEXT: %[[CST:.+]] = arith.constant 3
@@ -34,3 +37,37 @@ func.func @test_dont_reorder_constants() -> (i32, i32, i32) {
   %2 = arith.constant 3 : i32
   return %0, %1, %2 : i32, i32, i32
 }
+
+func.func @test_dont_fold() -> (i32, i32, i32, i32, i32, i32) {
+  // Test either not folding or deduping constants.
+
+  // CHECK-LABEL: test_dont_fold
+  // CHECK-NOT: arith.constant 0
+  // CHECK-DAG: %[[CST:.+]] = arith.constant 0
+  // CHECK-DAG: %[[CST:.+]] = arith.constant 1
+  // CHECK-DAG: %[[CST:.+]] = arith.constant 2
+  // CHECK-DAG: %[[CST:.+]] = arith.constant 3
+  // CHECK-NEXT: return
+
+  // NOCSE-LABEL: test_dont_fold
+  // NOCSE-DAG: arith.constant 0 : i32
+  // NOCSE-DAG: arith.constant 1 : i32
+  // NOCSE-DAG: arith.constant 2 : i32
+  // NOCSE-DAG: arith.constant 1 : i32
+  // NOCSE-DAG: arith.constant 2 : i32
+  // NOCSE-DAG: arith.constant 3 : i32
+  // NOCSE-NEXT: return
+
+  // NOFOLD-LABEL: test_dont_fold
+  // NOFOLD: arith.addi
+  // NOFOLD: arith.addi
+  // NOFOLD: arith.addi
+
+  %c0 = arith.constant 0 : i32
+  %c1 = arith.constant 1 : i32
+  %c2 = arith.constant 2 : i32
+  %0 = arith.addi %c0, %c1 : i32
+  %1 = arith.addi %0, %c1 : i32
+  %2 = arith.addi %c2, %c1 : i32
+  return %0, %1, %2, %c0, %c1, %c2 : i32, i32, i32, i32, i32, i32
+}

mlir/test/lib/Dialect/Test/TestPatterns.cpp

@@ -388,8 +388,9 @@ struct TestGreedyPatternDriver
     GreedyRewriteConfig config;
     config.useTopDownTraversal = this->useTopDownTraversal;
     config.maxIterations = this->maxIterations;
-    (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns),
-                                       config);
+    config.fold = this->fold;
+    config.cseConstants = this->cseConstants;
+    (void)applyPatternsGreedily(getOperation(), std::move(patterns), config);
   }
 
   Option<bool> useTopDownTraversal{
@@ -400,6 +401,11 @@ struct TestGreedyPatternDriver
       *this, "max-iterations",
       llvm::cl::desc("Max. iterations in the GreedyRewriteConfig"),
       llvm::cl::init(GreedyRewriteConfig().maxIterations)};
+  Option<bool> fold{*this, "fold", llvm::cl::desc("Whether to fold"),
+                    llvm::cl::init(GreedyRewriteConfig().fold)};
+  Option<bool> cseConstants{*this, "cse-constants",
+                            llvm::cl::desc("Whether to CSE constants"),
+                            llvm::cl::init(GreedyRewriteConfig().cseConstants)};
 };
 
 struct DumpNotifications : public RewriterBase::Listener {