Address Performance issue in NestedMatcher

A performance issue was reported due to the usage of NestedMatcher in
ComposeAffineMaps. The main culprit was the ubiquitous copies that were
occuring when appending even a single element in `matchOne`.

This CL generally simplifies the implementation and removes one level of indirection by getting rid of
auxiliary storage as well as simplifying the API.
The users of the API are updated accordingly.

The implementation was tested on a heavily unrolled example with
ComposeAffineMaps and is now close in performance with an implementation based
on stateless InstWalker.

As a reminder, the whole ComposeAffineMaps pass is slated to disappear but the bug report was very useful as a stress test for NestedMatchers.

Lastly, the following cleanups reported by @aminim were addressed:
1. make NestedPatternContext scoped within runFunction rather than at the Pass level. This was caused by a previous misunderstanding of Pass lifetime;
2. use defensive assertions in the constructor of NestedPatternContext to make it clear a unique such locally scoped context is allowed to exist.

PiperOrigin-RevId: 231781279

1 files changed, 42 insertions, 116 deletions

diff --git a/mlir/lib/Analysis/NestedMatcher.cpp b/mlir/lib/Analysis/NestedMatcher.cpp
index 491a9bef1b9..43e3b332a58 100644
--- a/mlir/lib/Analysis/NestedMatcher.cpp
+++ b/mlir/lib/Analysis/NestedMatcher.cpp
@@ -20,93 +20,49 @@
 #include "mlir/StandardOps/StandardOps.h"
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/raw_ostream.h"
 
-namespace mlir {
-
-/// Underlying storage for NestedMatch.
-struct NestedMatchStorage {
-  MutableArrayRef<NestedMatch::EntryType> matches;
-};
-
-/// Underlying storage for NestedPattern.
-struct NestedPatternStorage {
-  NestedPatternStorage(Instruction::Kind k, ArrayRef<NestedPattern> c,
-                       FilterFunctionType filter, Instruction *skip)
-      : kind(k), nestedPatterns(c), filter(filter), skip(skip) {}
-
-  Instruction::Kind kind;
-  ArrayRef<NestedPattern> nestedPatterns;
-  FilterFunctionType filter;
-  /// skip is needed so that we can implement match without switching on the
-  /// type of the Instruction.
-  /// The idea is that a NestedPattern first checks if it matches locally
-  /// and then recursively applies its nested matchers to its elem->nested.
-  /// Since we want to rely on the InstWalker impl rather than duplicate its
-  /// the logic, we allow an off-by-one traversal to account for the fact that
-  /// we write:
-  ///
-  ///  void match(Instruction *elem) {
-  ///    for (auto &c : getNestedPatterns()) {
-  ///      NestedPattern childPattern(...);
-  ///                                     ^~~~ Needs off-by-one skip.
-  ///
-  Instruction *skip;
-};
-
-} // end namespace mlir
-
 using namespace mlir;
 
 llvm::BumpPtrAllocator *&NestedMatch::allocator() {
-  static thread_local llvm::BumpPtrAllocator *allocator = nullptr;
+  thread_local llvm::BumpPtrAllocator *allocator = nullptr;
   return allocator;
 }
 
