[mlir] GreedyPatternRewriter: Reprocess modified ops

[springerm]

Summary

Ops that were modifed in-place (finalizeRootUpdate was called) should be reprocessed by the GreedyPatternRewriter. This is currently not happening with GreedyRewriteConfig::maxIterations = 1.

Note: If your project goes into an infinite loop because of this change, you likely have one or multiple faulty patterns that modify the same operations in-place (updateRootInplace) indefinitely.

[springerm]

(There is still a problem with this change. It is crashing vector-transfer-full-partial-split.mlir.)

[mehdi_amini]

Is this something we can test for?

[rriddle]

The change makes sense to me, but it'd be nice to add a test if we can.

[springerm]

address comments

[springerm]

typo

[rriddle]

mlir/unittests/Transforms/GreedyPatternRewriteDriver.cpp

19–31

It would be nice if we could just have something on one of the existing test passes, instead of using a unittest. We generally don't use unittests to test IR/pattern/transformation/etc. changes.

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

[springerm]

address comments

[springerm]

update

[rriddle]

LGTM, but the windows build failure there looks real. Please resolve that before landing.

[springerm]

fix Windows build

mlir/lib/Dialect/Vector/Transforms/VectorTransferSplitRewritePatterns.cpp

LogicalResult mlir::vector::splitFullAndPartialTransfer(
     RewriterBase &b, VectorTransferOpInterface xferOp,
     VectorTransformsOptions options, scf::IfOp *ifOp) {
   if (options.vectorTransferSplit == VectorTransferSplit::None)
     return failure();
 
   SmallVector<bool, 4> bools(xferOp.getTransferRank(), true);
   auto inBoundsAttr = b.getBoolArrayAttr(bools);
   if (options.vectorTransferSplit == VectorTransferSplit::ForceInBounds) {
+    b.updateRootInPlace(xferOp, [&]() {
       xferOp->setAttr(xferOp.getInBoundsAttrStrName(), inBoundsAttr);
+    });
     return success();
   }
 
   // Assert preconditions. Additionally, keep the variables in an inner scope to
   // ensure they aren't used in the wrong scopes further down.
   {
     assert(succeeded(splitFullAndPartialTransferPrecondition(xferOp)) &&
            "Expected splitFullAndPartialTransferPrecondition to hold");
scf::IfOp fullPartialIfOp =
                                           alloc);
     if (ifOp)
       *ifOp = fullPartialIfOp;
 
     // Set existing read op to in-bounds, it always reads from a full buffer.
     for (unsigned i = 0, e = returnTypes.size(); i != e; ++i)
       xferReadOp.setOperand(i, fullPartialIfOp.getResult(i));
 
+    b.updateRootInPlace(xferOp, [&]() {
       xferOp->setAttr(xferOp.getInBoundsAttrStrName(), inBoundsAttr);
+    });
 
     return success();
   }
 
   auto xferWriteOp = cast<vector::TransferWriteOp>(xferOp.getOperation());
 
   // Decide which location to write the entire vector to.
   auto memrefAndIndices = getLocationToWriteFullVec(
LogicalResult mlir::vector::splitFullAndPartialTransfer(
 
   // Create a potential copy from the allocated buffer to the final output in
   // the slow path case.
   if (options.vectorTransferSplit == VectorTransferSplit::VectorTransfer)
     createFullPartialVectorTransferWrite(b, xferWriteOp, inBoundsCond, alloc);
   else
     createFullPartialLinalgCopy(b, xferWriteOp, inBoundsCond, alloc);
 
-  xferOp->erase();
+  b.eraseOp(xferOp);
 
   return success();
 }
 
 LogicalResult mlir::vector::VectorTransferFullPartialRewriter::matchAndRewrite(
     Operation *op, PatternRewriter &rewriter) const {
   auto xferOp = dyn_cast<VectorTransferOpInterface>(op);
   if (!xferOp || failed(splitFullAndPartialTransferPrecondition(xferOp)) ||
       failed(filter(xferOp)))
     return failure();
-  rewriter.startRootUpdate(xferOp);
-  if (succeeded(splitFullAndPartialTransfer(rewriter, xferOp, options))) {
-    rewriter.finalizeRootUpdate(xferOp);
-    return success();
-  }
-  rewriter.cancelRootUpdate(xferOp);
-  return failure();
+  return splitFullAndPartialTransfer(rewriter, xferOp, options);
 }

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

public:
   virtual void addToWorklist(Operation *op);
 
   /// Pop the next operation from the worklist.
   Operation *popFromWorklist();
 
   /// If the specified operation is in the worklist, remove it.
   void removeFromWorklist(Operation *op);
 
+  /// Notifies the driver that the specified operation may have been modified
+  /// in-place.
+  void finalizeRootUpdate(Operation *op) override;
+
 protected:
   // Implement the hook for inserting operations, and make sure that newly
   // inserted ops are added to the worklist for processing.
   void notifyOperationInserted(Operation *op) override;
 
   // Look over the provided operands for any defining operations that should
   // be re-added to the worklist. This function should be called when an
   // operation is modified or removed, as it may trigger further

 void GreedyPatternRewriteDriver::notifyOperationInserted(Operation *op) {
   LLVM_DEBUG({
     logger.startLine() << "** Insert  : '" << op->getName() << "'(" << op
                        << ")\n";
   });
   addToWorklist(op);
 }
 
+void GreedyPatternRewriteDriver::finalizeRootUpdate(Operation *op) {
+  LLVM_DEBUG({
+    logger.startLine() << "** Modified: '" << op->getName() << "'(" << op
+                       << ")\n";
+  });
+  addToWorklist(op);
+}
+
 void GreedyPatternRewriteDriver::addOperandsToWorklist(ValueRange operands) {
   for (Value operand : operands) {
     // If the use count of this operand is now < 2, we re-add the defining
     // operation to the worklist.
     // TODO: This is based on the fact that zero use operations
     // may be deleted, and that single use values often have more
     // canonicalization opportunities.
     if (!operand || (!operand.use_empty() && !operand.hasOneUse()))

mlir/unittests/Transforms/CMakeLists.txt

 add_mlir_unittest(MLIRTransformsTests
   Canonicalizer.cpp
   DialectConversion.cpp
+  GreedyPatternRewriteDriver.cpp
 )
 target_link_libraries(MLIRTransformsTests
   PRIVATE
   MLIRParser
   MLIRTransforms)

mlir/unittests/Transforms/GreedyPatternRewriteDriver.cpp

+//===- GreedyPatternRewriteDriver.cpp - Greedy pattern rewriter tests -----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/IR/Builders.h"
+#include "gtest/gtest.h"
+
+#include "../../test/lib/Dialect/Test/TestDialect.h"
+
+using namespace mlir;
+
+namespace {
+template <int FromCst, int ToCst>
+struct RewriteAttribute : public OpRewritePattern<test::AnyAttrOfOp> {
+  using OpRewritePattern<test::AnyAttrOfOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(test::AnyAttrOfOp op,
+                                PatternRewriter &rewriter) const override {
+    int64_t constVal = op.getAttr().cast<IntegerAttr>().getInt();
+    if (constVal != FromCst)
+      return failure();
+    rewriter.updateRootInPlace(
+        op, [&]() { op.setAttrAttr(rewriter.getI32IntegerAttr(ToCst)); });
+    return success();
+  }
+};
+
+TEST(GreedyPatternRewriteDriver, EnqueueUpdatedOp) {
+  MLIRContext context;
+  context.allowUnregisteredDialects();
+  context.loadDialect<test::TestDialect>();
+
+  OpBuilder builder(&context);
+  Location loc = builder.getUnknownLoc();
+
+  // Build an op with a region. IsolatedRegionOp has an operand but it is not
+  // used in this test.
+  test::TestOpConstant dummyIndex = builder.create<test::TestOpConstant>(
+      loc, builder.getIndexType(), builder.getIndexAttr(0));
+  auto regionOp = builder.create<test::IsolatedRegionOp>(loc, dummyIndex);
+  Block &block = regionOp.getRegion().emplaceBlock();
+  block.addArgument(builder.getIndexType(), loc);
+  builder.setInsertionPointToStart(&block);
+
+  // Add a test op to the body.
+  test::AnyAttrOfOp testOp =
+      builder.create<test::AnyAttrOfOp>(loc, builder.getI32IntegerAttr(0));
+  builder.create<test::TerminatorOp>(loc);
+
+  // Patterns change the attribute of testOp: 0 -> 1 -> 2 -> 3
+  RewritePatternSet patterns(&context);
+  patterns.insert<RewriteAttribute<0, 1>, RewriteAttribute<1, 2>,
+                  RewriteAttribute<2, 3>>(&context);
+  GreedyRewriteConfig config;
+  config.maxIterations = 1;
+  LogicalResult status =
+      applyPatternsAndFoldGreedily(regionOp, std::move(patterns), config);
+  EXPECT_TRUE(status.succeeded());
+  EXPECT_EQ(testOp.getAttr().cast<IntegerAttr>().getInt(), 3);
+}
+} // namespace

utils/bazel/llvm-project-overlay/mlir/unittests/BUILD.bazel

cc_test(
     ]),
     deps = [
         "//mlir:AffineAnalysis",
         "//mlir:IR",
         "//mlir:Parser",
         "//mlir:Pass",
         "//mlir:TransformUtils",
         "//mlir:Transforms",
+        "//mlir/test:TestDialect",
         "//third-party/unittest:gtest_main",
     ],
 )
 
 cc_test(
     name = "analysis_tests",
     size = "small",
     srcs = glob([