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