From e4635e6328c88f1f403b316bd8340fe03f9835a8 Mon Sep 17 00:00:00 2001
From: River Riddle <riddleriver@gmail.com>
Date: Mon, 23 Aug 2021 18:14:18 +0000
Subject: [PATCH] [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         |  8 ++++++++
 mlir/test/Transforms/test-operation-folder.mlir | 12 +++++++++++
 mlir/test/lib/Dialect/Test/TestPatterns.cpp     | 27 ++++++++++++++++++++++++-
 3 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Transforms/Utils/FoldUtils.cpp b/mlir/lib/Transforms/Utils/FoldUtils.cpp
index af415d5..f9e1bb4 100644
--- a/mlir/lib/Transforms/Utils/FoldUtils.cpp
+++ b/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;
     }
diff --git a/mlir/test/Transforms/test-operation-folder.mlir b/mlir/test/Transforms/test-operation-folder.mlir
index 7ef6530..f5e86ed 100644
--- a/mlir/test/Transforms/test-operation-folder.mlir
+++ b/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
+}
diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
index 62bed7e..de141dc 100644
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -99,6 +99,29 @@ public:
   }
 };
 
+/// 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));
   }
-- 
2.7.4