From d4921f4a96a46e24c2fb76ca7feb66e6894395d2 Mon Sep 17 00:00:00 2001 From: Nicolas Vasilache Date: Thu, 31 Jan 2019 07:16:29 -0800 Subject: 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 --- mlir/lib/Analysis/NestedMatcher.cpp | 158 ++++++++++-------------------------- 1 file changed, 42 insertions(+), 116 deletions(-) (limited to 'mlir/lib/Analysis/NestedMatcher.cpp') 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 matches; -}; - -/// Underlying storage for NestedPattern. -struct NestedPatternStorage { - NestedPatternStorage(Instruction::Kind k, ArrayRef c, - FilterFunctionType filter, Instruction *skip) - : kind(k), nestedPatterns(c), filter(filter), skip(skip) {} - - Instruction::Kind kind; - ArrayRef 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 elements) { - auto *matches = - allocator()->Allocate(elements.size()); - std::uninitialized_copy(elements.begin(), elements.end(), matches); - auto *storage = allocator()->Allocate(); - new (storage) NestedMatchStorage(); - storage->matches = - MutableArrayRef(matches, elements.size()); +NestedMatch NestedMatch::build(Instruction *instruction, + ArrayRef nestedMatches) { auto *result = allocator()->Allocate(); - new (result) NestedMatch(storage); + auto *children = allocator()->Allocate(nestedMatches.size()); + std::uninitialized_copy(nestedMatches.begin(), nestedMatches.end(), children); + new (result) NestedMatch(); + result->matchedInstruction = instruction; + result->matchedChildren = + ArrayRef(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 nested, + FilterFunctionType filter) + : kind(k), nestedPatterns(ArrayRef(nested)), filter(filter) { + auto *newNested = allocator()->Allocate(nested.size()); + std::uninitialized_copy(nested.begin(), nested.end(), newNested); + nestedPatterns = ArrayRef(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 *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 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 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 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 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 nested, - FilterFunctionType filter) - : storage(allocator()->Allocate()), - matches(NestedMatch::build({})) { - auto *newChildren = allocator()->Allocate(nested.size()); - std::uninitialized_copy(nested.begin(), nested.end(), newChildren); - // Initialize with placement new. - new (storage) NestedPatternStorage( - k, ArrayRef(newChildren, nested.size()), filter, - nullptr /* skip */); -} - -Instruction::Kind NestedPattern::getKind() { return storage->kind; } - -ArrayRef NestedPattern::getNestedPatterns() { - return storage->nestedPatterns; -} - -FilterFunctionType NestedPattern::getFilterFunction() { - return storage->filter; } static bool isAffineIfOp(const Instruction &inst) { -- cgit v1.2.3