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 | |
52 | using ForwardSliceOptions = SliceOptions; |
53 | |
54 | /// Fills `forwardSlice` with the computed forward slice (i.e. all |
55 | /// the transitive uses of op), **without** including that operation. |
56 | /// |
57 | /// This additionally takes a TransitiveFilter which acts as a frontier: |
58 | /// when looking at uses transitively, an operation that does not pass the |
59 | /// filter is never propagated through. This allows in particular to carve out |
60 | /// the scope within a ForOp or the scope within an IfOp. |
61 | /// |
62 | /// The implementation traverses the use chains in postorder traversal for |
63 | /// efficiency reasons: if an operation is already in `forwardSlice`, no |
64 | /// need to traverse its uses again. Since use-def chains form a DAG, this |
65 | /// terminates. |
66 | /// |
67 | /// Upon return to the root call, `forwardSlice` is filled with a |
68 | /// postorder list of uses (i.e. a reverse topological order). To get a proper |
69 | /// topological order, we just reverse the order in `forwardSlice` before |
70 | /// returning. |
71 | /// |
72 | /// Example starting from node 0 |
73 | /// ============================ |
74 | /// |
75 | /// 0 |
76 | /// ___________|___________ |
77 | /// 1 2 3 4 |
78 | /// |_______| |______| |
79 | /// | | | |
80 | /// | 5 6 |
81 | /// |___|_____________| |
82 | /// | | |
83 | /// 7 8 |
84 | /// |_______________| |
85 | /// | |
86 | /// 9 |
87 | /// |
88 | /// Assuming all local orders match the numbering order: |
89 | /// 1. after getting back to the root getForwardSlice, `forwardSlice` may |
90 | /// contain: |
91 | /// {9, 7, 8, 5, 1, 2, 6, 3, 4} |
92 | /// 2. reversing the result of 1. gives: |
93 | /// {4, 3, 6, 2, 1, 5, 8, 7, 9} |
94 | /// |
95 | void getForwardSlice(Operation *op, SetVector<Operation *> *forwardSlice, |
96 | const ForwardSliceOptions &options = {}); |
97 | |
98 | /// Value-rooted version of `getForwardSlice`. Return the union of all forward |
99 | /// slices for the uses of the value `root`. |
100 | void getForwardSlice(Value root, SetVector<Operation *> *forwardSlice, |
101 | const ForwardSliceOptions &options = {}); |
102 | |
103 | /// Fills `backwardSlice` with the computed backward slice (i.e. |
104 | /// all the transitive defs of op), **without** including that operation. |
105 | /// |
106 | /// This additionally takes a TransitiveFilter which acts as a frontier: |
107 | /// when looking at defs transitively, an operation that does not pass the |
108 | /// filter is never propagated through. This allows in particular to carve out |
109 | /// the scope within a ForOp or the scope within an IfOp. |
110 | /// |
111 | /// The implementation traverses the def chains in postorder traversal for |
112 | /// efficiency reasons: if an operation is already in `backwardSlice`, no |
113 | /// need to traverse its definitions again. Since useuse-def chains form a DAG, |
114 | /// this terminates. |
115 | /// |
116 | /// Upon return to the root call, `backwardSlice` is filled with a |
117 | /// postorder list of defs. This happens to be a topological order, from the |
118 | /// point of view of the use-def chains. |
119 | /// |
120 | /// Example starting from node 8 |
121 | /// ============================ |
122 | /// |
123 | /// 1 2 3 4 |
124 | /// |_______| |______| |
125 | /// | | | |
126 | /// | 5 6 |
127 | /// |___|_____________| |
128 | /// | | |
129 | /// 7 8 |
130 | /// |_______________| |
131 | /// | |
132 | /// 9 |
133 | /// |
134 | /// Assuming all local orders match the numbering order: |
135 | /// {1, 2, 5, 3, 4, 6} |
136 | /// |
137 | void getBackwardSlice(Operation *op, SetVector<Operation *> *backwardSlice, |
138 | const BackwardSliceOptions &options = {}); |
139 | |
140 | /// Value-rooted version of `getBackwardSlice`. Return the union of all backward |
141 | /// slices for the op defining or owning the value `root`. |
142 | void getBackwardSlice(Value root, SetVector<Operation *> *backwardSlice, |
143 | const BackwardSliceOptions &options = {}); |
144 | |
145 | /// Iteratively computes backward slices and forward slices until |
146 | /// a fixed point is reached. Returns an `SetVector<Operation *>` which |
147 | /// **includes** the original operation. |
148 | /// |
149 | /// This allows building a slice (i.e. multi-root DAG where everything |
150 | /// that is reachable from an Value in forward and backward direction is |
151 | /// contained in the slice). |
152 | /// This is the abstraction we need to materialize all the operations for |
153 | /// supervectorization without worrying about orderings and Value |
154 | /// replacements. |
155 | /// |
156 | /// Example starting from any node |
157 | /// ============================== |
158 | /// |
159 | /// 1 2 3 4 |
160 | /// |_______| |______| |
161 | /// | | | | |
162 | /// | 5 6___| |
163 | /// |___|_____________| | |
164 | /// | | | |
165 | /// 7 8 | |
166 | /// |_______________| | |
167 | /// | | |
168 | /// 9 10 |
169 | /// |
170 | /// Return the whole DAG in some topological order. |
171 | /// |
172 | /// The implementation works by just filling up a worklist with iterative |
173 | /// alternate calls to `getBackwardSlice` and `getForwardSlice`. |
174 | /// |
175 | /// The following section describes some additional implementation |
176 | /// considerations for a potentially more efficient implementation but they are |
177 | /// just an intuition without proof, we still use a worklist for now. |
178 | /// |
179 | /// Additional implementation considerations |
180 | /// ======================================== |
181 | /// Consider the defs-op-uses hourglass. |
182 | /// ____ |
183 | /// \ / defs (in some topological order) |
184 | /// \/ |
185 | /// op |
186 | /// /\ |
187 | /// / \ uses (in some topological order) |
188 | /// /____\ |
189 | /// |
190 | /// We want to iteratively apply `getSlice` to construct the whole |
191 | /// list of Operation that are reachable by (use|def)+ from op. |
192 | /// We want the resulting slice in topological order. |
193 | /// Ideally we would like the ordering to be maintained in-place to avoid |
194 | /// copying Operation at each step. Keeping this ordering by construction |
195 | /// seems very unclear, so we list invariants in the hope of seeing whether |
196 | /// useful properties pop up. |
197 | /// |
198 | /// In the following: |
199 | /// we use |= for set inclusion; |
200 | /// we use << for set topological ordering (i.e. each pair is ordered). |
201 | /// |
202 | /// Assumption: |
203 | /// =========== |
204 | /// We wish to maintain the following property by a recursive argument: |
205 | /// """ |
206 | /// defs << {op} <<uses are in topological order. |
207 | /// """ |
208 | /// The property clearly holds for 0 and 1-sized uses and defs; |
209 | /// |
210 | /// Invariants: |
211 | /// 2. defs and uses are in topological order internally, by construction; |
212 | /// 3. for any {x} |= defs, defs(x) |= defs; because all go through op |
213 | /// 4. for any {x} |= uses, defs |= defs(x); because all go through op |
214 | /// 5. for any {x} |= defs, uses |= uses(x); because all go through op |
215 | /// 6. for any {x} |= uses, uses(x) |= uses; because all go through op |
216 | /// |
217 | /// Intuitively, we should be able to recurse like: |
218 | /// preorder(defs) - op - postorder(uses) |
219 | /// and keep things ordered but this is still hand-wavy and not worth the |
220 | /// trouble for now: punt to a simple worklist-based solution. |
221 | /// |
222 | SetVector<Operation *> |
223 | getSlice(Operation *op, const BackwardSliceOptions &backwardSliceOptions = {}, |
224 | const ForwardSliceOptions &forwardSliceOptions = {}); |
225 | |
226 | /// Multi-root DAG topological sort. |
227 | /// Performs a topological sort of the Operation in the `toSort` SetVector. |
228 | /// Returns a topologically sorted SetVector. |
229 | SetVector<Operation *> topologicalSort(const SetVector<Operation *> &toSort); |
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 | |