//===- PatternMatch.cpp - Base classes for pattern match ------------------===//

//

// 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/IR/PatternMatch.h"

#include "mlir/IR/BlockAndValueMapping.h"

#include "mlir/IR/Operation.h"

#include "mlir/IR/Value.h"

using namespace mlir;

PatternBenefit::PatternBenefit(unsigned benefit) : representation(benefit) {

  assert(representation == benefit && benefit != ImpossibleToMatchSentinel &&

         "This pattern match benefit is too large to represent");

}

unsigned short PatternBenefit::getBenefit() const {

  assert(representation != ImpossibleToMatchSentinel &&

         "Pattern doesn't match");

  return representation;

}

//===----------------------------------------------------------------------===//

// Pattern implementation

//===----------------------------------------------------------------------===//

Pattern::Pattern(StringRef rootName, PatternBenefit benefit,

                 MLIRContext *context)

    : rootKind(OperationName(rootName, context)), benefit(benefit) {}

// Out-of-line vtable anchor.

void Pattern::anchor() {}

//===----------------------------------------------------------------------===//

// RewritePattern and PatternRewriter implementation

//===----------------------------------------------------------------------===//

void RewritePattern::rewrite(Operation *op, PatternRewriter &rewriter) const {

  llvm_unreachable("need to implement either matchAndRewrite or one of the "

                   "rewrite functions!");

}

LogicalResult RewritePattern::match(Operation *op) const {

  llvm_unreachable("need to implement either match or matchAndRewrite!");

}

/// Patterns must specify the root operation name they match against, and can

/// also specify the benefit of the pattern matching. They can also specify the

/// names of operations that may be generated during a successful rewrite.

RewritePattern::RewritePattern(StringRef rootName,

                               ArrayRef<StringRef> generatedNames,

                               PatternBenefit benefit, MLIRContext *context)

    : Pattern(rootName, benefit, context) {

  generatedOps.reserve(generatedNames.size());

  std::transform(generatedNames.begin(), generatedNames.end(),

                 std::back_inserter(generatedOps), [context](StringRef name) {

                   return OperationName(name, context);

                 });

}

PatternRewriter::~PatternRewriter() {

  // Out of line to provide a vtable anchor for the class.

}

/// This method performs the final replacement for a pattern, where the

/// results of the operation are updated to use the specified list of SSA

/// values.

void PatternRewriter::replaceOp(Operation *op, ValueRange newValues) {

  // Notify the rewriter subclass that we're about to replace this root.

  notifyRootReplaced(op);

  assert(op->getNumResults() == newValues.size() &&

         "incorrect # of replacement values");

  op->replaceAllUsesWith(newValues);

  notifyOperationRemoved(op);

  op->erase();

}

/// This method erases an operation that is known to have no uses. The uses of

/// the given operation *must* be known to be dead.

void PatternRewriter::eraseOp(Operation *op) {

  assert(op->use_empty() && "expected 'op' to have no uses");

  notifyOperationRemoved(op);

  op->erase();

}

void PatternRewriter::eraseBlock(Block *block) {

  for (auto &op : llvm::make_early_inc_range(llvm::reverse(*block))) {

    assert(op.use_empty() && "expected 'op' to have no uses");

    eraseOp(&op);

  }

  block->erase();

}

/// Merge the operations of block 'source' into the end of block 'dest'.

/// 'source's predecessors must be empty or only contain 'dest`.

/// 'argValues' is used to replace the block arguments of 'source' after

/// merging.

void PatternRewriter::mergeBlocks(Block *source, Block *dest,

                                  ValueRange argValues) {

  assert(llvm::all_of(source->getPredecessors(),

                      [dest](Block *succ) { return succ == dest; }) &&

         "expected 'source' to have no predecessors or only 'dest'");

  assert(argValues.size() == source->getNumArguments() &&

         "incorrect # of argument replacement values");

  // Replace all of the successor arguments with the provided values.

  for (auto it : llvm::zip(source->getArguments(), argValues))

    std::get<0>(it).replaceAllUsesWith(std::get<1>(it));

  // Splice the operations of the 'source' block into the 'dest' block and erase

  // it.

  dest->getOperations().splice(dest->end(), source->getOperations());

  source->dropAllUses();

  source->erase();

}

/// Split the operations starting at "before" (inclusive) out of the given

/// block into a new block, and return it.

Block *PatternRewriter::splitBlock(Block *block, Block::iterator before) {

  return block->splitBlock(before);

}

/// 'op' and 'newOp' are known to have the same number of results, replace the

/// uses of op with uses of newOp

void PatternRewriter::replaceOpWithResultsOfAnotherOp(Operation *op,

                                                      Operation *newOp) {

  assert(op->getNumResults() == newOp->getNumResults() &&

         "replacement op doesn't match results of original op");

  if (op->getNumResults() == 1)

    return replaceOp(op, newOp->getResult(0));

  return replaceOp(op, newOp->getResults());

}

/// Move the blocks that belong to "region" before the given position in

/// another region.  The two regions must be different.  The caller is in

/// charge to update create the operation transferring the control flow to the

/// region and pass it the correct block arguments.

void PatternRewriter::inlineRegionBefore(Region &region, Region &parent,

                                         Region::iterator before) {

  parent.getBlocks().splice(before, region.getBlocks());

}

void PatternRewriter::inlineRegionBefore(Region &region, Block *before) {

  inlineRegionBefore(region, *before->getParent(), before->getIterator());

}

/// Clone the blocks that belong to "region" before the given position in

/// another region "parent". The two regions must be different. The caller is

/// responsible for creating or updating the operation transferring flow of

/// control to the region and passing it the correct block arguments.

void PatternRewriter::cloneRegionBefore(Region &region, Region &parent,

                                        Region::iterator before,

                                        BlockAndValueMapping &mapping) {

  region.cloneInto(&parent, before, mapping);

}

void PatternRewriter::cloneRegionBefore(Region &region, Region &parent,

                                        Region::iterator before) {

  BlockAndValueMapping mapping;

  cloneRegionBefore(region, parent, before, mapping);

}

void PatternRewriter::cloneRegionBefore(Region &region, Block *before) {

  cloneRegionBefore(region, *before->getParent(), before->getIterator());

}

//===----------------------------------------------------------------------===//

// PatternMatcher implementation

//===----------------------------------------------------------------------===//

RewritePatternMatcher::RewritePatternMatcher(

    const OwningRewritePatternList &patterns) {

  for (auto &pattern : patterns)

    this->patterns.push_back(pattern.get());

  // Sort the patterns by benefit to simplify the matching logic.

  std::stable_sort(this->patterns.begin(), this->patterns.end(),

                   [](RewritePattern *l, RewritePattern *r) {

                     return r->getBenefit() < l->getBenefit();

                   });

}

/// Try to match the given operation to a pattern and rewrite it.

bool RewritePatternMatcher::matchAndRewrite(Operation *op,

                                            PatternRewriter &rewriter) {

  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 (succeeded(pattern->matchAndRewrite(op, rewriter)))

      return true;

  }

  return false;

}