PatternMatch.h

Go to the documentation of this file.

//===- PatternMatch.h - PatternMatcher classes -------==---------*- C++ -*-===//
//
// Part of the MLIR 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
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_PATTERNMATCHER_H
#define MLIR_PATTERNMATCHER_H

#include "mlir/IR/Builders.h"

namespace mlir {

class PatternRewriter;

//===----------------------------------------------------------------------===//
// PatternBenefit class
//===----------------------------------------------------------------------===//

class PatternBenefit {
  enum { ImpossibleToMatchSentinel = 65535 };

public:
  /*implicit*/ PatternBenefit(unsigned benefit);
  PatternBenefit(const PatternBenefit &) = default;
  PatternBenefit &operator=(const PatternBenefit &) = default;

  static PatternBenefit impossibleToMatch() { return PatternBenefit(); }
  bool isImpossibleToMatch() const { return *this == impossibleToMatch(); }

  // corresponding pattern isImpossibleToMatch() then this aborts.
  unsigned short getBenefit() const;

  bool operator==(const PatternBenefit &rhs) const {
    return representation == rhs.representation;
  }
  bool operator!=(const PatternBenefit &rhs) const { return !(*this == rhs); }
  bool operator<(const PatternBenefit &rhs) const {
    return representation < rhs.representation;
  }

private:
  PatternBenefit() : representation(ImpossibleToMatchSentinel) {}
  unsigned short representation;
};

class PatternState {
public:
  virtual ~PatternState() {}

protected:
  // Must be subclassed.
  PatternState() {}
};

using PatternMatchResult = Optional<std::unique_ptr<PatternState>>;

//===----------------------------------------------------------------------===//
// Pattern class
//===----------------------------------------------------------------------===//

class Pattern {
public:
  PatternBenefit getBenefit() const { return benefit; }

  OperationName getRootKind() const { return rootKind; }

  //===--------------------------------------------------------------------===//
  // Implementation hooks for patterns to implement.
  //===--------------------------------------------------------------------===//

  virtual PatternMatchResult match(Operation *op) const = 0;

  virtual ~Pattern() {}

  //===--------------------------------------------------------------------===//
  // Helper methods to simplify pattern implementations
  //===--------------------------------------------------------------------===//

  static PatternMatchResult matchFailure() { return None; }

  PatternMatchResult
  matchSuccess(std::unique_ptr<PatternState> state = {}) const {
    return PatternMatchResult(std::move(state));
  }

protected:
  Pattern(StringRef rootName, PatternBenefit benefit, MLIRContext *context);

private:
  const OperationName rootKind;
  const PatternBenefit benefit;

  virtual void anchor();
};

class RewritePattern : public Pattern {
public:
  virtual void rewrite(Operation *op, std::unique_ptr<PatternState> state,
                       PatternRewriter &rewriter) const;

  virtual void rewrite(Operation *op, PatternRewriter &rewriter) const;

  PatternMatchResult match(Operation *op) const override;

  virtual PatternMatchResult matchAndRewrite(Operation *op,
                                             PatternRewriter &rewriter) const {
    if (auto matchResult = match(op)) {
      rewrite(op, std::move(*matchResult), rewriter);
      return matchSuccess();
    }
    return matchFailure();
  }

  ArrayRef<OperationName> getGeneratedOps() const { return generatedOps; }

protected:
  RewritePattern(StringRef rootName, PatternBenefit benefit,
                 MLIRContext *context)
      : Pattern(rootName, benefit, context) {}
  RewritePattern(StringRef rootName, ArrayRef<StringRef> generatedNames,
                 PatternBenefit benefit, MLIRContext *context);

  SmallVector<OperationName, 2> generatedOps;
};

template <typename SourceOp> struct OpRewritePattern : public RewritePattern {
  OpRewritePattern(MLIRContext *context, PatternBenefit benefit = 1)
      : RewritePattern(SourceOp::getOperationName(), benefit, context) {}

  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);
  }

  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");
  }
  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);
      return matchSuccess();
    }
    return matchFailure();
  }
};

//===----------------------------------------------------------------------===//
// PatternRewriter class
//===----------------------------------------------------------------------===//

class PatternRewriter : public OpBuilder {
public:
  template <typename OpTy, typename... Args>
  OpTy create(Location location, Args... args) {
    OperationState state(location, OpTy::getOperationName());
    OpTy::build(this, state, args...);
    auto *op = createOperation(state);
    auto result = dyn_cast<OpTy>(op);
    assert(result && "Builder didn't return the right type");
    return result;
  }

  template <typename OpTy, typename... Args>
  OpTy createChecked(Location location, Args... args) {
    OperationState state(location, OpTy::getOperationName());
    OpTy::build(this, state, args...);
    auto *op = createOperation(state);

    // If the Operation we produce is valid, return it.
    if (!OpTy::verifyInvariants(op)) {
      auto result = dyn_cast<OpTy>(op);
      assert(result && "Builder didn't return the right type");
      return result;
    }

    // Otherwise, the error message got emitted.  Just remove the operation
    // we made.
    op->erase();
    return OpTy();
  }

