[GreedyPatternRewriter] Avoid reversing constant order

[rkayaith]

Summary

The previous fix from af371f9f98da only applied when using a bottom-up
traversal. The change here applies the constant preprocessing logic to the
top-down case as well. This resolves the issue with the canonicalizer pass still
reordering constants, since it uses a top-down traversal by default.

Fixes #51892

[rkayaith]

use config default as option default

[rriddle]

Nice! Can you look into the flang test failures?

[rkayaith]

fix flang failures

[rkayaith]

Could you land this for me?

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

auto logResult = [&](StringRef result, const llvm::Twine &msg = {}) {
     logger.getOStream() << "\n";
   };
   auto logResultWithLine = [&](StringRef result, const llvm::Twine &msg = {}) {
     logResult(result, msg);
     logger.startLine() << logLineComment;
   };
 #endif
 
+  auto insertKnownConstant = [&](Operation *op) {
+    // Check for existing constants when populating the worklist. This avoids
+    // accidentally reversing the constant order during processing.
+    Attribute constValue;
+    if (matchPattern(op, m_Constant(&constValue)))
+      if (!folder.insertKnownConstant(op, constValue))
+        return true;
+    return false;
+  };
+
   bool changed = false;
   unsigned iteration = 0;
   do {
     worklist.clear();
     worklistMap.clear();
 
     if (!config.useTopDownTraversal) {
       // Add operations to the worklist in postorder.
       for (auto &region : regions) {
-        region.walk([this](Operation *op) {
-          // If we aren't processing top-down, check for existing constants when
-          // populating the worklist. This avoids accidentally reversing the
-          // constant order during processing.
-          Attribute constValue;
-          if (matchPattern(op, m_Constant(&constValue)))
-            if (!folder.insertKnownConstant(op, constValue))
-              return;
+        region.walk([&](Operation *op) {
+          if (!insertKnownConstant(op))
             addToWorklist(op);
         });
       }
     } else {
       // Add all nested operations to the worklist in preorder.
       for (auto &region : regions)
-        region.walk<WalkOrder::PreOrder>(
-            [this](Operation *op) { worklist.push_back(op); });
+        region.walk<WalkOrder::PreOrder>([&](Operation *op) {
+          if (!insertKnownConstant(op))
+            worklist.push_back(op);
+        });
 
       // Reverse the list so our pop-back loop processes them in-order.
       std::reverse(worklist.begin(), worklist.end());
       // Remember the reverse index.
       for (size_t i = 0, e = worklist.size(); i != e; ++i)
         worklistMap[worklist[i]] = i;
     }
 

mlir/test/Dialect/Arithmetic/canonicalize.mlir

func.func @bitcastOfBitcast(%arg : i16) -> i16 {
   %bf = arith.bitcast %f : f16 to bf16
   %res = arith.bitcast %bf : bf16 to i16
   return %res : i16
 }
 
 // -----
 
 // CHECK-LABEL: test_maxsi
-// CHECK: %[[C0:.+]] = arith.constant 42
-// CHECK: %[[MAX_INT_CST:.+]] = arith.constant 127
+// CHECK-DAG: %[[C0:.+]] = arith.constant 42
+// CHECK-DAG: %[[MAX_INT_CST:.+]] = arith.constant 127
 // CHECK: %[[X:.+]] = arith.maxsi %arg0, %[[C0]]
 // CHECK: return %arg0, %[[MAX_INT_CST]], %arg0, %[[X]]
 func.func @test_maxsi(%arg0 : i8) -> (i8, i8, i8, i8) {
   %maxIntCst = arith.constant 127 : i8
   %minIntCst = arith.constant -128 : i8
   %c0 = arith.constant 42 : i8
   %0 = arith.maxsi %arg0, %arg0 : i8
   %1 = arith.maxsi %arg0, %maxIntCst : i8
   %2 = arith.maxsi %arg0, %minIntCst : i8
   %3 = arith.maxsi %arg0, %c0 : i8
   return %0, %1, %2, %3: i8, i8, i8, i8
 }
 
 // CHECK-LABEL: test_maxsi2
-// CHECK: %[[C0:.+]] = arith.constant 42
-// CHECK: %[[MAX_INT_CST:.+]] = arith.constant 127
+// CHECK-DAG: %[[C0:.+]] = arith.constant 42
+// CHECK-DAG: %[[MAX_INT_CST:.+]] = arith.constant 127
 // CHECK: %[[X:.+]] = arith.maxsi %arg0, %[[C0]]
 // CHECK: return %arg0, %[[MAX_INT_CST]], %arg0, %[[X]]
 func.func @test_maxsi2(%arg0 : i8) -> (i8, i8, i8, i8) {
   %maxIntCst = arith.constant 127 : i8
   %minIntCst = arith.constant -128 : i8
   %c0 = arith.constant 42 : i8
   %0 = arith.maxsi %arg0, %arg0 : i8
   %1 = arith.maxsi %maxIntCst, %arg0: i8
   %2 = arith.maxsi %minIntCst, %arg0: i8
   %3 = arith.maxsi %c0, %arg0 : i8
   return %0, %1, %2, %3: i8, i8, i8, i8
 }
 
 // -----
 
 // CHECK-LABEL: test_maxui
-// CHECK: %[[C0:.+]] = arith.constant 42
-// CHECK: %[[MAX_INT_CST:.+]] = arith.constant -1
+// CHECK-DAG: %[[C0:.+]] = arith.constant 42
+// CHECK-DAG: %[[MAX_INT_CST:.+]] = arith.constant -1
 // CHECK: %[[X:.+]] = arith.maxui %arg0, %[[C0]]
 // CHECK: return %arg0, %[[MAX_INT_CST]], %arg0, %[[X]]
 func.func @test_maxui(%arg0 : i8) -> (i8, i8, i8, i8) {
   %maxIntCst = arith.constant 255 : i8
   %minIntCst = arith.constant 0 : i8
   %c0 = arith.constant 42 : i8
   %0 = arith.maxui %arg0, %arg0 : i8
   %1 = arith.maxui %arg0, %maxIntCst : i8
   %2 = arith.maxui %arg0, %minIntCst : i8
   %3 = arith.maxui %arg0, %c0 : i8
   return %0, %1, %2, %3: i8, i8, i8, i8
 }
 
 // CHECK-LABEL: test_maxui
-// CHECK: %[[C0:.+]] = arith.constant 42
-// CHECK: %[[MAX_INT_CST:.+]] = arith.constant -1
+// CHECK-DAG: %[[C0:.+]] = arith.constant 42
+// CHECK-DAG: %[[MAX_INT_CST:.+]] = arith.constant -1
 // CHECK: %[[X:.+]] = arith.maxui %arg0, %[[C0]]
 // CHECK: return %arg0, %[[MAX_INT_CST]], %arg0, %[[X]]
 func.func @test_maxui2(%arg0 : i8) -> (i8, i8, i8, i8) {
   %maxIntCst = arith.constant 255 : i8
   %minIntCst = arith.constant 0 : i8
   %c0 = arith.constant 42 : i8
   %0 = arith.maxui %arg0, %arg0 : i8
   %1 = arith.maxui %maxIntCst, %arg0 : i8
   %2 = arith.maxui %minIntCst, %arg0 : i8
   %3 = arith.maxui %c0, %arg0 : i8
   return %0, %1, %2, %3: i8, i8, i8, i8
 }
 
 // -----
 
 // CHECK-LABEL: test_minsi
-// CHECK: %[[C0:.+]] = arith.constant 42
-// CHECK: %[[MIN_INT_CST:.+]] = arith.constant -128
+// CHECK-DAG: %[[C0:.+]] = arith.constant 42
+// CHECK-DAG: %[[MIN_INT_CST:.+]] = arith.constant -128
 // CHECK: %[[X:.+]] = arith.minsi %arg0, %[[C0]]
 // CHECK: return %arg0, %arg0, %[[MIN_INT_CST]], %[[X]]
 func.func @test_minsi(%arg0 : i8) -> (i8, i8, i8, i8) {
   %maxIntCst = arith.constant 127 : i8
   %minIntCst = arith.constant -128 : i8
   %c0 = arith.constant 42 : i8
   %0 = arith.minsi %arg0, %arg0 : i8
   %1 = arith.minsi %arg0, %maxIntCst : i8
   %2 = arith.minsi %arg0, %minIntCst : i8
   %3 = arith.minsi %arg0, %c0 : i8
   return %0, %1, %2, %3: i8, i8, i8, i8
 }
 
 // CHECK-LABEL: test_minsi
-// CHECK: %[[C0:.+]] = arith.constant 42
-// CHECK: %[[MIN_INT_CST:.+]] = arith.constant -128
+// CHECK-DAG: %[[C0:.+]] = arith.constant 42
+// CHECK-DAG: %[[MIN_INT_CST:.+]] = arith.constant -128
 // CHECK: %[[X:.+]] = arith.minsi %arg0, %[[C0]]
 // CHECK: return %arg0, %arg0, %[[MIN_INT_CST]], %[[X]]
 func.func @test_minsi2(%arg0 : i8) -> (i8, i8, i8, i8) {
   %maxIntCst = arith.constant 127 : i8
   %minIntCst = arith.constant -128 : i8
   %c0 = arith.constant 42 : i8
   %0 = arith.minsi %arg0, %arg0 : i8
   %1 = arith.minsi %maxIntCst, %arg0 : i8
   %2 = arith.minsi %minIntCst, %arg0 : i8
   %3 = arith.minsi %c0, %arg0 : i8
   return %0, %1, %2, %3: i8, i8, i8, i8
 }
 
 // -----
 
 // CHECK-LABEL: test_minui
-// CHECK: %[[C0:.+]] = arith.constant 42
-// CHECK: %[[MIN_INT_CST:.+]] = arith.constant 0
+// CHECK-DAG: %[[C0:.+]] = arith.constant 42
+// CHECK-DAG: %[[MIN_INT_CST:.+]] = arith.constant 0
 // CHECK: %[[X:.+]] = arith.minui %arg0, %[[C0]]
 // CHECK: return %arg0, %arg0, %[[MIN_INT_CST]], %[[X]]
 func.func @test_minui(%arg0 : i8) -> (i8, i8, i8, i8) {
   %maxIntCst = arith.constant 255 : i8
   %minIntCst = arith.constant 0 : i8
   %c0 = arith.constant 42 : i8
   %0 = arith.minui %arg0, %arg0 : i8
   %1 = arith.minui %arg0, %maxIntCst : i8
   %2 = arith.minui %arg0, %minIntCst : i8
   %3 = arith.minui %arg0, %c0 : i8
   return %0, %1, %2, %3: i8, i8, i8, i8
 }
 
 // CHECK-LABEL: test_minui
-// CHECK: %[[C0:.+]] = arith.constant 42
-// CHECK: %[[MIN_INT_CST:.+]] = arith.constant 0
+// CHECK-DAG: %[[C0:.+]] = arith.constant 42
+// CHECK-DAG: %[[MIN_INT_CST:.+]] = arith.constant 0
 // CHECK: %[[X:.+]] = arith.minui %arg0, %[[C0]]
 // CHECK: return %arg0, %arg0, %[[MIN_INT_CST]], %[[X]]
 func.func @test_minui2(%arg0 : i8) -> (i8, i8, i8, i8) {
   %maxIntCst = arith.constant 255 : i8
   %minIntCst = arith.constant 0 : i8
   %c0 = arith.constant 42 : i8
   %0 = arith.minui %arg0, %arg0 : i8
   %1 = arith.minui %maxIntCst, %arg0 : i8

mlir/test/Dialect/SCF/canonicalize.mlir

func.func @invariant_loop_args_in_same_order(%f_arg0: tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) {
   ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>, %arg4: tensor<i32>): // no predecessors
     // %arg1 here will get replaced by %cst_1
     %1 = arith.addi %arg0, %arg1 : tensor<i32>
     %2 = arith.addi %arg2, %arg3 : tensor<i32>
     scf.yield %1, %arg1, %2, %2, %arg4 : tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>
   }
   return %0#0, %0#1, %0#2, %0#3, %0#4 : tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>
 }
-// CHECK:    %[[CST42:.*]] = arith.constant dense<42>
-// CHECK:    %[[ONE:.*]] = arith.constant dense<1>
 // CHECK:    %[[ZERO:.*]] = arith.constant dense<0>
+// CHECK:    %[[ONE:.*]] = arith.constant dense<1>
+// CHECK:    %[[CST42:.*]] = arith.constant dense<42>
 // CHECK:    %[[WHILE:.*]]:3 = scf.while (%[[ARG0:.*]] = %[[ZERO]], %[[ARG2:.*]] = %[[ONE]], %[[ARG3:.*]] = %[[ONE]])
 // CHECK:       arith.cmpi slt, %[[ARG0]], %{{.*}}
 // CHECK:       tensor.extract %{{.*}}[]
 // CHECK:       scf.condition(%{{.*}}) %[[ARG0]], %[[ARG2]], %[[ARG3]]
 // CHECK:    } do {
 // CHECK:     ^{{.*}}(%[[ARG0:.*]]: tensor<i32>, %[[ARG2:.*]]: tensor<i32>, %[[ARG3:.*]]: tensor<i32>):
 // CHECK:       %[[VAL0:.*]] = arith.addi %[[ARG0]], %[[FUNC_ARG0]]
 // CHECK:       %[[VAL1:.*]] = arith.addi %[[ARG2]], %[[ARG3]]
func.func @while_loop_invariant_argument_different_order() -> (tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) {
   } do {
   ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<i32>, %arg4: tensor<i32>, %arg5: tensor<i32>): // no predecessors
     %1 = arith.addi %arg0, %cst_1 : tensor<i32>
     %2 = arith.addi %arg2, %arg3 : tensor<i32>
     scf.yield %arg3, %arg1, %2, %2, %arg4 : tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>
   }
   return %0#0, %0#1, %0#2, %0#3, %0#4, %0#5 : tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>
 }
