[mlir] Remove the PatternState class and simplify PatternMatchResult.

[rriddle]

Summary

PatternState was a mechanism to pass state between the match and rewrite calls of a RewritePattern. With the rise of matchAndRewrite, this class is unused and unnecessary. This revision removes PatternState and simplifies PatternMatchResult to just be a LogicalResult. A future revision will replace all usages of PatternMatchResult/matchSuccess/matchFailure with LogicalResult equivalents.

[jpienaar]

LG given it is unused.

[rriddle]

Rebase

[rriddle]

Rebase

mlir/include/mlir/IR/PatternMatch.h

bool operator<(const PatternBenefit &rhs) const {
     return representation < rhs.representation;
   }
 
 private:
   PatternBenefit() : representation(ImpossibleToMatchSentinel) {}
   unsigned short representation;
 };
 
-/// Pattern state is used by patterns that want to maintain state between their
-/// match and rewrite phases.  Patterns can define a pattern-specific subclass
-/// of this.
-class PatternState {
-public:
-  virtual ~PatternState() {}
-
-protected:
-  // Must be subclassed.
-  PatternState() {}
-};
-
-/// This is the type returned by a pattern match.  A match failure returns a
-/// None value.  A match success returns a Some value with any state the pattern
-/// may need to maintain (but may also be null).
-using PatternMatchResult = Optional<std::unique_ptr<PatternState>>;
+/// This is the type returned by a pattern match.
+/// TODO: Replace usages with LogicalResult directly.
+using PatternMatchResult = LogicalResult;
 
 //===----------------------------------------------------------------------===//
 // Pattern class
 //===----------------------------------------------------------------------===//
 
 /// Instances of Pattern can be matched against SSA IR.  These matches get used
 /// in ways dependent on their subclasses and the driver doing the matching.
 /// For example, RewritePatterns implement a rewrite from one matched pattern
public:
   /// match multiple root types are instantiated once per root.
   OperationName getRootKind() const { return rootKind; }
 
   //===--------------------------------------------------------------------===//
   // Implementation hooks for patterns to implement.
   //===--------------------------------------------------------------------===//
 
   /// Attempt to match against code rooted at the specified operation,
-  /// which is the same operation code as getRootKind().  On failure, this
-  /// returns a None value.  On success it returns a (possibly null)
-  /// pattern-specific state wrapped in an Optional.
+  /// which is the same operation code as getRootKind().
   virtual PatternMatchResult match(Operation *op) const = 0;
 
   virtual ~Pattern() {}
 
   //===--------------------------------------------------------------------===//
   // Helper methods to simplify pattern implementations
   //===--------------------------------------------------------------------===//
 
-  /// This method indicates that no match was found.
-  static PatternMatchResult matchFailure() { return None; }
+  /// Return a result, indicating that no match was found.
+  PatternMatchResult matchFailure() const { return failure(); }
 
-  /// This method indicates that a match was found and has the specified cost.
-  PatternMatchResult
-  matchSuccess(std::unique_ptr<PatternState> state = {}) const {
-    return PatternMatchResult(std::move(state));
-  }
+  /// This method indicates that a match was found.
+  PatternMatchResult matchSuccess() const { return success(); }
 
 protected:
   /// Patterns must specify the root operation name they match against, and can
   /// also specify the benefit of the pattern matching.
   Pattern(StringRef rootName, PatternBenefit benefit, MLIRContext *context);
 
 private:
   const OperationName rootKind;
   const PatternBenefit benefit;
 
   virtual void anchor();
 };
 
 /// RewritePattern is the common base class for all DAG to DAG replacements.
 /// There are two possible usages of this class:
 ///   * Multi-step RewritePattern with "match" and "rewrite"
 ///     - By overloading the "match" and "rewrite" functions, the user can
 ///       separate the concerns of matching and rewriting.
 ///   * Single-step RewritePattern with "matchAndRewrite"
 ///     - By overloading the "matchAndRewrite" function, the user can perform