  virtual Operation *insert(Operation *op) = 0;

  virtual void inlineRegionBefore(Region &region, Region &parent,
                                  Region::iterator before);
  void inlineRegionBefore(Region &region, Block *before);

  virtual void cloneRegionBefore(Region &region, Region &parent,
                                 Region::iterator before,
                                 BlockAndValueMapping &mapping);
  void cloneRegionBefore(Region &region, Region &parent,
                         Region::iterator before);
  void cloneRegionBefore(Region &region, Block *before);

  virtual void replaceOp(Operation *op, ValueRange newValues,
                         ValueRange valuesToRemoveIfDead);
  void replaceOp(Operation *op, ValueRange newValues) {
    replaceOp(op, newValues, llvm::None);
  }

  template <typename OpTy, typename... Args>
  void replaceOpWithNewOp(Operation *op, Args &&... args) {
    auto newOp = create<OpTy>(op->getLoc(), std::forward<Args>(args)...);
    replaceOpWithResultsOfAnotherOp(op, newOp.getOperation(), {});
  }

  template <typename OpTy, typename... Args>
  void replaceOpWithNewOp(ValueRange valuesToRemoveIfDead, Operation *op,
                          Args &&... args) {
    auto newOp = create<OpTy>(op->getLoc(), std::forward<Args>(args)...);
    replaceOpWithResultsOfAnotherOp(op, newOp.getOperation(),
                                    valuesToRemoveIfDead);
  }

  virtual void eraseOp(Operation *op);

  virtual void mergeBlocks(Block *source, Block *dest,
                           ValueRange argValues = llvm::None);

  virtual Block *splitBlock(Block *block, Block::iterator before);

  virtual void startRootUpdate(Operation *op) {}

  virtual void finalizeRootUpdate(Operation *op) {}

  virtual void cancelRootUpdate(Operation *op) {}

  template <typename CallableT>
  void updateRootInPlace(Operation *root, CallableT &&callable) {
    startRootUpdate(root);
    callable();
    finalizeRootUpdate(root);
  }

protected:
  explicit PatternRewriter(MLIRContext *ctx) : OpBuilder(ctx) {}
  virtual ~PatternRewriter();

  // These are the callback methods that subclasses can choose to implement if
  // they would like to be notified about certain types of mutations.

  virtual void notifyRootReplaced(Operation *op) {}

  virtual void notifyOperationRemoved(Operation *op) {}

private:
  void replaceOpWithResultsOfAnotherOp(Operation *op, Operation *newOp,
                                       ValueRange valuesToRemoveIfDead);
};

//===----------------------------------------------------------------------===//
// Pattern-driven rewriters
//===----------------------------------------------------------------------===//

class OwningRewritePatternList {
  using PatternListT = std::vector<std::unique_ptr<RewritePattern>>;

public:
  PatternListT::iterator begin() { return patterns.begin(); }
  PatternListT::iterator end() { return patterns.end(); }
  PatternListT::const_iterator begin() const { return patterns.begin(); }
  PatternListT::const_iterator end() const { return patterns.end(); }
  void clear() { patterns.clear(); }

  //===--------------------------------------------------------------------===//
  // Pattern Insertion
  //===--------------------------------------------------------------------===//

  template <typename... Ts, typename ConstructorArg,
            typename... ConstructorArgs,
            typename = std::enable_if_t<sizeof...(Ts) != 0>>
  void insert(ConstructorArg &&arg, ConstructorArgs &&... args) {
    // The following expands a call to emplace_back for each of the pattern
    // types 'Ts'. This magic is necessary due to a limitation in the places
    // that a parameter pack can be expanded in c++11.
    // FIXME: In c++17 this can be simplified by using 'fold expressions'.
    using dummy = int[];
    (void)dummy{
        0, (patterns.emplace_back(std::make_unique<Ts>(arg, args...)), 0)...};
  }

private:
  PatternListT patterns;
};

class RewritePatternMatcher {
public:
  explicit RewritePatternMatcher(const OwningRewritePatternList &patterns);

  bool matchAndRewrite(Operation *op, PatternRewriter &rewriter);

private:
  RewritePatternMatcher(const RewritePatternMatcher &) = delete;
  void operator=(const RewritePatternMatcher &) = delete;

  std::vector<RewritePattern *> patterns;
};

bool applyPatternsGreedily(Operation *op,
                           const OwningRewritePatternList &patterns);
bool applyPatternsGreedily(MutableArrayRef<Region> regions,
                           const OwningRewritePatternList &patterns);
} // end namespace mlir

#endif // MLIR_PATTERN_MATCH_H