[mlir][FoldUtils] Ensure the created constant dominates the replaced op

This revision fixes a bug where an operation would get replaced with
a pre-existing constant that didn't dominate it. This can occur when
a pattern inserts operations to be folded at the beginning of the
constants insertion block. This revision fixes the bug by moving the
existing constant before the replaced operation in such cases. This is
fine because if a constant didn't already exist, a new one would have
been inserted before this operation anyways.

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

mlir/lib/Transforms/Utils/FoldUtils.cpp

@@ -233,6 +233,14 @@ LogicalResult OperationFolder::tryToFold(
     if (auto *constOp =
             tryGetOrCreateConstant(uniquedConstants, dialect, builder, attrRepl,
                                    res.getType(), op->getLoc())) {
+      // Ensure that this constant dominates the operation we are replacing it
+      // with. This may not automatically happen if the operation being folded
+      // was inserted before the constant within the insertion block.
+      if (constOp->getBlock() == op->getBlock() &&
+          !constOp->isBeforeInBlock(op)) {
+        constOp->moveBefore(op);
+      }
+
       results.push_back(constOp->getResult(0));
       continue;
     }

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

@@ -10,3 +10,15 @@ func @foo() -> i32 {
   %0 = "test.op_in_place_fold_anchor"(%c42) : (i32) -> (i32)
   return %0 : i32
 }
+
+func @test_fold_before_previously_folded_op() -> (i32, i32) {
+  // When folding two constants will be generated and uniqued. Check that the
+  // uniqued constant properly dominates both uses.
+  // CHECK: %[[CST:.+]] = constant true
+  // CHECK-NEXT: "test.cast"(%[[CST]]) : (i1) -> i32
+  // CHECK-NEXT: "test.cast"(%[[CST]]) : (i1) -> i32
+
+  %0 = "test.cast"() {test_fold_before_previously_folded_op} : () -> (i32)
+  %1 = "test.cast"() {test_fold_before_previously_folded_op} : () -> (i32)
+  return %0, %1 : i32, i32
+}

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

@@ -99,6 +99,29 @@ struct FoldingPattern : public RewritePattern {
   }
 };
 
+/// This pattern creates a foldable operation at the entry point of the block.
+/// This tests the situation where the operation folder will need to replace an
+/// operation with a previously created constant that does not initially
+/// dominate the operation to replace.
+struct FolderInsertBeforePreviouslyFoldedConstantPattern
+    : public OpRewritePattern<TestCastOp> {
+public:
+  using OpRewritePattern<TestCastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(TestCastOp op,
+                                PatternRewriter &rewriter) const override {
+    if (!op->hasAttr("test_fold_before_previously_folded_op"))
+      return failure();
+    rewriter.setInsertionPointToStart(op->getBlock());
+
+    auto constOp =
+        rewriter.create<ConstantOp>(op.getLoc(), rewriter.getBoolAttr(true));
+    rewriter.replaceOpWithNewOp<TestCastOp>(op, rewriter.getI32Type(),
+                                            Value(constOp));
+    return success();
+  }
+};
+
 struct TestPatternDriver : public PassWrapper<TestPatternDriver, FunctionPass> {
   StringRef getArgument() const final { return "test-patterns"; }
   StringRef getDescription() const final { return "Run test dialect patterns"; }
@@ -107,7 +130,9 @@ struct TestPatternDriver : public PassWrapper<TestPatternDriver, FunctionPass> {
     populateWithGenerated(patterns);
 
     // Verify named pattern is generated with expected name.
-    patterns.add<FoldingPattern, TestNamedPatternRule>(&getContext());
+    patterns.add<FoldingPattern, TestNamedPatternRule,
+                 FolderInsertBeforePreviouslyFoldedConstantPattern>(
+        &getContext());
 
     (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
   }