1 //===- UseDefAnalysis.cpp - Analysis for Transitive UseDef chains ---------===//
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 // This file implements Analysis functions specific to slicing in Function.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "mlir/Analysis/SliceAnalysis.h"
14 #include "mlir/Dialect/Affine/IR/AffineOps.h"
15 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
16 #include "mlir/Dialect/SCF/SCF.h"
17 #include "mlir/IR/BuiltinOps.h"
18 #include "mlir/IR/Operation.h"
19 #include "mlir/Support/LLVM.h"
20 #include "llvm/ADT/SetVector.h"
21 
22 ///
23 /// Implements Analysis functions specific to slicing in Function.
24 ///
25 
26 using namespace mlir;
27 
28 using llvm::SetVector;
29 
30 static void getForwardSliceImpl(Operation *op,
31                                 SetVector<Operation *> *forwardSlice,
32                                 TransitiveFilter filter) {
33   if (!op) {
34     return;
35   }
36 
37   // Evaluate whether we should keep this use.
38   // This is useful in particular to implement scoping; i.e. return the
39   // transitive forwardSlice in the current scope.
40   if (!filter(op)) {
41     return;
42   }
43 
44   if (auto forOp = dyn_cast<AffineForOp>(op)) {
45     for (Operation *userOp : forOp.getInductionVar().getUsers())
46       if (forwardSlice->count(userOp) == 0)
47         getForwardSliceImpl(userOp, forwardSlice, filter);
48   } else if (auto forOp = dyn_cast<scf::ForOp>(op)) {
49     for (Operation *userOp : forOp.getInductionVar().getUsers())
50       if (forwardSlice->count(userOp) == 0)
51         getForwardSliceImpl(userOp, forwardSlice, filter);
52     for (Value result : forOp.getResults())
53       for (Operation *userOp : result.getUsers())
54         if (forwardSlice->count(userOp) == 0)
55           getForwardSliceImpl(userOp, forwardSlice, filter);
56   } else {
57     assert(op->getNumRegions() == 0 && "unexpected generic op with regions");
58     for (Value result : op->getResults()) {
59       for (Operation *userOp : result.getUsers())
60         if (forwardSlice->count(userOp) == 0)
61           getForwardSliceImpl(userOp, forwardSlice, filter);
62     }
63   }
64 
65   forwardSlice->insert(op);
66 }
67 
68 void mlir::getForwardSlice(Operation *op, SetVector<Operation *> *forwardSlice,
69                            TransitiveFilter filter) {
70   getForwardSliceImpl(op, forwardSlice, filter);
71   // Don't insert the top level operation, we just queried on it and don't
72   // want it in the results.
73   forwardSlice->remove(op);
74 
75   // Reverse to get back the actual topological order.
76   // std::reverse does not work out of the box on SetVector and I want an
77   // in-place swap based thing (the real std::reverse, not the LLVM adapter).
78   std::vector<Operation *> v(forwardSlice->takeVector());
79   forwardSlice->insert(v.rbegin(), v.rend());
80 }
81 
82 static void getBackwardSliceImpl(Operation *op,
83                                  SetVector<Operation *> *backwardSlice,
84                                  TransitiveFilter filter) {
85   if (!op)
86     return;
87 
88   assert((op->getNumRegions() == 0 ||
89           isa<AffineForOp, scf::ForOp, linalg::LinalgOp, linalg::PadTensorOp>(
90               op)) &&
91          "unexpected generic op with regions");
92 
93   // Evaluate whether we should keep this def.
94   // This is useful in particular to implement scoping; i.e. return the
95   // transitive forwardSlice in the current scope.
96   if (!filter(op)) {
97     return;
98   }
99 
100   for (auto en : llvm::enumerate(op->getOperands())) {
101     auto operand = en.value();
102     if (auto blockArg = operand.dyn_cast<BlockArgument>()) {
103       if (auto affIv = getForInductionVarOwner(operand)) {
104         auto *affOp = affIv.getOperation();
105         if (backwardSlice->count(affOp) == 0)
106           getBackwardSliceImpl(affOp, backwardSlice, filter);
107       } else if (auto loopIv = scf::getForInductionVarOwner(operand)) {
108         auto *loopOp = loopIv.getOperation();
109         if (backwardSlice->count(loopOp) == 0)
110           getBackwardSliceImpl(loopOp, backwardSlice, filter);
111       } else if (blockArg.getOwner() !=
112                  &op->getParentOfType<FuncOp>().getBody().front()) {
113         op->emitError("unsupported CF for operand ") << en.index();
114         llvm_unreachable("Unsupported control flow");
115       }
116       continue;
117     }
118     auto *op = operand.getDefiningOp();
119     if (backwardSlice->count(op) == 0) {
120       getBackwardSliceImpl(op, backwardSlice, filter);
121     }
122   }
123 
124   backwardSlice->insert(op);
125 }
126 
127 void mlir::getBackwardSlice(Operation *op,
128                             SetVector<Operation *> *backwardSlice,
129                             TransitiveFilter filter) {
130   getBackwardSliceImpl(op, backwardSlice, filter);
131 
132   // Don't insert the top level operation, we just queried on it and don't
133   // want it in the results.
134   backwardSlice->remove(op);
135 }
136 
137 SetVector<Operation *> mlir::getSlice(Operation *op,
138                                       TransitiveFilter backwardFilter,
139                                       TransitiveFilter forwardFilter) {
140   SetVector<Operation *> slice;
141   slice.insert(op);
142 
143   unsigned currentIndex = 0;
144   SetVector<Operation *> backwardSlice;
145   SetVector<Operation *> forwardSlice;
146   while (currentIndex != slice.size()) {
147     auto *currentOp = (slice)[currentIndex];
148     // Compute and insert the backwardSlice starting from currentOp.
149     backwardSlice.clear();
150     getBackwardSlice(currentOp, &backwardSlice, backwardFilter);
151     slice.insert(backwardSlice.begin(), backwardSlice.end());
152 
153     // Compute and insert the forwardSlice starting from currentOp.
154     forwardSlice.clear();
155     getForwardSlice(currentOp, &forwardSlice, forwardFilter);
156     slice.insert(forwardSlice.begin(), forwardSlice.end());
157     ++currentIndex;
158   }
159   return topologicalSort(slice);
160 }
161 
162 namespace {
163 /// DFS post-order implementation that maintains a global count to work across
164 /// multiple invocations, to help implement topological sort on multi-root DAGs.
165 /// We traverse all operations but only record the ones that appear in
166 /// `toSort` for the final result.
167 struct DFSState {
168   DFSState(const SetVector<Operation *> &set)
169       : toSort(set), topologicalCounts(), seen() {}
170   const SetVector<Operation *> &toSort;
171   SmallVector<Operation *, 16> topologicalCounts;
172   DenseSet<Operation *> seen;
173 };
174 } // namespace
175 
176 static void DFSPostorder(Operation *current, DFSState *state) {
177   for (Value result : current->getResults()) {
178     for (Operation *op : result.getUsers())
179       DFSPostorder(op, state);
180   }
181   bool inserted;
182   using IterTy = decltype(state->seen.begin());
183   IterTy iter;
184   std::tie(iter, inserted) = state->seen.insert(current);
185   if (inserted) {
186     if (state->toSort.count(current) > 0) {
187       state->topologicalCounts.push_back(current);
188     }
189   }
190 }
191 
192 SetVector<Operation *>
193 mlir::topologicalSort(const SetVector<Operation *> &toSort) {
194   if (toSort.empty()) {
195     return toSort;
196   }
197 
198   // Run from each root with global count and `seen` set.
199   DFSState state(toSort);
200   for (auto *s : toSort) {
201     assert(toSort.count(s) == 1 && "NYI: multi-sets not supported");
202     DFSPostorder(s, &state);
203   }
204 
205   // Reorder and return.
206   SetVector<Operation *> res;
207   for (auto it = state.topologicalCounts.rbegin(),
208             eit = state.topologicalCounts.rend();
209        it != eit; ++it) {
210     res.insert(*it);
211   }
212   return res;
213 }
214