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-exception
6//
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>(Num: nestedMatches.size());
29 llvm::uninitialized_copy(Src&: nestedMatches, Dst: 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>(Num: nested.size());
47 llvm::uninitialized_copy(Src&: nested, Dst: 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(nested: other.nestedPatterns);
65}
66
67NestedPattern &NestedPattern::operator=(const NestedPattern &other) {
68 freeNested();
69 filter = other.filter;
70 skip = other.skip;
71 copyNestedToThis(nested: 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(a: depth, b: 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 then
88/// there is no match;
89/// 2. calls the customizable filter function to refine the single operation
90/// 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 is
93/// appended to the list of matches;
94/// 5. TODO: Optionally applies actions (lambda), in which case we will want
95/// 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 function
102 if (!filter(*op)) {
103 return;
104 }
105
106 if (nestedPatterns.empty()) {
107 SmallVector<NestedMatch, 8> nestedMatches;
108 matches->push_back(Elt: NestedMatch::build(operation: 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 would
115 // essentially need to reimplement walk here.
116 nestedPattern.skip = op;
117 nestedPattern.match(op, matches: &nestedMatches);
118 // If we could not match even one of the specified nestedPattern, early exit
119 // as this whole branch is not a match.
120 if (nestedMatches.empty()) {
121 return;
122 }
123 matches->push_back(Elt: NestedMatch::build(operation: op, nestedMatches));
124 }
125}
126
127static bool isAffineForOp(Operation &op) { return isa<AffineForOp>(Val: op); }
128
129static bool isAffineIfOp(Operation &op) { return isa<AffineIfOp>(Val: 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>(Val: op);
176}
177
178} // namespace matcher
179} // namespace affine
180} // namespace mlir
181

source code of mlir/lib/Dialect/Affine/Analysis/NestedMatcher.cpp