-///       the rewrite in the same call as the match. This removes the need for
-///       any PatternState.
+///       the rewrite in the same call as the match.
 ///
 class RewritePattern : public Pattern {
 public:
   /// Rewrite the IR rooted at the specified operation with the result of
   /// this pattern, generating any new operations with the specified
-  /// rewriter.  If an unexpected error is encountered (an internal
-  /// compiler error), it is emitted through the normal MLIR diagnostic
-  /// hooks and the IR is left in a valid state.
-  virtual void rewrite(Operation *op, std::unique_ptr<PatternState> state,
-                       PatternRewriter &rewriter) const;
-
-  /// Rewrite the IR rooted at the specified operation with the result of
-  /// this pattern, generating any new operations with the specified
   /// builder.  If an unexpected error is encountered (an internal
   /// compiler error), it is emitted through the normal MLIR diagnostic
   /// hooks and the IR is left in a valid state.
   virtual void rewrite(Operation *op, PatternRewriter &rewriter) const;
 
   /// Attempt to match against code rooted at the specified operation,
   /// which is the same operation code as getRootKind().  On failure, this
   /// returns a None value.  On success, it returns a (possibly null)
   /// pattern-specific state wrapped in an Optional.  This state is passed back
   /// into the rewrite function if this match is selected.
   PatternMatchResult match(Operation *op) const override;
 
   /// Attempt to match against code rooted at the specified operation,
   /// which is the same operation code as getRootKind(). If successful, this
   /// function will automatically perform the rewrite.
   virtual PatternMatchResult matchAndRewrite(Operation *op,
                                              PatternRewriter &rewriter) const {
-    if (auto matchResult = match(op)) {
-      rewrite(op, std::move(*matchResult), rewriter);
+    if (succeeded(match(op))) {
+      rewrite(op, rewriter);
       return matchSuccess();
     }
     return matchFailure();
   }
 
   /// Return a list of operations that may be generated when rewriting an
   /// operation instance with this pattern.
   ArrayRef<OperationName> getGeneratedOps() const { return generatedOps; }

 /// opposed to a raw Operation.
 template <typename SourceOp> struct OpRewritePattern : public RewritePattern {
   /// Patterns must specify the root operation name they match against, and can
   /// also specify the benefit of the pattern matching.
   OpRewritePattern(MLIRContext *context, PatternBenefit benefit = 1)
       : RewritePattern(SourceOp::getOperationName(), benefit, context) {}
 
   /// Wrappers around the RewritePattern methods that pass the derived op type.
-  void rewrite(Operation *op, std::unique_ptr<PatternState> state,
-               PatternRewriter &rewriter) const final {
-    rewrite(cast<SourceOp>(op), std::move(state), rewriter);
-  }
   void rewrite(Operation *op, PatternRewriter &rewriter) const final {
     rewrite(cast<SourceOp>(op), rewriter);
   }
   PatternMatchResult match(Operation *op) const final {
     return match(cast<SourceOp>(op));
   }
   PatternMatchResult matchAndRewrite(Operation *op,
                                      PatternRewriter &rewriter) const final {
     return matchAndRewrite(cast<SourceOp>(op), rewriter);
   }
 
   /// Rewrite and Match methods that operate on the SourceOp type. These must be
   /// overridden by the derived pattern class.
-  virtual void rewrite(SourceOp op, std::unique_ptr<PatternState> state,
-                       PatternRewriter &rewriter) const {
-    rewrite(op, rewriter);
-  }
   virtual void rewrite(SourceOp op, PatternRewriter &rewriter) const {
-    llvm_unreachable("must override matchAndRewrite or a rewrite method");
+    llvm_unreachable("must override rewrite or matchAndRewrite");
   }
   virtual PatternMatchResult match(SourceOp op) const {
     llvm_unreachable("must override match or matchAndRewrite");
   }
   virtual PatternMatchResult matchAndRewrite(SourceOp op,
                                              PatternRewriter &rewriter) const {
-    if (auto matchResult = match(op)) {
-      rewrite(op, std::move(*matchResult), rewriter);
+    if (succeeded(match(op))) {
+      rewrite(op, rewriter);
       return matchSuccess();
     }
     return matchFailure();
   }
 };
 
 //===----------------------------------------------------------------------===//
 // PatternRewriter class

mlir/include/mlir/Transforms/DialectConversion.h

virtual void rewrite(Operation *op, ArrayRef<Value> operands,
                        ConversionPatternRewriter &rewriter) const {
     llvm_unreachable("unimplemented rewrite");
   }
 
   /// Hook for derived classes to implement combined matching and rewriting.
   virtual PatternMatchResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const {
-    if (!match(op))
+    if (failed(match(op)))
       return matchFailure();
     rewrite(op, operands, rewriter);
     return matchSuccess();
   }
 
   /// Attempt to match and rewrite the IR root at the specified operation.
   PatternMatchResult matchAndRewrite(Operation *op,
                                      PatternRewriter &rewriter) const final;
struct OpConversionPattern : public ConversionPattern {
   virtual void rewrite(SourceOp op, ArrayRef<Value> operands,
                        ConversionPatternRewriter &rewriter) const {
     llvm_unreachable("must override matchAndRewrite or a rewrite method");
   }
 
   virtual PatternMatchResult
   matchAndRewrite(SourceOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const {
-    if (!match(op))
+    if (failed(match(op)))
       return matchFailure();
     rewrite(op, operands, rewriter);
     return matchSuccess();
   }
 
 private:
   using ConversionPattern::matchAndRewrite;
 };

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

for (int64_t off = offset, e = offset + size * stride, idx = 0; off < e;
         // be called multiple times by the PatternRewrite infrastructure (to
         // avoid infinite recursion, but in this case, infinite recursion
         // cannot happen because the rank is strictly decreasing).
         // TODO(rriddle, nicolasvasilache) Implement something like a hook for
         // a potential function that must decrease and allow the same pattern
         // multiple times.
         auto success = matchAndRewrite(insertStridedSliceOp, rewriter);
         (void)success;
-        assert(success && "Unexpected failure");
+        assert(succeeded(success) && "Unexpected failure");
         extractedSource = insertStridedSliceOp;
       }
       // 4. Insert the extractedSource into the res vector.
       res = insertOne(rewriter, loc, extractedSource, res, off);
     }
 
     rewriter.replaceOp(op, res);
     return matchSuccess();
for (int64_t off = offset, e = offset + size * stride, idx = 0; off < e;
         // be called multiple times by the PatternRewrite infrastructure (to
         // avoid infinite recursion, but in this case, infinite recursion
         // cannot happen because the rank is strictly decreasing).
         // TODO(rriddle, nicolasvasilache) Implement something like a hook for
         // a potential function that must decrease and allow the same pattern
         // multiple times.
         auto success = matchAndRewrite(stridedSliceOp, rewriter);
         (void)success;
-        assert(success && "Unexpected failure");
+        assert(succeeded(success) && "Unexpected failure");
         extracted = stridedSliceOp;
       }
       res = insertOne(rewriter, loc, extracted, res, idx);
     }
     rewriter.replaceOp(op, {res});
     return matchSuccess();
   }
 };

