[mlir][Canonicalize] Fix command-line options

The canonicalize command-line options currently have no effect, as the pass is
reading the pass options in its constructor, before they're actually
initialized. This results in the default values of the options always being used.

The change here moves the initialization of the `GreedyRewriteConfig` out of the
constructor, so that it runs after the pass options have been parsed.

Fixes #55466

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D125621

2 files changed, 23 insertions, 13 deletions

diff --git a/mlir/lib/Transforms/Canonicalizer.cpp b/mlir/lib/Transforms/Canonicalizer.cpp
index cb532746c448..3f6dbe933b2c 100644
--- a/mlir/lib/Transforms/Canonicalizer.cpp
+++ b/mlir/lib/Transforms/Canonicalizer.cpp
@@ -21,22 +21,17 @@ using namespace mlir;
 namespace {
 /// Canonicalize operations in nested regions.
 struct Canonicalizer : public CanonicalizerBase<Canonicalizer> {
+  Canonicalizer() = default;
   Canonicalizer(const GreedyRewriteConfig &config,
                 ArrayRef<std::string> disabledPatterns,
-                ArrayRef<std::string> enabledPatterns)
-      : config(config) {
+                ArrayRef<std::string> enabledPatterns) {
+    this->topDownProcessingEnabled = config.useTopDownTraversal;
+    this->enableRegionSimplification = config.enableRegionSimplification;
+    this->maxIterations = config.maxIterations;
     this->disabledPatterns = disabledPatterns;
     this->enabledPatterns = enabledPatterns;
   }
 
-  Canonicalizer() {
-    // Default constructed Canonicalizer takes its settings from command line
-    // options.
-    config.useTopDownTraversal = topDownProcessingEnabled;
-    config.enableRegionSimplification = enableRegionSimplification;
-    config.maxIterations = maxIterations;
-  }
-
   /// Initialize the canonicalizer by building the set of patterns used during
   /// execution.
   LogicalResult initialize(MLIRContext *context) override {
@@ -51,11 +46,13 @@ struct Canonicalizer : public CanonicalizerBase<Canonicalizer> {
     return success();
   }
   void runOnOperation() override {
-    (void)applyPatternsAndFoldGreedily(getOperation()->getRegions(), patterns,
-                                       config);
+    GreedyRewriteConfig config;
+    config.useTopDownTraversal = topDownProcessingEnabled;
+    config.enableRegionSimplification = enableRegionSimplification;
+    config.maxIterations = maxIterations;
+    (void)applyPatternsAndFoldGreedily(getOperation(), patterns, config);
   }
 
-  GreedyRewriteConfig config;
   FrozenRewritePatternSet patterns;
 };
 } // namespace
diff --git a/mlir/test/Transforms/test-canonicalize.mlir b/mlir/test/Transforms/test-canonicalize.mlir
index b845293b3520..2181d1856d3a 100644
--- a/mlir/test/Transforms/test-canonicalize.mlir
+++ b/mlir/test/Transforms/test-canonicalize.mlir
@@ -1,4 +1,5 @@
 // RUN: mlir-opt %s -pass-pipeline='func.func(canonicalize)' | FileCheck %s
+// RUN: mlir-opt %s -pass-pipeline='func.func(canonicalize{region-simplify=false})' | FileCheck %s --check-prefixes=CHECK,NO-RS
 
 // CHECK-LABEL: func @remove_op_with_inner_ops_pattern
 func.func @remove_op_with_inner_ops_pattern() {
@@ -89,3 +90,15 @@ func.func @test_dialect_canonicalizer() -> (i32) {
   // CHECK: return %[[CST]]
   return %0 : i32
 }
+
+// Check that the option to control region simplification actually works
+// CHECK-LABEL: test_region_simplify
+func.func @test_region_simplify() {
+  // CHECK-NEXT:   return
+  // NO-RS-NEXT: ^bb1
+  // NO-RS-NEXT:   return
+  // CHECK-NEXT: }
+  return
+^bb1:
+  return
+}