xref: /llvm-project/mlir/include/mlir/Analysis/SliceAnalysis.h (revision d97bc388fd9ef8bc38353f93ff42d894ddc4a271)
1 //===- SliceAnalysis.h - Analysis for Transitive UseDef chains --*- 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_ANALYSIS_SLICEANALYSIS_H_
10 #define MLIR_ANALYSIS_SLICEANALYSIS_H_
11 
12 #include <functional>
13 #include <vector>
14 
15 #include "mlir/Support/LLVM.h"
16 
17 #include "llvm/ADT/SetVector.h"
18 
19 namespace mlir {
20 class BlockArgument;
21 class Operation;
22 class Value;
23 
24 struct SliceOptions {
25   /// Type of the condition to limit the propagation of transitive use-defs.
26   /// This can be used in particular to limit the propagation to a given Scope
27   /// or to avoid passing through certain types of operation in a configurable
28   /// manner.
29   using TransitiveFilter = std::function<bool(Operation *)>;
30   TransitiveFilter filter = nullptr;
31 
32   /// Include the top level op in the slice.
33   bool inclusive = false;
34 
35   // TODO: Remove this alias once downstream users are updated.
36   SliceOptions() {}
37   SliceOptions(TransitiveFilter filter) : filter(std::move(filter)) {}
38 };
39 
40 // TODO: Remove this alias once downstream users are updated.
41 using TransitiveFilter = SliceOptions::TransitiveFilter;
42 
43 struct BackwardSliceOptions : public SliceOptions {
44   using SliceOptions::SliceOptions;
45   /// When omitBlockArguments is true, the backward slice computation omits
46   /// traversing any block arguments. When omitBlockArguments is false, the
47   /// backward slice computation traverses block arguments and asserts that the
48   /// parent op has a single region with a single block.
49   bool omitBlockArguments = false;
50 
51   /// When omitUsesFromAbove is true, the backward slice computation omits
52   /// traversing values that are captured from above.
53   /// TODO: this should default to `false` after users have been updated.
54   bool omitUsesFromAbove = true;
55 };
56 
57 using ForwardSliceOptions = SliceOptions;
58 
59 /// Fills `forwardSlice` with the computed forward slice (i.e. all
60 /// the transitive uses of op), **without** including that operation.
61 ///
62 /// This additionally takes a TransitiveFilter which acts as a frontier:
63 /// when looking at uses transitively, an operation that does not pass the
64 /// filter is never propagated through. This allows in particular to carve out
65 /// the scope within a ForOp or the scope within an IfOp.
66 ///
67 /// The implementation traverses the use chains in postorder traversal for
68 /// efficiency reasons: if an operation is already in `forwardSlice`, no
69 /// need to traverse its uses again. Since use-def chains form a DAG, this
70 /// terminates.
71 ///
72 /// Upon return to the root call, `forwardSlice` is filled with a
73 /// postorder list of uses (i.e. a reverse topological order). To get a proper
74 /// topological order, we just reverse the order in `forwardSlice` before
75 /// returning.
76 ///
77 /// Example starting from node 0
78 /// ============================
79 ///
80 ///               0
81 ///    ___________|___________
82 ///    1       2      3      4
83 ///    |_______|      |______|
84 ///    |   |             |
85 ///    |   5             6
86 ///    |___|_____________|
87 ///      |               |
88 ///      7               8
89 ///      |_______________|
90 ///              |
91 ///              9
92 ///
93 /// Assuming all local orders match the numbering order:
94 /// 1. after getting back to the root getForwardSlice, `forwardSlice` may
95 ///    contain:
96 ///      {9, 7, 8, 5, 1, 2, 6, 3, 4}
97 /// 2. reversing the result of 1. gives:
98 ///      {4, 3, 6, 2, 1, 5, 8, 7, 9}
99 ///
100 void getForwardSlice(Operation *op, SetVector<Operation *> *forwardSlice,
101                      const ForwardSliceOptions &options = {});
102 
103 /// Value-rooted version of `getForwardSlice`. Return the union of all forward
104 /// slices for the uses of the value `root`.
105 void getForwardSlice(Value root, SetVector<Operation *> *forwardSlice,
106                      const ForwardSliceOptions &options = {});
107 
108 /// Fills `backwardSlice` with the computed backward slice (i.e.
109 /// all the transitive defs of op), **without** including that operation.
110 ///
111 /// This additionally takes a TransitiveFilter which acts as a frontier:
112 /// when looking at defs transitively, an operation that does not pass the
113 /// filter is never propagated through. This allows in particular to carve out
114 /// the scope within a ForOp or the scope within an IfOp.
115 ///
116 /// The implementation traverses the def chains in postorder traversal for
117 /// efficiency reasons: if an operation is already in `backwardSlice`, no
118 /// need to traverse its definitions again. Since useuse-def chains form a DAG,
119 /// this terminates.
120 ///
121 /// Upon return to the root call, `backwardSlice` is filled with a
122 /// postorder list of defs. This happens to be a topological order, from the
123 /// point of view of the use-def chains.
124 ///
125 /// Example starting from node 8
126 /// ============================
127 ///
128 ///    1       2      3      4
129 ///    |_______|      |______|
130 ///    |   |             |
131 ///    |   5             6
132 ///    |___|_____________|
133 ///      |               |
134 ///      7               8
135 ///      |_______________|
136 ///              |
137 ///              9
138 ///
139 /// Assuming all local orders match the numbering order:
140 ///    {1, 2, 5, 3, 4, 6}
141 ///
142 void getBackwardSlice(Operation *op, SetVector<Operation *> *backwardSlice,
143                       const BackwardSliceOptions &options = {});
144 
145 /// Value-rooted version of `getBackwardSlice`. Return the union of all backward
146 /// slices for the op defining or owning the value `root`.
147 void getBackwardSlice(Value root, SetVector<Operation *> *backwardSlice,
148                       const BackwardSliceOptions &options = {});
149 
150 /// Iteratively computes backward slices and forward slices until
151 /// a fixed point is reached. Returns an `SetVector<Operation *>` which
152 /// **includes** the original operation.
153 ///
154 /// This allows building a slice (i.e. multi-root DAG where everything
155 /// that is reachable from an Value in forward and backward direction is
156 /// contained in the slice).
157 /// This is the abstraction we need to materialize all the operations for
158 /// supervectorization without worrying about orderings and Value
159 /// replacements.
160 ///
161 /// Example starting from any node
162 /// ==============================
163 ///
164 ///    1       2      3      4
165 ///    |_______|      |______|
166 ///    |   |             |   |
167 ///    |   5             6___|
168 ///    |___|_____________|   |
169 ///      |               |   |
170 ///      7               8   |
171 ///      |_______________|   |
172 ///              |           |
173 ///              9          10
174 ///
175 /// Return the whole DAG in some topological order.
176 ///
177 /// The implementation works by just filling up a worklist with iterative
178 /// alternate calls to `getBackwardSlice` and `getForwardSlice`.
179 ///
180 /// The following section describes some additional implementation
181 /// considerations for a potentially more efficient implementation but they are
182 /// just an intuition without proof, we still use a worklist for now.
183 ///
184 /// Additional implementation considerations
185 /// ========================================
186 /// Consider the defs-op-uses hourglass.
187 ///    ____
188 ///    \  /  defs (in some topological order)
189 ///     \/
190 ///     op
191 ///     /\
192 ///    /  \  uses (in some topological order)
193 ///   /____\
194 ///
195 /// We want to iteratively apply `getSlice` to construct the whole
196 /// list of Operation that are reachable by (use|def)+ from op.
197 /// We want the resulting slice in topological order.
198 /// Ideally we would like the ordering to be maintained in-place to avoid
199 /// copying Operation at each step. Keeping this ordering by construction
200 /// seems very unclear, so we list invariants in the hope of seeing whether
201 /// useful properties pop up.
202 ///
203 /// In the following:
204 ///   we use |= for set inclusion;
205 ///   we use << for set topological ordering (i.e. each pair is ordered).
206 ///
207 /// Assumption:
208 /// ===========
209 /// We wish to maintain the following property by a recursive argument:
210 ///   """
211 ///      defs << {op} <<uses are in topological order.
212 ///   """
213 /// The property clearly holds for 0 and 1-sized uses and defs;
214 ///
215 /// Invariants:
216 ///   2. defs and uses are in topological order internally, by construction;
217 ///   3. for any {x} |= defs, defs(x) |= defs;    because all go through op
218 ///   4. for any {x} |= uses,    defs |= defs(x); because all go through op
219 ///   5. for any {x} |= defs,    uses |= uses(x); because all go through op
220 ///   6. for any {x} |= uses, uses(x) |= uses;    because all go through op
221 ///
222 /// Intuitively, we should be able to recurse like:
223 ///   preorder(defs) - op - postorder(uses)
224 /// and keep things ordered but this is still hand-wavy and not worth the
225 /// trouble for now: punt to a simple worklist-based solution.
226 ///
227 SetVector<Operation *>
228 getSlice(Operation *op, const BackwardSliceOptions &backwardSliceOptions = {},
229          const ForwardSliceOptions &forwardSliceOptions = {});
230 
231 /// Utility to match a generic reduction given a list of iteration-carried
232 /// arguments, `iterCarriedArgs` and the position of the potential reduction
233 /// argument within the list, `redPos`. If a reduction is matched, returns the
234 /// reduced value and the topologically-sorted list of combiner operations
235 /// involved in the reduction. Otherwise, returns a null value.
236 ///
237 /// The matching algorithm relies on the following invariants, which are subject
238 /// to change:
239 ///  1. The first combiner operation must be a binary operation with the
240 ///     iteration-carried value and the reduced value as operands.
241 ///  2. The iteration-carried value and combiner operations must be side
242 ///     effect-free, have single result and a single use.
243 ///  3. Combiner operations must be immediately nested in the region op
244 ///     performing the reduction.
245 ///  4. Reduction def-use chain must end in a terminator op that yields the
246 ///     next iteration/output values in the same order as the iteration-carried
247 ///     values in `iterCarriedArgs`.
248 ///  5. `iterCarriedArgs` must contain all the iteration-carried/output values
249 ///     of the region op performing the reduction.
250 ///
251 /// This utility is generic enough to detect reductions involving multiple
252 /// combiner operations (disabled for now) across multiple dialects, including
253 /// Linalg, Affine and SCF. For the sake of genericity, it does not return
254 /// specific enum values for the combiner operations since its goal is also
255 /// matching reductions without pre-defined semantics in core MLIR. It's up to
256 /// each client to make sense out of the list of combiner operations. It's also
257 /// up to each client to check for additional invariants on the expected
258 /// reductions not covered by this generic matching.
259 Value matchReduction(ArrayRef<BlockArgument> iterCarriedArgs, unsigned redPos,
260                      SmallVectorImpl<Operation *> &combinerOps);
261 
262 } // namespace mlir
263 
264 #endif // MLIR_ANALYSIS_SLICEANALYSIS_H_
265