1 //===- NestedMacher.h - Nested matcher for Function -------------*- C++ -*-===//
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 #ifndef MLIR_DIALECT_AFFINE_ANALYSIS_NESTEDMATCHER_H
10 #define MLIR_DIALECT_AFFINE_ANALYSIS_NESTEDMATCHER_H
11
12 #include "mlir/IR/BuiltinOps.h"
13 #include "mlir/IR/Operation.h"
14 #include "llvm/Support/Allocator.h"
15
16 namespace mlir {
17 class Operation;
18
19 namespace affine {
20 class NestedPattern;
21
22 /// An NestedPattern captures nested patterns in the IR.
23 /// It is used in conjunction with a scoped NestedPatternContext which is an
24 /// llvm::BumpPtrAllocator that handles memory allocations efficiently and
25 /// avoids ownership issues.
26 ///
27 /// In order to use NestedPatterns, first create a scoped context.
28 /// When the context goes out of scope, everything is freed.
29 /// This design simplifies the API by avoiding references to the context and
30 /// makes it clear that references to matchers must not escape.
31 ///
32 /// Example:
33 /// {
34 /// NestedPatternContext context;
35 /// auto gemmLike = Doall(Doall(Red(LoadStores())));
36 /// auto matches = gemmLike.match(f);
37 /// // do work on matches
38 /// } // everything is freed
39 ///
40 ///
41 /// Nested abstraction for matching results.
42 /// Provides access to the nested Operation* captured by a Matcher.
43 ///
44 /// A NestedMatch contains an Operation* and the children NestedMatch and is
45 /// thus cheap to copy. NestedMatch is stored in a scoped bumper allocator whose
46 /// lifetime is managed by an RAII NestedPatternContext.
47 class NestedMatch {
48 public:
49 static NestedMatch build(Operation *operation,
50 ArrayRef<NestedMatch> nestedMatches);
51 NestedMatch(const NestedMatch &) = default;
52 NestedMatch &operator=(const NestedMatch &) = default;
53
54 explicit operator bool() { return matchedOperation != nullptr; }
55
getMatchedOperation()56 Operation *getMatchedOperation() const { return matchedOperation; }
getMatchedChildren()57 ArrayRef<NestedMatch> getMatchedChildren() { return matchedChildren; }
58
59 private:
60 friend class NestedPattern;
61 friend class NestedPatternContext;
62
63 /// Underlying global bump allocator managed by a NestedPatternContext.
64 static llvm::BumpPtrAllocator *&allocator();
65
66 NestedMatch() = default;
67
68 /// Payload, holds a NestedMatch and all its children along this branch.
69 Operation *matchedOperation = nullptr;
70 ArrayRef<NestedMatch> matchedChildren;
71 };
72
73 /// A NestedPattern is a nested operation walker that:
74 /// 1. recursively matches a substructure in the tree;
75 /// 2. uses a filter function to refine matches with extra semantic
76 /// constraints (passed via a lambda of type FilterFunctionType);
77 /// 3. TODO: optionally applies actions (lambda).
78 ///
79 /// Nested patterns are meant to capture imperfectly nested loops while matching
80 /// properties over the whole loop nest. For instance, in vectorization we are
81 /// interested in capturing all the imperfectly nested loops of a certain type
82 /// and such that all the load and stores have certain access patterns along the
83 /// loops' induction variables). Such NestedMatches are first captured using the
84 /// `match` function and are later processed to analyze properties and apply
85 /// transformations in a non-greedy way.
86 ///
87 /// The NestedMatches captured in the IR can grow large, especially after
88 /// aggressive unrolling. As experience has shown, it is generally better to use
89 /// a plain walk over operations to match flat patterns but the current
90 /// implementation is competitive nonetheless.
91 using FilterFunctionType = std::function<bool(Operation &)>;
defaultFilterFunction(Operation &)92 inline bool defaultFilterFunction(Operation &) { return true; }
93 class NestedPattern {
94 public:
95 NestedPattern(ArrayRef<NestedPattern> nested,
96 FilterFunctionType filter = defaultFilterFunction);
97 NestedPattern(const NestedPattern &other);
98 NestedPattern &operator=(const NestedPattern &other);
99
~NestedPattern()100 ~NestedPattern() {
101 // Call destructors manually, ArrayRef is non-owning so it wouldn't call
102 // them, but we should free the memory allocated by std::function outside of
103 // the arena allocator.
104 freeNested();
105 }
106
107 /// Returns all the top-level matches in `op`.
match(Operation * op,SmallVectorImpl<NestedMatch> * matches)108 void match(Operation *op, SmallVectorImpl<NestedMatch> *matches) {
109 op->walk([&](Operation *child) { matchOne(child, matches); });
110 }
111
112 /// Returns the depth of the pattern.
113 unsigned getDepth() const;
114
115 private:
116 friend class NestedPatternContext;
117 friend class NestedMatch;
118 friend struct State;
119
120 /// Copies the list of nested patterns to the arena allocator associated with
121 /// this pattern.
122 void copyNestedToThis(ArrayRef<NestedPattern> nested);
123
124 /// Calls destructors on nested patterns.
125 void freeNested();
126
127 /// Underlying global bump allocator managed by a NestedPatternContext.
128 static llvm::BumpPtrAllocator *&allocator();
129
130 /// Matches this pattern against a single `op` and fills matches with the
131 /// result.
132 void matchOne(Operation *op, SmallVectorImpl<NestedMatch> *matches);
133
134 /// Nested patterns to be matched.
135 ArrayRef<NestedPattern> nestedPatterns;
136
137 /// Extra filter function to apply to prune patterns as the IR is walked.
138 FilterFunctionType filter;
139
140 /// skip is an implementation detail needed so that we can implement match
141 /// without switching on the type of the Operation. The idea is that a
142 /// NestedPattern first checks if it matches locally and then recursively
143 /// applies its nested matchers to its elem->nested. Since we want to rely on
144 /// the existing operation walking functionality rather than duplicate
145 /// it, we allow an off-by-one traversal to account for the fact that we
146 /// write:
147 ///
148 /// void match(Operation *elem) {
149 /// for (auto &c : getNestedPatterns()) {
150 /// NestedPattern childPattern(...);
151 /// ^~~~ Needs off-by-one skip.
152 ///
153 Operation *skip;
154 };
155
156 /// RAII structure to transparently manage the bump allocator for
157 /// NestedPattern and NestedMatch classes. This avoids passing a context to
158 /// all the API functions.
159 class NestedPatternContext {
160 public:
NestedPatternContext()161 NestedPatternContext() {
162 assert(NestedMatch::allocator() == nullptr &&
163 "Only a single NestedPatternContext is supported");
164 assert(NestedPattern::allocator() == nullptr &&
165 "Only a single NestedPatternContext is supported");
166 NestedMatch::allocator() = &allocator;
167 NestedPattern::allocator() = &allocator;
168 }
~NestedPatternContext()169 ~NestedPatternContext() {
170 NestedMatch::allocator() = nullptr;
171 NestedPattern::allocator() = nullptr;
172 }
173 llvm::BumpPtrAllocator allocator;
174 };
175
176 namespace matcher {
177 // Syntactic sugar NestedPattern builder functions.
178 NestedPattern Op(FilterFunctionType filter = defaultFilterFunction);
179 NestedPattern If(const NestedPattern &child);
180 NestedPattern If(const FilterFunctionType &filter, const NestedPattern &child);
181 NestedPattern If(ArrayRef<NestedPattern> nested = {});
182 NestedPattern If(const FilterFunctionType &filter,
183 ArrayRef<NestedPattern> nested = {});
184 NestedPattern For(const NestedPattern &child);
185 NestedPattern For(const FilterFunctionType &filter, const NestedPattern &child);
186 NestedPattern For(ArrayRef<NestedPattern> nested = {});
187 NestedPattern For(const FilterFunctionType &filter,
188 ArrayRef<NestedPattern> nested = {});
189
190 bool isParallelLoop(Operation &op);
191 bool isReductionLoop(Operation &op);
192 bool isLoadOrStore(Operation &op);
193
194 } // namespace matcher
195 } // namespace affine
196 } // namespace mlir
197
198 #endif // MLIR_DIALECT_AFFINE_ANALYSIS_NESTEDMATCHER_H
199