1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
|
//======- BufferViewFlowAnalysis.cpp - Buffer alias analysis -*- C++ -*-======//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Interfaces/ViewLikeInterface.h"
#include "llvm/ADT/SetOperations.h"
using namespace mlir;
/// Constructs a new alias analysis using the op provided.
BufferViewFlowAnalysis::BufferViewFlowAnalysis(Operation *op) { build(op); }
/// Find all immediate and indirect dependent buffers this value could
/// potentially have. Note that the resulting set will also contain the value
/// provided as it is a dependent alias of itself.
BufferViewFlowAnalysis::ValueSetT
BufferViewFlowAnalysis::resolve(Value rootValue) const {
ValueSetT result;
SmallVector<Value, 8> queue;
queue.push_back(rootValue);
while (!queue.empty()) {
Value currentValue = queue.pop_back_val();
if (result.insert(currentValue).second) {
auto it = dependencies.find(currentValue);
if (it != dependencies.end()) {
for (Value aliasValue : it->second)
queue.push_back(aliasValue);
}
}
}
return result;
}
/// Removes the given values from all alias sets.
void BufferViewFlowAnalysis::remove(const SmallPtrSetImpl<Value> &aliasValues) {
for (auto &entry : dependencies)
llvm::set_subtract(entry.second, aliasValues);
}
/// This function constructs a mapping from values to its immediate
/// dependencies. It iterates over all blocks, gets their predecessors,
/// determines the values that will be passed to the corresponding block
/// arguments and inserts them into the underlying map. Furthermore, it wires
/// successor regions and branch-like return operations from nested regions.
void BufferViewFlowAnalysis::build(Operation *op) {
// Registers all dependencies of the given values.
auto registerDependencies = [&](ValueRange values, ValueRange dependencies) {
for (auto [value, dep] : llvm::zip(values, dependencies))
this->dependencies[value].insert(dep);
};
// Add additional dependencies created by view changes to the alias list.
op->walk([&](ViewLikeOpInterface viewInterface) {
dependencies[viewInterface.getViewSource()].insert(
viewInterface->getResult(0));
});
// Query all branch interfaces to link block argument dependencies.
op->walk([&](BranchOpInterface branchInterface) {
Block *parentBlock = branchInterface->getBlock();
for (auto it = parentBlock->succ_begin(), e = parentBlock->succ_end();
it != e; ++it) {
// Query the branch op interface to get the successor operands.
auto successorOperands =
branchInterface.getSuccessorOperands(it.getIndex());
// Build the actual mapping of values to their immediate dependencies.
registerDependencies(successorOperands.getForwardedOperands(),
(*it)->getArguments().drop_front(
successorOperands.getProducedOperandCount()));
}
});
// Query the RegionBranchOpInterface to find potential successor regions.
op->walk([&](RegionBranchOpInterface regionInterface) {
// Extract all entry regions and wire all initial entry successor inputs.
SmallVector<RegionSuccessor, 2> entrySuccessors;
regionInterface.getSuccessorRegions(/*index=*/std::nullopt,
entrySuccessors);
for (RegionSuccessor &entrySuccessor : entrySuccessors) {
// Wire the entry region's successor arguments with the initial
// successor inputs.
assert(entrySuccessor.getSuccessor() &&
"Invalid entry region without an attached successor region");
registerDependencies(
regionInterface.getSuccessorEntryOperands(
entrySuccessor.getSuccessor()->getRegionNumber()),
entrySuccessor.getSuccessorInputs());
}
// Wire flow between regions and from region exits.
for (Region ®ion : regionInterface->getRegions()) {
// Iterate over all successor region entries that are reachable from the
// current region.
SmallVector<RegionSuccessor, 2> successorRegions;
regionInterface.getSuccessorRegions(region.getRegionNumber(),
successorRegions);
for (RegionSuccessor &successorRegion : successorRegions) {
// Determine the current region index (if any).
std::optional<unsigned> regionIndex;
Region *regionSuccessor = successorRegion.getSuccessor();
if (regionSuccessor)
regionIndex = regionSuccessor->getRegionNumber();
// Iterate over all immediate terminator operations and wire the
// successor inputs with the successor operands of each terminator.
for (Block &block : region) {
auto successorOperands = getRegionBranchSuccessorOperands(
block.getTerminator(), regionIndex);
if (successorOperands) {
registerDependencies(*successorOperands,
successorRegion.getSuccessorInputs());
}
}
}
}
});
// TODO: This should be an interface.
op->walk([&](arith::SelectOp selectOp) {
registerDependencies({selectOp.getOperand(1)}, {selectOp.getResult()});
registerDependencies({selectOp.getOperand(2)}, {selectOp.getResult()});
});
}
|
