1 //===- DependenceAnalysis.cpp - Dependence analysis on SSA views ----------===//
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 view-based alias and dependence analyses.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
14 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
15 #include "mlir/Dialect/MemRef/IR/MemRef.h"
16 #include "mlir/Dialect/StandardOps/IR/Ops.h"
17 #include "mlir/IR/BuiltinOps.h"
18 
19 #include "llvm/Support/CommandLine.h"
20 #include "llvm/Support/Debug.h"
21 
22 #define DEBUG_TYPE "linalg-dependence-analysis"
23 
24 using namespace mlir;
25 using namespace mlir::linalg;
26 
27 using llvm::dbgs;
28 
29 Value Aliases::find(Value v) {
30   if (v.isa<BlockArgument>())
31     return v;
32 
33   auto it = aliases.find(v);
34   if (it != aliases.end()) {
35     assert(it->getSecond().getType().isa<BaseMemRefType>() &&
36            "Memref expected");
37     return it->getSecond();
38   }
39 
40   while (true) {
41     if (v.isa<BlockArgument>())
42       return v;
43 
44     Operation *defOp = v.getDefiningOp();
45     if (!defOp)
46       return v;
47 
48     // Treat RegionBranchOpInterfaces like an allocate and don't try to follow
49     // the aliasing further.
50     if (isa<RegionBranchOpInterface>(defOp))
51       return v;
52     if (isa<memref::BufferCastOp>(defOp))
53       return v;
54 
55     if (auto memEffect = dyn_cast<MemoryEffectOpInterface>(defOp)) {
56       // Collect all memory effects on `v`.
57       SmallVector<MemoryEffects::EffectInstance, 1> effects;
58       memEffect.getEffectsOnValue(v, effects);
59 
60       // If we have the 'Allocate' memory effect on `v`, then `v` should be the
61       // original buffer.
62       if (llvm::any_of(
63               effects, [](const MemoryEffects::EffectInstance &instance) {
64                 return isa<MemoryEffects::Allocate>(instance.getEffect());
65               }))
66         return v;
67     }
68 
69     if (auto viewLikeOp = dyn_cast<ViewLikeOpInterface>(defOp)) {
70       auto it =
71           aliases.insert(std::make_pair(v, find(viewLikeOp.getViewSource())));
72       return it.first->second;
73     }
74 
75     llvm::errs() << "View alias analysis reduces to: " << v << "\n";
76     llvm_unreachable("unsupported view alias case");
77   }
78 }
79 
80 StringRef LinalgDependenceGraph::getDependenceTypeStr(DependenceType depType) {
81   switch (depType) {
82   case LinalgDependenceGraph::DependenceType::RAW:
83     return "RAW";
84   case LinalgDependenceGraph::DependenceType::RAR:
85     return "RAR";
86   case LinalgDependenceGraph::DependenceType::WAR:
87     return "WAR";
88   case LinalgDependenceGraph::DependenceType::WAW:
89     return "WAW";
90   default:
91     break;
92   }
93   llvm_unreachable("Unexpected DependenceType");
94 }
95 
96 LinalgDependenceGraph
97 LinalgDependenceGraph::buildDependenceGraph(Aliases &aliases, FuncOp f) {
98   SmallVector<LinalgOp, 8> linalgOps;
99   f.walk([&](LinalgOp op) { linalgOps.push_back(op); });
100   return LinalgDependenceGraph(aliases, linalgOps);
101 }
102 
103 LinalgDependenceGraph::LinalgDependenceGraph(Aliases &aliases,
104                                              ArrayRef<LinalgOp> ops)
105     : aliases(aliases), linalgOps(ops.begin(), ops.end()) {
106   for (auto en : llvm::enumerate(linalgOps)) {
107     linalgOpPositions.insert(
108         std::make_pair(en.value().getOperation(), en.index()));
109   }
110   for (unsigned i = 0, e = ops.size(); i < e; ++i) {
111     for (unsigned j = i + 1; j < e; ++j) {
112       addDependencesBetween(ops[i], ops[j]);
113     }
114   }
115 }
116 
117 void LinalgDependenceGraph::addDependenceElem(
118     DependenceType dt, LinalgDependenceGraphElem::OpView indexingOpView,
119     LinalgDependenceGraphElem::OpView dependentOpView) {
120   LLVM_DEBUG(dbgs() << "\nAdd dep type " << getDependenceTypeStr(dt) << ":\t ("
121                     << LinalgDependenceGraphElem::getValue(indexingOpView)
122                     << " @) -> \n\t\t("
123                     << LinalgDependenceGraphElem::getValue(dependentOpView)
124                     << " @)");
125   dependencesFromGraphs[dt][LinalgDependenceGraphElem::getOwner(indexingOpView)]
126       .push_back(
127           LinalgDependenceGraphElem{dependentOpView, indexingOpView, dt});
128   dependencesIntoGraphs[dt]
129                        [LinalgDependenceGraphElem::getOwner(dependentOpView)]
130                            .push_back(LinalgDependenceGraphElem{
131                                indexingOpView, dependentOpView, dt});
132 }
133 
134 LinalgDependenceGraph::dependence_range
135 LinalgDependenceGraph::getDependencesFrom(
136     LinalgOp src, LinalgDependenceGraph::DependenceType dt) const {
137   return getDependencesFrom(src.getOperation(), dt);
138 }
139 
140 LinalgDependenceGraph::dependence_range
141 LinalgDependenceGraph::getDependencesFrom(
142     Operation *src, LinalgDependenceGraph::DependenceType dt) const {
143   auto iter = dependencesFromGraphs[dt].find(src);
144   if (iter == dependencesFromGraphs[dt].end())
145     return llvm::make_range(nullptr, nullptr);
146   return llvm::make_range(iter->second.begin(), iter->second.end());
147 }
148 
149 LinalgDependenceGraph::dependence_range
150 LinalgDependenceGraph::getDependencesInto(
151     LinalgOp dst, LinalgDependenceGraph::DependenceType dt) const {
152   return getDependencesInto(dst.getOperation(), dt);
153 }
154 
155 LinalgDependenceGraph::dependence_range
156 LinalgDependenceGraph::getDependencesInto(
157     Operation *dst, LinalgDependenceGraph::DependenceType dt) const {
158   auto iter = dependencesIntoGraphs[dt].find(dst);
159   if (iter == dependencesIntoGraphs[dt].end())
160     return llvm::make_range(nullptr, nullptr);
161   return llvm::make_range(iter->second.begin(), iter->second.end());
162 }
163 
164 void LinalgDependenceGraph::addDependencesBetween(LinalgOp src, LinalgOp dst) {
165   if (src.hasTensorSemantics() && dst.hasTensorSemantics()) {
166     for (OpOperand &dstOpOperand : dst.getInputOpOperands()) {
167       // Check if the operand is defined by the src.
168       auto definingOp = dstOpOperand.get().getDefiningOp<LinalgOp>();
169       if (definingOp && definingOp == src)
170         addDependenceElem(DependenceType::RAW, dstOpOperand.get(),
171                           &dstOpOperand);
172     }
173     return;
174   }
175   assert(src.hasBufferSemantics() && dst.hasBufferSemantics() &&
176          "unhandled dependence tracking for mixed buffer/tensor operations");
177   for (OpOperand *srcOpOperand : src.getOutputBuffersOpOperands()) { // W
178     // RAW graph
179     for (OpOperand *dstOpOperand : dst.getInputBuffersOpOperands()) // R
180       if (aliases.alias(srcOpOperand->get(), dstOpOperand->get()))  // RAW alias
181         addDependenceElem(DependenceType::RAW, srcOpOperand, dstOpOperand);
182     // WAW graph
183     for (OpOperand *dstOpOperand : dst.getOutputBuffersOpOperands()) // W
184       if (aliases.alias(srcOpOperand->get(), dstOpOperand->get())) // WAW alias
185         addDependenceElem(DependenceType::WAW, srcOpOperand, dstOpOperand);
186   }
187   for (OpOperand *srcOpOperand : src.getInputBuffersOpOperands()) { // R
188     // RAR graph
189     for (OpOperand *dstOpOperand : dst.getInputBuffersOpOperands()) // R
190       if (aliases.alias(srcOpOperand->get(), dstOpOperand->get()))  // RAR alias
191         addDependenceElem(DependenceType::RAR, srcOpOperand, dstOpOperand);
192     // WAR graph
193     for (OpOperand *dstOpOperand : dst.getOutputBuffersOpOperands()) // W
194       if (aliases.alias(srcOpOperand->get(), dstOpOperand->get())) // WAR alias
195         addDependenceElem(DependenceType::WAR, srcOpOperand, dstOpOperand);
196   }
197 }
198 
199 SmallVector<Operation *, 8>
200 LinalgDependenceGraph::findCoveringDependences(LinalgOp srcLinalgOp,
201                                                LinalgOp dstLinalgOp) const {
202   return findOperationsWithCoveringDependences(
203       srcLinalgOp, dstLinalgOp, nullptr,
204       {DependenceType::WAW, DependenceType::WAR, DependenceType::RAW});
205 }
206 
207 SmallVector<Operation *, 8> LinalgDependenceGraph::findCoveringWrites(
208     LinalgOp srcLinalgOp, LinalgOp dstLinalgOp, Value view) const {
209   return findOperationsWithCoveringDependences(
210       srcLinalgOp, dstLinalgOp, view,
211       {DependenceType::WAW, DependenceType::WAR});
212 }
213 
214 SmallVector<Operation *, 8> LinalgDependenceGraph::findCoveringReads(
215     LinalgOp srcLinalgOp, LinalgOp dstLinalgOp, Value view) const {
216   return findOperationsWithCoveringDependences(
217       srcLinalgOp, dstLinalgOp, view,
218       {DependenceType::RAR, DependenceType::RAW});
219 }
220 
221 SmallVector<Operation *, 8>
222 LinalgDependenceGraph::findOperationsWithCoveringDependences(
223     LinalgOp srcLinalgOp, LinalgOp dstLinalgOp, Value view,
224     ArrayRef<DependenceType> types) const {
225   auto *src = srcLinalgOp.getOperation();
226   auto *dst = dstLinalgOp.getOperation();
227   auto srcPos = linalgOpPositions.lookup(src);
228   auto dstPos = linalgOpPositions.lookup(dst);
229   assert(srcPos < dstPos && "expected dst after src in IR traversal order");
230 
231   SmallVector<Operation *, 8> res;
232   // Consider an intermediate interleaved `interim` op, look for any dependence
233   // to an aliasing view on a src -> op -> dst path.
234   // TODO: we are not considering paths yet, just interleaved positions.
235   for (auto dt : types) {
236     for (auto dependence : getDependencesFrom(src, dt)) {
237       auto interimPos = linalgOpPositions.lookup(dependence.getDependentOp());
238       // Skip if not interleaved.
239       if (interimPos >= dstPos || interimPos <= srcPos)
240         continue;
241       Value consumerView = dependence.getIndexingValue();
242       if (view && !aliases.alias(view, consumerView))
243         continue;
244       auto *op = dependence.getDependentOp();
245       LLVM_DEBUG(dbgs() << "\n***Found covering dependence of type "
246                         << getDependenceTypeStr(dt) << ": " << *src << " -> "
247                         << *op << " on " << consumerView);
248       res.push_back(op);
249     }
250   }
251   return res;
252 }
253 
254 bool LinalgDependenceGraph::hasDependenceFrom(
255     LinalgOp srcLinalgOp, LinalgOp dstLinalgOp,
256     ArrayRef<LinalgDependenceGraph::DependenceType> depTypes) const {
257   for (auto dep : depTypes)
258     for (auto dependence : getDependencesInto(dstLinalgOp, dep))
259       if (dependence.getDependentOp() == srcLinalgOp)
260         return true;
261   return false;
262 }
263 
264 bool LinalgDependenceGraph::hasDependentOperationsFrom(
265     LinalgOp linalgOp,
266     ArrayRef<LinalgDependenceGraph::DependenceType> depTypes) const {
267   for (auto dep : depTypes) {
268     if (!getDependencesFrom(linalgOp, dep).empty())
269       return true;
270   }
271   return false;
272 }
273 
274 bool LinalgDependenceGraph::hasDependentOperationsInto(
275     LinalgOp linalgOp,
276     ArrayRef<LinalgDependenceGraph::DependenceType> depTypes) const {
277   for (auto dep : depTypes) {
278     if (!getDependencesInto(linalgOp, dep).empty())
279       return true;
280   }
281   return false;
282 }
283 
284 bool LinalgDependenceGraph::hasDependentOperations(
285     LinalgOp linalgOp, ArrayRef<DependenceType> depTypes) const {
286   return hasDependentOperationsInto(linalgOp, depTypes) ||
287          hasDependentOperationsFrom(linalgOp, depTypes);
288 }
289 
290 SmallVector<LinalgDependenceGraph::LinalgDependenceGraphElem, 2>
291 LinalgDependenceGraph::getDependentOperationsInto(
292     LinalgOp linalgOp, ArrayRef<DependenceType> depTypes) const {
293   SmallVector<LinalgDependenceGraph::LinalgDependenceGraphElem, 2>
294       dependentOperations;
295   for (auto dependenceType : depTypes) {
296     auto dependencies = getDependencesInto(linalgOp, dependenceType);
297     dependentOperations.append(dependencies.begin(), dependencies.end());
298   }
299   return dependentOperations;
300 }
301 
302 SmallVector<LinalgDependenceGraph::LinalgDependenceGraphElem, 2>
303 LinalgDependenceGraph::getDependentOperationsFrom(
304     LinalgOp linalgOp, ArrayRef<DependenceType> depTypes) const {
305   SmallVector<LinalgDependenceGraph::LinalgDependenceGraphElem, 2>
306       dependentOperations;
307   for (auto dependenceType : depTypes) {
308     auto dependencies = getDependencesFrom(linalgOp, dependenceType);
309     dependentOperations.append(dependencies.begin(), dependencies.end());
310   }
311   return dependentOperations;
312 }
313 
314 /// Returns all dependent operations (into and from) given `operation`.
315 SmallVector<LinalgDependenceGraph::LinalgDependenceGraphElem, 2>
316 LinalgDependenceGraph::getDependentOperations(
317     LinalgOp linalgOp, ArrayRef<DependenceType> depTypes) const {
318   SmallVector<LinalgDependenceGraphElem, 2> dependentOperations =
319       getDependentOperationsInto(linalgOp, depTypes);
320   SmallVector<LinalgDependenceGraphElem, 2> t =
321       getDependentOperationsFrom(linalgOp, depTypes);
322   dependentOperations.append(t.begin(), t.end());
323   return dependentOperations;
324 }
325