-// CHECK:    %[[CST42:.*]] = arith.constant dense<42>
-// CHECK:    %[[ONE:.*]] = arith.constant dense<1>
 // CHECK:    %[[ZERO:.*]] = arith.constant dense<0>
+// CHECK:    %[[ONE:.*]] = arith.constant dense<1>
+// CHECK:    %[[CST42:.*]] = arith.constant dense<42>
 // CHECK:    %[[WHILE:.*]]:2 = scf.while (%[[ARG1:.*]] = %[[ONE]], %[[ARG4:.*]] = %[[ZERO]])
 // CHECK:       arith.cmpi slt, %[[ZERO]], %[[CST42]]
 // CHECK:       tensor.extract %{{.*}}[]
 // CHECK:       scf.condition(%{{.*}}) %[[ARG1]], %[[ARG4]]
 // CHECK:    } do {
 // CHECK:     ^{{.*}}(%{{.*}}: tensor<i32>, %{{.*}}: tensor<i32>):
 // CHECK:       scf.yield %[[ZERO]], %[[ONE]]
 // CHECK:    }

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

-// RUN: mlir-opt -test-patterns -test-patterns %s | FileCheck %s
+// RUN: mlir-opt -test-patterns='top-down=false' %s | FileCheck %s
+// RUN: mlir-opt -test-patterns='top-down=true' %s | FileCheck %s
 
 func.func @foo() -> i32 {
   %c42 = arith.constant 42 : i32
 
   // The new operation should be present in the output and contain an attribute
   // with value "42" that results from folding.
 
   // CHECK: "test.op_in_place_fold"(%{{.*}}) {attr = 42 : i32}

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

LogicalResult matchAndRewrite(TestCommutative2Op op,
     return success();
   }
 };
 
 struct TestPatternDriver
     : public PassWrapper<TestPatternDriver, OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPatternDriver)
 
+  TestPatternDriver() = default;
+  TestPatternDriver(const TestPatternDriver &other) : PassWrapper(other) {}
+
   StringRef getArgument() const final { return "test-patterns"; }
   StringRef getDescription() const final { return "Run test dialect patterns"; }
   void runOnOperation() override {
     mlir::RewritePatternSet patterns(&getContext());
     populateWithGenerated(patterns);
 
     // Verify named pattern is generated with expected name.
     patterns.add<FoldingPattern, TestNamedPatternRule,
                  FolderInsertBeforePreviouslyFoldedConstantPattern,
                  FolderCommutativeOp2WithConstant>(&getContext());
 
-    (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
+    GreedyRewriteConfig config;
+    config.useTopDownTraversal = this->useTopDownTraversal;
+    (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns),
+                                       config);
   }
+
+  Option<bool> useTopDownTraversal{
+      *this, "top-down",
+      llvm::cl::desc("Seed the worklist in general top-down order"),
+      llvm::cl::init(GreedyRewriteConfig().useTopDownTraversal)};
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // ReturnType Driver.
 //===----------------------------------------------------------------------===//
 
 namespace {