-NestedMatch NestedMatch::build(ArrayRef<NestedMatch::EntryType> elements) {
-  auto *matches =
-      allocator()->Allocate<NestedMatch::EntryType>(elements.size());
-  std::uninitialized_copy(elements.begin(), elements.end(), matches);
-  auto *storage = allocator()->Allocate<NestedMatchStorage>();
-  new (storage) NestedMatchStorage();
-  storage->matches =
-      MutableArrayRef<NestedMatch::EntryType>(matches, elements.size());
+NestedMatch NestedMatch::build(Instruction *instruction,
+                               ArrayRef<NestedMatch> nestedMatches) {
   auto *result = allocator()->Allocate<NestedMatch>();
-  new (result) NestedMatch(storage);
+  auto *children = allocator()->Allocate<NestedMatch>(nestedMatches.size());
+  std::uninitialized_copy(nestedMatches.begin(), nestedMatches.end(), children);
+  new (result) NestedMatch();
+  result->matchedInstruction = instruction;
+  result->matchedChildren =
+      ArrayRef<NestedMatch>(children, nestedMatches.size());
   return *result;
 }
 
-NestedMatch::iterator NestedMatch::begin() { return storage->matches.begin(); }
-NestedMatch::iterator NestedMatch::end() { return storage->matches.end(); }
-NestedMatch::EntryType &NestedMatch::front() {
-  return *storage->matches.begin();
-}
-NestedMatch::EntryType &NestedMatch::back() {
-  return *(storage->matches.begin() + size() - 1);
-}
-
-/// Calls walk on `function`.
-NestedMatch NestedPattern::match(Function *function) {
-  assert(!matches && "NestedPattern already matched!");
-  this->walkPostOrder(function);
-  return matches;
+llvm::BumpPtrAllocator *&NestedPattern::allocator() {
+  thread_local llvm::BumpPtrAllocator *allocator = nullptr;
+  return allocator;
 }
 
-/// Calls walk on `instruction`.
-NestedMatch NestedPattern::match(Instruction *instruction) {
-  assert(!matches && "NestedPattern already matched!");
-  this->walkPostOrder(instruction);
-  return matches;
+NestedPattern::NestedPattern(Instruction::Kind k,
+                             ArrayRef<NestedPattern> nested,
+                             FilterFunctionType filter)
+    : kind(k), nestedPatterns(ArrayRef<NestedPattern>(nested)), filter(filter) {
+  auto *newNested = allocator()->Allocate<NestedPattern>(nested.size());
+  std::uninitialized_copy(nested.begin(), nested.end(), newNested);
+  nestedPatterns = ArrayRef<NestedPattern>(newNested, nested.size());
 }
 
-unsigned NestedPattern::getDepth() {
-  auto nested = getNestedPatterns();
-  if (nested.empty()) {
+unsigned NestedPattern::getDepth() const {
+  if (nestedPatterns.empty()) {
     return 1;
   }
   unsigned depth = 0;
-  for (auto c : nested) {
+  for (auto &c : nestedPatterns) {
     depth = std::max(depth, c.getDepth());
   }
   return depth + 1;
@@ -122,69 +78,39 @@ unsigned NestedPattern::getDepth() {
 ///      appended to the list of matches;
 ///   5. TODO(ntv) Optionally applies actions (lambda), in which case we will
 ///      want to traverse in post-order DFS to avoid invalidating iterators.
-void NestedPattern::matchOne(Instruction *elem) {
-  if (storage->skip == elem) {
+void NestedPattern::matchOne(Instruction *inst,
+                             SmallVectorImpl<NestedMatch> *matches) {
+  if (skip == inst) {
     return;
   }
   // Structural filter
-  if (elem->getKind() != getKind()) {
+  if (inst->getKind() != kind) {
     return;
   }
   // Local custom filter function
-  if (!getFilterFunction()(*elem)) {
+  if (!filter(*inst)) {
     return;
   }
 
-  SmallVector<NestedMatch::EntryType, 8> nestedEntries;
-  for (auto c : getNestedPatterns()) {
-    /// We create a new nestedPattern here because a matcher holds its
-    /// results. So we concretely need multiple copies of a given matcher, one
-    /// for each matching result.
-    NestedPattern nestedPattern(c);
+  if (nestedPatterns.empty()) {
+    SmallVector<NestedMatch, 8> nestedMatches;
+    matches->push_back(NestedMatch::build(inst, nestedMatches));
+    return;
+  }
+  // Take a copy of each nested pattern so we can match it.
+  for (auto nestedPattern : nestedPatterns) {
+    SmallVector<NestedMatch, 8> nestedMatches;
     // Skip elem in the walk immediately following. Without this we would
     // essentially need to reimplement walkPostOrder here.
-    nestedPattern.storage->skip = elem;
-    nestedPattern.walkPostOrder(elem);
-    if (!nestedPattern.matches) {
+    nestedPattern.skip = inst;
+    nestedPattern.match(inst, &nestedMatches);
+    // If we could not match even one of the specified nestedPattern, early exit
+    // as this whole branch is not a match.
+    if (nestedMatches.empty()) {
       return;
     }
-    for (auto m : nestedPattern.matches) {
-      nestedEntries.push_back(m);
-    }
+    matches->push_back(NestedMatch::build(inst, nestedMatches));
   }
-
-  SmallVector<NestedMatch::EntryType, 8> newEntries(
-      matches.storage->matches.begin(), matches.storage->matches.end());
-  newEntries.push_back(std::make_pair(elem, NestedMatch::build(nestedEntries)));
-  matches = NestedMatch::build(newEntries);
-}
-
-llvm::BumpPtrAllocator *&NestedPattern::allocator() {
-  static thread_local llvm::BumpPtrAllocator *allocator = nullptr;
-  return allocator;
-}
-
-NestedPattern::NestedPattern(Instruction::Kind k,
-                             ArrayRef<NestedPattern> nested,
-                             FilterFunctionType filter)
-    : storage(allocator()->Allocate<NestedPatternStorage>()),
-      matches(NestedMatch::build({})) {
-  auto *newChildren = allocator()->Allocate<NestedPattern>(nested.size());
-  std::uninitialized_copy(nested.begin(), nested.end(), newChildren);
-  // Initialize with placement new.
-  new (storage) NestedPatternStorage(
-      k, ArrayRef<NestedPattern>(newChildren, nested.size()), filter,
-      nullptr /* skip */);
-}
-
-Instruction::Kind NestedPattern::getKind() { return storage->kind; }
-
-ArrayRef<NestedPattern> NestedPattern::getNestedPatterns() {
-  return storage->nestedPatterns;
-}
-
-FilterFunctionType NestedPattern::getFilterFunction() {
-  return storage->filter;
 }
 
 static bool isAffineIfOp(const Instruction &inst) {