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

[rriddle]

Summary

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.

[bondhugula]

Nice catch and fix!

constant before the replaced operation in such cases, something
that would happen anyways if the constant didn't already exist.

This clause isn't clear to me. Is there a typo?

mlir/lib/Transforms/Utils/FoldUtils.cpp

237

This may happen -> This may not automatically happen ... or a slight rephrasing?

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

102

Short comment here.

[rriddle]

update

[rriddle]

In D108498#2958620, @bondhugula wrote:

Nice catch and fix!

constant before the replaced operation in such cases, something
that would happen anyways if the constant didn't already exist.

This clause isn't clear to me. Is there a typo?

Reworded a bit.

[rriddle]

Thanks for the review Uday! Appreciated.

[lattner]

FYI, this bug introduced a major performance regression in canonicalize:
https://bugs.llvm.org/show_bug.cgi?id=51738

-Chris

mlir/lib/Transforms/Utils/FoldUtils.cpp

for (unsigned i = 0, e = op->getNumResults(); i != e; ++i) {
     }
 
     // Check to see if there is a canonicalized version of this constant.
     auto res = op->getResult(i);
     Attribute attrRepl = foldResults[i].get<Attribute>();
     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;
     }
     // If materialization fails, cleanup any operations generated for the
     // previous results and return failure.
     for (Operation &op : llvm::make_early_inc_range(
              llvm::make_range(entry.begin(), builder.getInsertionPoint()))) {
       notifyRemoval(&op);

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

 // RUN: mlir-opt -test-patterns %s | FileCheck %s
 
 func @foo() -> i32 {
   %c42 = 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}
   %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

Value result = folder.create<TestOpInPlaceFold>(
         rewriter, op->getLoc(), rewriter.getIntegerType(32), op->getOperand(0),
         rewriter.getI32IntegerAttr(0));
     assert(result);
     rewriter.replaceOp(op, result);
     return success();
   }
 };
 
+/// 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"; }
   void runOnFunction() override {
     mlir::RewritePatternSet patterns(&getContext());
     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));
   }
 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // ReturnType Driver.