Iterate on the operations to fold in TestConstantFold in reverse to remove the need for ConstantFoldHelper to have a flag for insertion at the head of the entry block. This also fixes an asan bug in TestConstantFold due to the iteration order of operations and ConstantFoldHelper's constant insertion placement.

    Note: This now means that we cannot fold chains of operations, i.e. where constant foldable operations feed into each other. Given that this is a testing pass solely for constant folding, this isn't really something that we want anyways. Constant fold tests should be simple and direct, with more advanced folding/feeding being tested with the canonicalizer.

--

PiperOrigin-RevId: 242011744

2 files changed, 13 insertions, 17 deletions

diff --git a/mlir/lib/Transforms/TestConstantFold.cpp b/mlir/lib/Transforms/TestConstantFold.cpp
index 60407cd49d1..0990d7a73f6 100644
--- a/mlir/lib/Transforms/TestConstantFold.cpp
+++ b/mlir/lib/Transforms/TestConstantFold.cpp
@@ -61,10 +61,18 @@ void TestConstantFold::runOnFunction() {
   opsToErase.clear();
 
   auto &f = getFunction();
+  ConstantFoldHelper helper(&f);
 
-  ConstantFoldHelper helper(&f, /*insertAtHead=*/false);
+  // Collect and fold the operations within the function.
+  SmallVector<Operation *, 8> ops;
+  f.walk([&](Operation *op) { ops.push_back(op); });
 
-  f.walk([&](Operation *op) { foldOperation(op, helper); });
+  // Fold the constants in reverse so that the last generated constants from
+  // folding are at the beginning. This creates somewhat of a linear ordering to
+  // the newly generated constants that matches the operation order and improves
+  // the readability of test cases.
+  for (Operation *op : llvm::reverse(ops))
+    foldOperation(op, helper);
 
   // At this point, these operations are dead, remove them.
   for (auto *op : opsToErase) {
diff --git a/mlir/lib/Transforms/Utils/ConstantFoldUtils.cpp b/mlir/lib/Transforms/Utils/ConstantFoldUtils.cpp
index 5908ec251d0..fc8209be872 100644
--- a/mlir/lib/Transforms/Utils/ConstantFoldUtils.cpp
+++ b/mlir/lib/Transforms/Utils/ConstantFoldUtils.cpp
@@ -29,8 +29,7 @@
 
 using namespace mlir;
 
-ConstantFoldHelper::ConstantFoldHelper(Function *f, bool insertAtHead)
-    : function(f), isInsertAtHead(insertAtHead) {}
+ConstantFoldHelper::ConstantFoldHelper(Function *f) : function(f) {}
 
 bool ConstantFoldHelper::tryToConstantFold(
     Operation *op, std::function<void(Operation *)> preReplaceAction) {
@@ -146,18 +145,7 @@ bool ConstantFoldHelper::tryToUnify(Operation *op) {
 }
 
 void ConstantFoldHelper::moveConstantToEntryBlock(Operation *op) {
+  // Insert at the very top of the entry block.
   auto &entryBB = function->front();
-  if (isInsertAtHead || entryBB.empty()) {
-    // Insert at the very top of the entry block.
-    op->moveBefore(&entryBB, entryBB.begin());
-  } else {
-    // TODO: This is only used by TestConstantFold and not very clean. We should
-    // figure out a better way to work around this.
-
-    // Move to be ahead of the first non-constant op.
-    auto it = entryBB.begin();
-    while (it != entryBB.end() && it->isa<ConstantOp>())
-      ++it;
-    op->moveBefore(&entryBB, it);
-  }
+  op->moveBefore(&entryBB, entryBB.begin());
 }