mlir/lib/Dialect/SPIRV/SPIRVCanonicalization.cpp

PatternMatchResult matchAndRewrite(spirv::SelectionOp selectionOp,
 
     auto brConditionalOp =
         cast<spirv::BranchConditionalOp>(headerBlock->front());
 
     auto *trueBlock = brConditionalOp.getSuccessor(0);
     auto *falseBlock = brConditionalOp.getSuccessor(1);
     auto *mergeBlock = selectionOp.getMergeBlock();
 
-    if (!canCanonicalizeSelection(trueBlock, falseBlock, mergeBlock)) {
+    if (failed(canCanonicalizeSelection(trueBlock, falseBlock, mergeBlock)))
       return matchFailure();
-    }
 
     auto trueValue = getSrcValue(trueBlock);
     auto falseValue = getSrcValue(falseBlock);
     auto ptrValue = getDstPtr(trueBlock);
     auto storeOpAttributes =
         cast<spirv::StoreOp>(trueBlock->front()).getOperation()->getAttrs();
 
     auto selectOp = rewriter.create<spirv::SelectOp>(

mlir/lib/IR/PatternMatch.cpp

 
 // Out-of-line vtable anchor.
 void Pattern::anchor() {}
 
 //===----------------------------------------------------------------------===//
 // RewritePattern and PatternRewriter implementation
 //===----------------------------------------------------------------------===//
 
-void RewritePattern::rewrite(Operation *op, std::unique_ptr<PatternState> state,
-                             PatternRewriter &rewriter) const {
-  rewrite(op, rewriter);
-}
-
 void RewritePattern::rewrite(Operation *op, PatternRewriter &rewriter) const {
   llvm_unreachable("need to implement either matchAndRewrite or one of the "
                    "rewrite functions!");
 }
 
 PatternMatchResult RewritePattern::match(Operation *op) const {
   llvm_unreachable("need to implement either match or matchAndRewrite!");
 }
bool RewritePatternMatcher::matchAndRewrite(Operation *op,
   for (auto *pattern : patterns) {
     // Ignore patterns that are for the wrong root or are impossible to match.
     if (pattern->getRootKind() != op->getName() ||
         pattern->getBenefit().isImpossibleToMatch())
       continue;
 
     // Try to match and rewrite this pattern. The patterns are sorted by
     // benefit, so if we match we can immediately rewrite and return.
-    if (pattern->matchAndRewrite(op, rewriter))
+    if (succeeded(pattern->matchAndRewrite(op, rewriter)))
       return true;
   }
   return false;
 }

mlir/lib/Transforms/DialectConversion.cpp

auto cleanupFailure = [&] {
     // Reset the rewriter state and pop this pattern.
     rewriterImpl.resetState(curState);
     appliedPatterns.erase(pattern);
     return failure();
   };
 
   // Try to rewrite with the given pattern.
   rewriter.setInsertionPoint(op);
-  auto matchedPattern = pattern->matchAndRewrite(op, rewriter);
+  LogicalResult matchedPattern = pattern->matchAndRewrite(op, rewriter);
 #ifndef NDEBUG
   assert(rewriterImpl.pendingRootUpdates.empty() && "dangling root updates");
 #endif
 
-  if (!matchedPattern) {
+  if (failed(matchedPattern)) {
     LLVM_DEBUG(logFailure(rewriterImpl.logger, "pattern failed to match"));
     return cleanupFailure();
   }
 
   // If the pattern moved or created any blocks, try to legalize their types.
   // This ensures that the types of the block arguments are legal for the region
   // they were moved into.
   for (unsigned i = curState.numBlockActions,