[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 ↗

(On Diff #475732)

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/test/IR/greedy-pattern-rewriter-driver.mlir

+// RUN: mlir-opt %s -test-patterns="max-iterations=1" | FileCheck %s
+
+// CHECK-LABEL: func @add_to_worklist_after_inplace_update()
+func.func @add_to_worklist_after_inplace_update() {
+  // The following op is updated in-place and should be added back to the
+  // worklist of the GreedyPatternRewriteDriver (regardless of the value of
+  // config.max_iterations).
+
+  // CHECK: "test.any_attr_of_i32_str"() {attr = 3 : i32} : () -> ()
+  "test.any_attr_of_i32_str"() {attr = 0 : i32} : () -> ()
+  return
+}
 No newline at end of file

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

LogicalResult matchAndRewrite(TestCommutative2Op op,
     Attribute constInput;
     if (!matchPattern(operand->getOperand(1), m_Constant(&constInput)))
       return failure();
     rewriter.replaceOp(op, operand->getOperand(1));
     return success();
   }
 };
 
+/// This pattern matches test.any_attr_of_i32_str ops. In case of an integer
+/// attribute with value smaller than MaxVal, it increments the value by 1.
+template <int MaxVal>
+struct IncrementIntAttribute : public OpRewritePattern<test::AnyAttrOfOp> {
+  using OpRewritePattern<test::AnyAttrOfOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(test::AnyAttrOfOp op,
+                                PatternRewriter &rewriter) const override {
+    auto intAttr = op.getAttr().dyn_cast<IntegerAttr>();
+    if (!intAttr)
+      return failure();
+    int64_t val = intAttr.getInt();
+    if (val >= MaxVal)
+      return failure();
+    rewriter.updateRootInPlace(
+        op, [&]() { op.setAttrAttr(rewriter.getI32IntegerAttr(val + 1)); });
+    return success();
+  }
+};
+
 struct TestPatternDriver
     : public PassWrapper<TestPatternDriver, OperationPass<func::FuncOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPatternDriver)
 
   TestPatternDriver() = default;
   TestPatternDriver(const TestPatternDriver &other) : PassWrapper(other) {}
 
   StringRef getArgument() const final { return "test-patterns"; }
   StringRef getDescription() const final { return "Run test dialect patterns"; }
   void runOnOperation() override {
     mlir::RewritePatternSet patterns(&getContext());
     populateWithGenerated(patterns);
 
     // Verify named pattern is generated with expected name.
     patterns.add<FoldingPattern, TestNamedPatternRule,
                  FolderInsertBeforePreviouslyFoldedConstantPattern,
                  FolderCommutativeOp2WithConstant>(&getContext());
 
+    // Additional patterns for testing the GreedyPatternRewriteDriver.
+    patterns.insert<IncrementIntAttribute<3>>(&getContext());
+
     GreedyRewriteConfig config;
     config.useTopDownTraversal = this->useTopDownTraversal;
+    config.maxIterations = this->maxIterations;
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns),
                                        config);
   }
 
   Option<bool> useTopDownTraversal{
       *this, "top-down",
       llvm::cl::desc("Seed the worklist in general top-down order"),
       llvm::cl::init(GreedyRewriteConfig().useTopDownTraversal)};
+  Option<int> maxIterations{
+      *this, "max-iterations",
+      llvm::cl::desc("Max. iterations in the GreedyRewriteConfig"),
+      llvm::cl::init(GreedyRewriteConfig().maxIterations)};
 };
 
 struct TestStrictPatternDriver
     : public PassWrapper<TestStrictPatternDriver, OperationPass<func::FuncOp>> {
 public:
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestStrictPatternDriver)
 
   TestStrictPatternDriver() = default;