181 lines · cpp
1//===- NestedMatcher.cpp - NestedMatcher Impl ----------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include <utility>10 11#include "mlir/Dialect/Affine/Analysis/NestedMatcher.h"12#include "mlir/Dialect/Affine/IR/AffineOps.h"13 14#include "llvm/ADT/STLExtras.h"15#include "llvm/Support/Allocator.h"16 17using namespace mlir;18using namespace mlir::affine;19 20llvm::BumpPtrAllocator *&NestedMatch::allocator() {21 thread_local llvm::BumpPtrAllocator *allocator = nullptr;22 return allocator;23}24 25NestedMatch NestedMatch::build(Operation *operation,26 ArrayRef<NestedMatch> nestedMatches) {27 auto *result = allocator()->Allocate<NestedMatch>();28 auto *children = allocator()->Allocate<NestedMatch>(nestedMatches.size());29 llvm::uninitialized_copy(nestedMatches, children);30 new (result) NestedMatch();31 result->matchedOperation = operation;32 result->matchedChildren =33 ArrayRef<NestedMatch>(children, nestedMatches.size());34 return *result;35}36 37llvm::BumpPtrAllocator *&NestedPattern::allocator() {38 thread_local llvm::BumpPtrAllocator *allocator = nullptr;39 return allocator;40}41 42void NestedPattern::copyNestedToThis(ArrayRef<NestedPattern> nested) {43 if (nested.empty())44 return;45 46 auto *newNested = allocator()->Allocate<NestedPattern>(nested.size());47 llvm::uninitialized_copy(nested, newNested);48 nestedPatterns = ArrayRef<NestedPattern>(newNested, nested.size());49}50 51void NestedPattern::freeNested() {52 for (const auto &p : nestedPatterns)53 p.~NestedPattern();54}55 56NestedPattern::NestedPattern(ArrayRef<NestedPattern> nested,57 FilterFunctionType filter)58 : filter(std::move(filter)), skip(nullptr) {59 copyNestedToThis(nested);60}61 62NestedPattern::NestedPattern(const NestedPattern &other)63 : filter(other.filter), skip(other.skip) {64 copyNestedToThis(other.nestedPatterns);65}66 67NestedPattern &NestedPattern::operator=(const NestedPattern &other) {68 freeNested();69 filter = other.filter;70 skip = other.skip;71 copyNestedToThis(other.nestedPatterns);72 return *this;73}74 75unsigned NestedPattern::getDepth() const {76 if (nestedPatterns.empty()) {77 return 1;78 }79 unsigned depth = 0;80 for (auto &c : nestedPatterns) {81 depth = std::max(depth, c.getDepth());82 }83 return depth + 1;84}85 86/// Matches a single operation in the following way:87/// 1. checks the kind of operation against the matcher, if different then88/// there is no match;89/// 2. calls the customizable filter function to refine the single operation90/// match with extra semantic constraints;91/// 3. if all is good, recursively matches the nested patterns;92/// 4. if all nested match then the single operation matches too and is93/// appended to the list of matches;94/// 5. TODO: Optionally applies actions (lambda), in which case we will want95/// to traverse in post-order DFS to avoid invalidating iterators.96void NestedPattern::matchOne(Operation *op,97 SmallVectorImpl<NestedMatch> *matches) {98 if (skip == op) {99 return;100 }101 // Local custom filter function102 if (!filter(*op)) {103 return;104 }105 106 if (nestedPatterns.empty()) {107 SmallVector<NestedMatch, 8> nestedMatches;108 matches->push_back(NestedMatch::build(op, nestedMatches));109 return;110 }111 // Take a copy of each nested pattern so we can match it.112 for (auto nestedPattern : nestedPatterns) {113 SmallVector<NestedMatch, 8> nestedMatches;114 // Skip elem in the walk immediately following. Without this we would115 // essentially need to reimplement walk here.116 nestedPattern.skip = op;117 nestedPattern.match(op, &nestedMatches);118 // If we could not match even one of the specified nestedPattern, early exit119 // as this whole branch is not a match.120 if (nestedMatches.empty()) {121 return;122 }123 matches->push_back(NestedMatch::build(op, nestedMatches));124 }125}126 127static bool isAffineForOp(Operation &op) { return isa<AffineForOp>(op); }128 129static bool isAffineIfOp(Operation &op) { return isa<AffineIfOp>(op); }130 131namespace mlir {132namespace affine {133namespace matcher {134 135NestedPattern Op(FilterFunctionType filter) {136 return NestedPattern({}, std::move(filter));137}138 139NestedPattern If(const NestedPattern &child) {140 return NestedPattern(child, isAffineIfOp);141}142NestedPattern If(const FilterFunctionType &filter, const NestedPattern &child) {143 return NestedPattern(child, [filter](Operation &op) {144 return isAffineIfOp(op) && filter(op);145 });146}147NestedPattern If(ArrayRef<NestedPattern> nested) {148 return NestedPattern(nested, isAffineIfOp);149}150NestedPattern If(const FilterFunctionType &filter,151 ArrayRef<NestedPattern> nested) {152 return NestedPattern(nested, [filter](Operation &op) {153 return isAffineIfOp(op) && filter(op);154 });155}156 157NestedPattern For(const NestedPattern &child) {158 return NestedPattern(child, isAffineForOp);159}160NestedPattern For(const FilterFunctionType &filter,161 const NestedPattern &child) {162 return NestedPattern(163 child, [=](Operation &op) { return isAffineForOp(op) && filter(op); });164}165NestedPattern For(ArrayRef<NestedPattern> nested) {166 return NestedPattern(nested, isAffineForOp);167}168NestedPattern For(const FilterFunctionType &filter,169 ArrayRef<NestedPattern> nested) {170 return NestedPattern(171 nested, [=](Operation &op) { return isAffineForOp(op) && filter(op); });172}173 174bool isLoadOrStore(Operation &op) {175 return isa<AffineLoadOp, AffineStoreOp>(op);176}177 178} // namespace matcher179} // namespace affine180} // namespace mlir181