1 | //===- LocalAliasAnalysis.cpp - Local stateless alias Analysis for MLIR ---===// |
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 | #include "mlir/Analysis/AliasAnalysis/LocalAliasAnalysis.h" |
10 | |
11 | #include "mlir/Analysis/AliasAnalysis.h" |
12 | #include "mlir/IR/Attributes.h" |
13 | #include "mlir/IR/Block.h" |
14 | #include "mlir/IR/Matchers.h" |
15 | #include "mlir/IR/OpDefinition.h" |
16 | #include "mlir/IR/Operation.h" |
17 | #include "mlir/IR/Region.h" |
18 | #include "mlir/IR/Value.h" |
19 | #include "mlir/IR/ValueRange.h" |
20 | #include "mlir/Interfaces/ControlFlowInterfaces.h" |
21 | #include "mlir/Interfaces/FunctionInterfaces.h" |
22 | #include "mlir/Interfaces/SideEffectInterfaces.h" |
23 | #include "mlir/Interfaces/ViewLikeInterface.h" |
24 | #include "mlir/Support/LLVM.h" |
25 | #include "mlir/Support/LogicalResult.h" |
26 | #include "llvm/Support/Casting.h" |
27 | #include <cassert> |
28 | #include <optional> |
29 | #include <utility> |
30 | |
31 | using namespace mlir; |
32 | |
33 | //===----------------------------------------------------------------------===// |
34 | // Underlying Address Computation |
35 | //===----------------------------------------------------------------------===// |
36 | |
37 | /// The maximum depth that will be searched when trying to find an underlying |
38 | /// value. |
39 | static constexpr unsigned maxUnderlyingValueSearchDepth = 10; |
40 | |
41 | /// Given a value, collect all of the underlying values being addressed. |
42 | static void collectUnderlyingAddressValues(Value value, unsigned maxDepth, |
43 | DenseSet<Value> &visited, |
44 | SmallVectorImpl<Value> &output); |
45 | |
46 | /// Given a successor (`region`) of a RegionBranchOpInterface, collect all of |
47 | /// the underlying values being addressed by one of the successor inputs. If the |
48 | /// provided `region` is null, as per `RegionBranchOpInterface` this represents |
49 | /// the parent operation. |
50 | static void collectUnderlyingAddressValues(RegionBranchOpInterface branch, |
51 | Region *region, Value inputValue, |
52 | unsigned inputIndex, |
53 | unsigned maxDepth, |
54 | DenseSet<Value> &visited, |
55 | SmallVectorImpl<Value> &output) { |
56 | // Given the index of a region of the branch (`predIndex`), or std::nullopt to |
57 | // represent the parent operation, try to return the index into the outputs of |
58 | // this region predecessor that correspond to the input values of `region`. If |
59 | // an index could not be found, std::nullopt is returned instead. |
60 | auto getOperandIndexIfPred = |
61 | [&](RegionBranchPoint pred) -> std::optional<unsigned> { |
62 | SmallVector<RegionSuccessor, 2> successors; |
63 | branch.getSuccessorRegions(pred, successors); |
64 | for (RegionSuccessor &successor : successors) { |
65 | if (successor.getSuccessor() != region) |
66 | continue; |
67 | // Check that the successor inputs map to the given input value. |
68 | ValueRange inputs = successor.getSuccessorInputs(); |
69 | if (inputs.empty()) { |
70 | output.push_back(Elt: inputValue); |
71 | break; |
72 | } |
73 | unsigned firstInputIndex, lastInputIndex; |
74 | if (region) { |
75 | firstInputIndex = cast<BlockArgument>(Val: inputs[0]).getArgNumber(); |
76 | lastInputIndex = cast<BlockArgument>(Val: inputs.back()).getArgNumber(); |
77 | } else { |
78 | firstInputIndex = cast<OpResult>(Val: inputs[0]).getResultNumber(); |
79 | lastInputIndex = cast<OpResult>(Val: inputs.back()).getResultNumber(); |
80 | } |
81 | if (firstInputIndex > inputIndex || lastInputIndex < inputIndex) { |
82 | output.push_back(Elt: inputValue); |
83 | break; |
84 | } |
85 | return inputIndex - firstInputIndex; |
86 | } |
87 | return std::nullopt; |
88 | }; |
89 | |
90 | // Check branches from the parent operation. |
91 | auto branchPoint = RegionBranchPoint::parent(); |
92 | if (region) |
93 | branchPoint = region; |
94 | |
95 | if (std::optional<unsigned> operandIndex = |
96 | getOperandIndexIfPred(/*predIndex=*/RegionBranchPoint::parent())) { |
97 | collectUnderlyingAddressValues( |
98 | branch.getEntrySuccessorOperands(branchPoint)[*operandIndex], maxDepth, |
99 | visited, output); |
100 | } |
101 | // Check branches from each child region. |
102 | Operation *op = branch.getOperation(); |
103 | for (Region ®ion : op->getRegions()) { |
104 | if (std::optional<unsigned> operandIndex = getOperandIndexIfPred(region)) { |
105 | for (Block &block : region) { |
106 | // Try to determine possible region-branch successor operands for the |
107 | // current region. |
108 | if (auto term = dyn_cast<RegionBranchTerminatorOpInterface>( |
109 | block.getTerminator())) { |
110 | collectUnderlyingAddressValues( |
111 | term.getSuccessorOperands(branchPoint)[*operandIndex], maxDepth, |
112 | visited, output); |
113 | } else if (block.getNumSuccessors()) { |
114 | // Otherwise, if this terminator may exit the region we can't make |
115 | // any assumptions about which values get passed. |
116 | output.push_back(inputValue); |
117 | return; |
118 | } |
119 | } |
120 | } |
121 | } |
122 | } |
123 | |
124 | /// Given a result, collect all of the underlying values being addressed. |
125 | static void collectUnderlyingAddressValues(OpResult result, unsigned maxDepth, |
126 | DenseSet<Value> &visited, |
127 | SmallVectorImpl<Value> &output) { |
128 | Operation *op = result.getOwner(); |
129 | |
130 | // If this is a view, unwrap to the source. |
131 | if (ViewLikeOpInterface view = dyn_cast<ViewLikeOpInterface>(op)) |
132 | return collectUnderlyingAddressValues(view.getViewSource(), maxDepth, |
133 | visited, output); |
134 | // Check to see if we can reason about the control flow of this op. |
135 | if (auto branch = dyn_cast<RegionBranchOpInterface>(op)) { |
136 | return collectUnderlyingAddressValues(branch, /*region=*/nullptr, result, |
137 | result.getResultNumber(), maxDepth, |
138 | visited, output); |
139 | } |
140 | |
141 | output.push_back(Elt: result); |
142 | } |
143 | |
144 | /// Given a block argument, collect all of the underlying values being |
145 | /// addressed. |
146 | static void collectUnderlyingAddressValues(BlockArgument arg, unsigned maxDepth, |
147 | DenseSet<Value> &visited, |
148 | SmallVectorImpl<Value> &output) { |
149 | Block *block = arg.getOwner(); |
150 | unsigned argNumber = arg.getArgNumber(); |
151 | |
152 | // Handle the case of a non-entry block. |
153 | if (!block->isEntryBlock()) { |
154 | for (auto it = block->pred_begin(), e = block->pred_end(); it != e; ++it) { |
155 | auto branch = dyn_cast<BranchOpInterface>((*it)->getTerminator()); |
156 | if (!branch) { |
157 | // We can't analyze the control flow, so bail out early. |
158 | output.push_back(Elt: arg); |
159 | return; |
160 | } |
161 | |
162 | // Try to get the operand passed for this argument. |
163 | unsigned index = it.getSuccessorIndex(); |
164 | Value operand = branch.getSuccessorOperands(index)[argNumber]; |
165 | if (!operand) { |
166 | // We can't analyze the control flow, so bail out early. |
167 | output.push_back(Elt: arg); |
168 | return; |
169 | } |
170 | collectUnderlyingAddressValues(value: operand, maxDepth, visited, output); |
171 | } |
172 | return; |
173 | } |
174 | |
175 | // Otherwise, check to see if we can reason about the control flow of this op. |
176 | Region *region = block->getParent(); |
177 | Operation *op = region->getParentOp(); |
178 | if (auto branch = dyn_cast<RegionBranchOpInterface>(op)) { |
179 | return collectUnderlyingAddressValues(branch, region, arg, argNumber, |
180 | maxDepth, visited, output); |
181 | } |
182 | |
183 | // We can't reason about the underlying address of this argument. |
184 | output.push_back(Elt: arg); |
185 | } |
186 | |
187 | /// Given a value, collect all of the underlying values being addressed. |
188 | static void collectUnderlyingAddressValues(Value value, unsigned maxDepth, |
189 | DenseSet<Value> &visited, |
190 | SmallVectorImpl<Value> &output) { |
191 | // Check that we don't infinitely recurse. |
192 | if (!visited.insert(V: value).second) |
193 | return; |
194 | if (maxDepth == 0) { |
195 | output.push_back(Elt: value); |
196 | return; |
197 | } |
198 | --maxDepth; |
199 | |
200 | if (BlockArgument arg = dyn_cast<BlockArgument>(Val&: value)) |
201 | return collectUnderlyingAddressValues(arg, maxDepth, visited, output); |
202 | collectUnderlyingAddressValues(result: cast<OpResult>(Val&: value), maxDepth, visited, |
203 | output); |
204 | } |
205 | |
206 | /// Given a value, collect all of the underlying values being addressed. |
207 | static void collectUnderlyingAddressValues(Value value, |
208 | SmallVectorImpl<Value> &output) { |
209 | DenseSet<Value> visited; |
210 | collectUnderlyingAddressValues(value, maxDepth: maxUnderlyingValueSearchDepth, visited, |
211 | output); |
212 | } |
213 | |
214 | //===----------------------------------------------------------------------===// |
215 | // LocalAliasAnalysis: alias |
216 | //===----------------------------------------------------------------------===// |
217 | |
218 | /// Given a value, try to get an allocation effect attached to it. If |
219 | /// successful, `allocEffect` is populated with the effect. If an effect was |
220 | /// found, `allocScopeOp` is also specified if a parent operation of `value` |
221 | /// could be identified that bounds the scope of the allocated value; i.e. if |
222 | /// non-null it specifies the parent operation that the allocation does not |
223 | /// escape. If no scope is found, `allocScopeOp` is set to nullptr. |
224 | static LogicalResult |
225 | getAllocEffectFor(Value value, |
226 | std::optional<MemoryEffects::EffectInstance> &effect, |
227 | Operation *&allocScopeOp) { |
228 | // Try to get a memory effect interface for the parent operation. |
229 | Operation *op; |
230 | if (BlockArgument arg = dyn_cast<BlockArgument>(Val&: value)) |
231 | op = arg.getOwner()->getParentOp(); |
232 | else |
233 | op = cast<OpResult>(Val&: value).getOwner(); |
234 | MemoryEffectOpInterface interface = dyn_cast<MemoryEffectOpInterface>(op); |
235 | if (!interface) |
236 | return failure(); |
237 | |
238 | // Try to find an allocation effect on the resource. |
239 | if (!(effect = interface.getEffectOnValue<MemoryEffects::Allocate>(value))) |
240 | return failure(); |
241 | |
242 | // If we found an allocation effect, try to find a scope for the allocation. |
243 | // If the resource of this allocation is automatically scoped, find the parent |
244 | // operation that bounds the allocation scope. |
245 | if (llvm::isa<SideEffects::AutomaticAllocationScopeResource>( |
246 | Val: effect->getResource())) { |
247 | allocScopeOp = op->getParentWithTrait<OpTrait::AutomaticAllocationScope>(); |
248 | return success(); |
249 | } |
250 | |
251 | // TODO: Here we could look at the users to see if the resource is either |
252 | // freed on all paths within the region, or is just not captured by anything. |
253 | // For now assume allocation scope to the function scope (we don't care if |
254 | // pointer escape outside function). |
255 | allocScopeOp = op->getParentOfType<FunctionOpInterface>(); |
256 | return success(); |
257 | } |
258 | |
259 | /// Given the two values, return their aliasing behavior. |
260 | AliasResult LocalAliasAnalysis::aliasImpl(Value lhs, Value rhs) { |
261 | if (lhs == rhs) |
262 | return AliasResult::MustAlias; |
263 | Operation *lhsAllocScope = nullptr, *rhsAllocScope = nullptr; |
264 | std::optional<MemoryEffects::EffectInstance> lhsAlloc, rhsAlloc; |
265 | |
266 | // Handle the case where lhs is a constant. |
267 | Attribute lhsAttr, rhsAttr; |
268 | if (matchPattern(value: lhs, pattern: m_Constant(bind_value: &lhsAttr))) { |
269 | // TODO: This is overly conservative. Two matching constants don't |
270 | // necessarily map to the same address. For example, if the two values |
271 | // correspond to different symbols that both represent a definition. |
272 | if (matchPattern(value: rhs, pattern: m_Constant(bind_value: &rhsAttr))) |
273 | return AliasResult::MayAlias; |
274 | |
275 | // Try to find an alloc effect on rhs. If an effect was found we can't |
276 | // alias, otherwise we might. |
277 | return succeeded(result: getAllocEffectFor(value: rhs, effect&: rhsAlloc, allocScopeOp&: rhsAllocScope)) |
278 | ? AliasResult::NoAlias |
279 | : AliasResult::MayAlias; |
280 | } |
281 | // Handle the case where rhs is a constant. |
282 | if (matchPattern(value: rhs, pattern: m_Constant(bind_value: &rhsAttr))) { |
283 | // Try to find an alloc effect on lhs. If an effect was found we can't |
284 | // alias, otherwise we might. |
285 | return succeeded(result: getAllocEffectFor(value: lhs, effect&: lhsAlloc, allocScopeOp&: lhsAllocScope)) |
286 | ? AliasResult::NoAlias |
287 | : AliasResult::MayAlias; |
288 | } |
289 | |
290 | // Otherwise, neither of the values are constant so check to see if either has |
291 | // an allocation effect. |
292 | bool lhsHasAlloc = succeeded(result: getAllocEffectFor(value: lhs, effect&: lhsAlloc, allocScopeOp&: lhsAllocScope)); |
293 | bool rhsHasAlloc = succeeded(result: getAllocEffectFor(value: rhs, effect&: rhsAlloc, allocScopeOp&: rhsAllocScope)); |
294 | if (lhsHasAlloc == rhsHasAlloc) { |
295 | // If both values have an allocation effect we know they don't alias, and if |
296 | // neither have an effect we can't make an assumptions. |
297 | return lhsHasAlloc ? AliasResult::NoAlias : AliasResult::MayAlias; |
298 | } |
299 | |
300 | // When we reach this point we have one value with a known allocation effect, |
301 | // and one without. Move the one with the effect to the lhs to make the next |
302 | // checks simpler. |
303 | if (rhsHasAlloc) { |
304 | std::swap(a&: lhs, b&: rhs); |
305 | lhsAlloc = rhsAlloc; |
306 | lhsAllocScope = rhsAllocScope; |
307 | } |
308 | |
309 | // If the effect has a scoped allocation region, check to see if the |
310 | // non-effect value is defined above that scope. |
311 | if (lhsAllocScope) { |
312 | // If the parent operation of rhs is an ancestor of the allocation scope, or |
313 | // if rhs is an entry block argument of the allocation scope we know the two |
314 | // values can't alias. |
315 | Operation *rhsParentOp = rhs.getParentRegion()->getParentOp(); |
316 | if (rhsParentOp->isProperAncestor(other: lhsAllocScope)) |
317 | return AliasResult::NoAlias; |
318 | if (rhsParentOp == lhsAllocScope) { |
319 | BlockArgument rhsArg = dyn_cast<BlockArgument>(Val&: rhs); |
320 | if (rhsArg && rhs.getParentBlock()->isEntryBlock()) |
321 | return AliasResult::NoAlias; |
322 | } |
323 | } |
324 | |
325 | // If we couldn't reason about the relationship between the two values, |
326 | // conservatively assume they might alias. |
327 | return AliasResult::MayAlias; |
328 | } |
329 | |
330 | /// Given the two values, return their aliasing behavior. |
331 | AliasResult LocalAliasAnalysis::alias(Value lhs, Value rhs) { |
332 | if (lhs == rhs) |
333 | return AliasResult::MustAlias; |
334 | |
335 | // Get the underlying values being addressed. |
336 | SmallVector<Value, 8> lhsValues, rhsValues; |
337 | collectUnderlyingAddressValues(value: lhs, output&: lhsValues); |
338 | collectUnderlyingAddressValues(value: rhs, output&: rhsValues); |
339 | |
340 | // If we failed to collect for either of the values somehow, conservatively |
341 | // assume they may alias. |
342 | if (lhsValues.empty() || rhsValues.empty()) |
343 | return AliasResult::MayAlias; |
344 | |
345 | // Check the alias results against each of the underlying values. |
346 | std::optional<AliasResult> result; |
347 | for (Value lhsVal : lhsValues) { |
348 | for (Value rhsVal : rhsValues) { |
349 | AliasResult nextResult = aliasImpl(lhs: lhsVal, rhs: rhsVal); |
350 | result = result ? result->merge(other: nextResult) : nextResult; |
351 | } |
352 | } |
353 | |
354 | // We should always have a valid result here. |
355 | return *result; |
356 | } |
357 | |
358 | //===----------------------------------------------------------------------===// |
359 | // LocalAliasAnalysis: getModRef |
360 | //===----------------------------------------------------------------------===// |
361 | |
362 | ModRefResult LocalAliasAnalysis::getModRef(Operation *op, Value location) { |
363 | // Check to see if this operation relies on nested side effects. |
364 | if (op->hasTrait<OpTrait::HasRecursiveMemoryEffects>()) { |
365 | // TODO: To check recursive operations we need to check all of the nested |
366 | // operations, which can result in a quadratic number of queries. We should |
367 | // introduce some caching of some kind to help alleviate this, especially as |
368 | // this caching could be used in other areas of the codebase (e.g. when |
369 | // checking `wouldOpBeTriviallyDead`). |
370 | return ModRefResult::getModAndRef(); |
371 | } |
372 | |
373 | // Otherwise, check to see if this operation has a memory effect interface. |
374 | MemoryEffectOpInterface interface = dyn_cast<MemoryEffectOpInterface>(op); |
375 | if (!interface) |
376 | return ModRefResult::getModAndRef(); |
377 | |
378 | // Build a ModRefResult by merging the behavior of the effects of this |
379 | // operation. |
380 | SmallVector<MemoryEffects::EffectInstance> effects; |
381 | interface.getEffects(effects); |
382 | |
383 | ModRefResult result = ModRefResult::getNoModRef(); |
384 | for (const MemoryEffects::EffectInstance &effect : effects) { |
385 | if (isa<MemoryEffects::Allocate, MemoryEffects::Free>(Val: effect.getEffect())) |
386 | continue; |
387 | |
388 | // Check for an alias between the effect and our memory location. |
389 | // TODO: Add support for checking an alias with a symbol reference. |
390 | AliasResult aliasResult = AliasResult::MayAlias; |
391 | if (Value effectValue = effect.getValue()) |
392 | aliasResult = alias(lhs: effectValue, rhs: location); |
393 | |
394 | // If we don't alias, ignore this effect. |
395 | if (aliasResult.isNo()) |
396 | continue; |
397 | |
398 | // Merge in the corresponding mod or ref for this effect. |
399 | if (isa<MemoryEffects::Read>(Val: effect.getEffect())) { |
400 | result = result.merge(other: ModRefResult::getRef()); |
401 | } else { |
402 | assert(isa<MemoryEffects::Write>(effect.getEffect())); |
403 | result = result.merge(other: ModRefResult::getMod()); |
404 | } |
405 | if (result.isModAndRef()) |
406 | break; |
407 | } |
408 | return result; |
409 | } |
410 | |