[mlir][NFC] GreedyPatternRewriteDriver: Consistent return values

[springerm]

Summary

All apply... functions now return a LogicalResult indicating whether the iterative process converged or not.

Depends On: D141844

[mehdi_amini]

1. List Item

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

101

mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h

inline LogicalResult applyPatternsAndFoldGreedily(
   return applyPatternsAndFoldGreedily(op->getRegions(), patterns, config);
 }
 
 /// Applies the specified patterns on `op` alone while also trying to fold it,
 /// by selecting the highest benefits patterns in a greedy manner. Returns
 /// success if no more patterns can be matched. `erased` is set to true if `op`
 /// was folded away or erased as a result of becoming dead. Note: This does not
 /// apply any patterns recursively to the regions of `op`.
+///
+/// Returns success if the iterative process converged and no more patterns can
+/// be matched.
 LogicalResult applyOpPatternsAndFold(Operation *op,
                                      const FrozenRewritePatternSet &patterns,
                                      bool *erased = nullptr);
 
 /// Applies the specified rewrite patterns on `ops` while also trying to fold
 /// these ops as well as any other ops that were in turn created due to such
 /// rewrites. Furthermore, any pre-existing ops in the IR outside of `ops`
 /// remain completely unmodified if `strict` is set to true. If `strict` is
 /// false, other operations that use results of rewritten ops or supply operands
 /// to such ops are in turn simplified; any other ops still remain unmodified
 /// (i.e., regardless of `strict`). Note that ops in `ops` could be erased as a
 /// result of folding, becoming dead, or via pattern rewrites. If more far
 /// reaching simplification is desired, applyPatternsAndFoldGreedily should be
-/// used. Returns true if at all any IR was rewritten.
-bool applyOpPatternsAndFold(ArrayRef<Operation *> ops,
+/// used.
+///
+/// `changed` is set to true if the IR was modified at all. Returns success if
+/// Returns success if the iterative process converged and no more patterns can
+/// be matched.
+LogicalResult applyOpPatternsAndFold(ArrayRef<Operation *> ops,
                                      const FrozenRewritePatternSet &patterns,
-                            bool strict);
+                                     bool strict, bool *changed = nullptr);
 
 } // namespace mlir
 
 #endif // MLIR_TRANSFORMS_GREEDYPATTERNREWRITEDRIVER_H_

mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp

SimplifyBoundedAffineOpsOp::apply(TransformResults &results,
   // Canonicalization patterns are needed so that affine.apply ops are composed
   // with the remaining affine.min/max ops.
   AffineMaxOp::getCanonicalizationPatterns(patterns, getContext());
   AffineMinOp::getCanonicalizationPatterns(patterns, getContext());
   patterns.insert<SimplifyAffineMinMaxOp<AffineMinOp>,
                   SimplifyAffineMinMaxOp<AffineMaxOp>>(getContext(), cstr);
   FrozenRewritePatternSet frozenPatterns(std::move(patterns));
   // Apply the simplification pattern to a fixpoint.
-  (void)applyOpPatternsAndFold(targets, frozenPatterns, /*strict=*/true);
+  if (failed(
+          applyOpPatternsAndFold(targets, frozenPatterns, /*strict=*/true))) {
+    auto diag = emitDefiniteFailure()
+                << "affine.min/max simplification did not converge";
+    return diag;
+  }
   return DiagnosedSilenceableFailure::success();
 }
 
 void SimplifyBoundedAffineOpsOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getTarget(), effects);
   for (Value v : getBoundedValues())
     onlyReadsHandle(v, effects);

mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp

 class MultiOpPatternRewriteDriver : public GreedyPatternRewriteDriver {
 public:
   explicit MultiOpPatternRewriteDriver(MLIRContext *ctx,
                                        const FrozenRewritePatternSet &patterns,
                                        bool strict)
       : GreedyPatternRewriteDriver(ctx, patterns, GreedyRewriteConfig()),
         strictMode(strict) {}
 
-  bool simplifyLocally(ArrayRef<Operation *> op);
+  LogicalResult simplifyLocally(ArrayRef<Operation *> op,
+                                bool *changed = nullptr);
 
   void addToWorklist(Operation *op) override {
     if (!strictMode || strictModeFilteredOps.contains(op))
       GreedyPatternRewriteDriver::addSingleOpToWorklist(op);
   }
 
 private:
   void notifyOperationInserted(Operation *op) override {

 /// result of folding, becoming dead, or via pattern rewrites. Returns true if
 /// at all any changes happened.
 // Unlike `OpPatternRewriteDriver::simplifyLocally` which works on a single op
 // or GreedyPatternRewriteDriver::simplify, this method just iterates until
 // the worklist is empty. As our objective is to keep simplification "local",
 // there is no strong rationale to re-add all operations into the worklist and
 // rerun until an iteration changes nothing. If more widereaching simplification
 // is desired, GreedyPatternRewriteDriver should be used.
-bool MultiOpPatternRewriteDriver::simplifyLocally(ArrayRef<Operation *> ops) {
+LogicalResult
+MultiOpPatternRewriteDriver::simplifyLocally(ArrayRef<Operation *> ops,
+                                             bool *changed) {
   if (strictMode) {
     strictModeFilteredOps.clear();
     strictModeFilteredOps.insert(ops.begin(), ops.end());
   }
 
-  bool changed = false;
+  if (changed)
+    *changed = false;
   worklist.clear();
   worklistMap.clear();
   for (Operation *op : ops)
     addToWorklist(op);
 
   // These are scratch vectors used in the folding loop below.
   SmallVector<Value, 8> originalOperands, resultValues;
   int64_t numRewrites = 0;
while (!worklist.empty() &&
 
     assert((!strictMode || strictModeFilteredOps.contains(op)) &&
            "unexpected op was inserted under strict mode");
 
     // If the operation is trivially dead - remove it.
     if (isOpTriviallyDead(op)) {
       notifyOperationRemoved(op);
       op->erase();
-      changed = true;
+      if (changed)
+        *changed = true;
       continue;
     }
 
     // Collects all the operands and result uses of the given `op` into work
     // list. Also remove `op` and nested ops from worklist.
     originalOperands.assign(op->operand_begin(), op->operand_end());
     auto preReplaceAction = [&](Operation *op) {
       // Add the operands to the worklist for visitation.
while (!worklist.empty() &&
     auto processGeneratedConstants = [this](Operation *op) {
       notifyOperationInserted(op);
     };
 
     // Try to fold this op.
     bool inPlaceUpdate;
     if (succeeded(folder.tryToFold(op, processGeneratedConstants,
                                    preReplaceAction, &inPlaceUpdate))) {
-      changed = true;
+      if (changed)
+        *changed = true;
       if (!inPlaceUpdate) {
         // Op has been erased.
         continue;
       }
     }
 
     // Try to match one of the patterns. The rewriter is automatically
     // notified of any necessary changes, so there is nothing else to do
     // here.
     if (succeeded(matcher.matchAndRewrite(op, *this))) {
-      changed = true;
+      if (changed)
+        *changed = true;
       ++numRewrites;
     }
   }
 
-  return changed;
+  return success(worklist.empty());
 }
 
 /// Rewrites only `op` using the supplied canonicalization patterns and
 /// folding. `erased` is set to true if the op is erased as a result of being
 /// folded, replaced, or dead.
 LogicalResult mlir::applyOpPatternsAndFold(
     Operation *op, const FrozenRewritePatternSet &patterns, bool *erased) {
   // Start the pattern driver.
   GreedyRewriteConfig config;
   OpPatternRewriteDriver driver(op->getContext(), patterns);
   bool opErased;
   LogicalResult converged =
       driver.simplifyLocally(op, config.maxNumRewrites, opErased);
   if (erased)
     *erased = opErased;
   LLVM_DEBUG(if (failed(converged)) {
     llvm::dbgs() << "The pattern rewrite did not converge after "
                  << config.maxNumRewrites << " rewrites";
   });
   return converged;
 }
 
-bool mlir::applyOpPatternsAndFold(ArrayRef<Operation *> ops,
-                                  const FrozenRewritePatternSet &patterns,
-                                  bool strict) {
-  if (ops.empty())
-    return false;
+LogicalResult
+mlir::applyOpPatternsAndFold(ArrayRef<Operation *> ops,
+                             bool strict, bool *changed) {
+  if (ops.empty()) {
+    if (changed)
+      *changed = false;
+    return success();
+  }
 
   // Start the pattern driver.
   MultiOpPatternRewriteDriver driver(ops.front()->getContext(), patterns,
                                      strict);
-  return driver.simplifyLocally(ops);
+  return driver.simplifyLocally(ops, changed);
 }