1 | //===- BufferDeallocationOpInterface.h --------------------------*- 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_DIALECT_BUFFERIZATION_IR_BUFFERDEALLOCATIONOPINTERFACE_H_ |
10 | #define MLIR_DIALECT_BUFFERIZATION_IR_BUFFERDEALLOCATIONOPINTERFACE_H_ |
11 | |
12 | #include "mlir/Analysis/Liveness.h" |
13 | #include "mlir/IR/Operation.h" |
14 | #include "mlir/IR/SymbolTable.h" |
15 | #include "mlir/Support/LLVM.h" |
16 | |
17 | namespace mlir { |
18 | namespace bufferization { |
19 | |
20 | /// Compare two SSA values in a deterministic manner. Two block arguments are |
21 | /// ordered by argument number, block arguments are always less than operation |
22 | /// results, and operation results are ordered by the `isBeforeInBlock` order of |
23 | /// their defining operation. |
24 | struct ValueComparator { |
25 | bool operator()(const Value &lhs, const Value &rhs) const; |
26 | }; |
27 | |
28 | /// This class is used to track the ownership of values. The ownership can |
29 | /// either be not initialized yet ('Uninitialized' state), set to a unique SSA |
30 | /// value which indicates the ownership at runtime (or statically if it is a |
31 | /// constant value) ('Unique' state), or it cannot be represented in a single |
32 | /// SSA value ('Unknown' state). An artificial example of a case where ownership |
33 | /// cannot be represented in a single i1 SSA value could be the following: |
34 | /// `%0 = test.non_deterministic_select %arg0, %arg1 : i32` |
35 | /// Since the operation does not provide us a separate boolean indicator on |
36 | /// which of the two operands was selected, we would need to either insert an |
37 | /// alias check at runtime to determine if `%0` aliases with `%arg0` or `%arg1`, |
38 | /// or insert a `bufferization.clone` operation to get a fresh buffer which we |
39 | /// could assign ownership to. |
40 | /// |
41 | /// The three states this class can represent form a lattice on a partial order: |
42 | /// forall X in SSA values. uninitialized < unique(X) < unknown |
43 | /// forall X, Y in SSA values. |
44 | /// unique(X) == unique(Y) iff X and Y always evaluate to the same value |
45 | /// unique(X) != unique(Y) otherwise |
46 | class Ownership { |
47 | public: |
48 | /// Constructor that creates an 'Uninitialized' ownership. This is needed for |
49 | /// default-construction when used in DenseMap. |
50 | Ownership() = default; |
51 | |
52 | /// Constructor that creates an 'Unique' ownership. This is a non-explicit |
53 | /// constructor to allow implicit conversion from 'Value'. |
54 | Ownership(Value indicator); |
55 | |
56 | /// Get an ownership value in 'Unknown' state. |
57 | static Ownership getUnknown(); |
58 | /// Get an ownership value in 'Unique' state with 'indicator' as parameter. |
59 | static Ownership getUnique(Value indicator); |
60 | /// Get an ownership value in 'Uninitialized' state. |
61 | static Ownership getUninitialized(); |
62 | |
63 | /// Check if this ownership value is in the 'Uninitialized' state. |
64 | bool isUninitialized() const; |
65 | /// Check if this ownership value is in the 'Unique' state. |
66 | bool isUnique() const; |
67 | /// Check if this ownership value is in the 'Unknown' state. |
68 | bool isUnknown() const; |
69 | |
70 | /// If this ownership value is in 'Unique' state, this function can be used to |
71 | /// get the indicator parameter. Using this function in any other state is UB. |
72 | Value getIndicator() const; |
73 | |
74 | /// Get the join of the two-element subset {this,other}. Does not modify |
75 | /// 'this'. |
76 | Ownership getCombined(Ownership other) const; |
77 | |
78 | /// Modify 'this' ownership to be the join of the current 'this' and 'other'. |
79 | void combine(Ownership other); |
80 | |
81 | private: |
82 | enum class State { |
83 | Uninitialized, |
84 | Unique, |
85 | Unknown, |
86 | }; |
87 | |
88 | // The indicator value is only relevant in the 'Unique' state. |
89 | Value indicator; |
90 | State state = State::Uninitialized; |
91 | }; |
92 | |
93 | /// Options for BufferDeallocationOpInterface-based buffer deallocation. |
94 | struct DeallocationOptions { |
95 | // A pass option indicating whether private functions should be modified to |
96 | // pass the ownership of MemRef values instead of adhering to the function |
97 | // boundary ABI. |
98 | bool privateFuncDynamicOwnership = false; |
99 | }; |
100 | |
101 | /// This class collects all the state that we need to perform the buffer |
102 | /// deallocation pass with associated helper functions such that we have easy |
103 | /// access to it in the BufferDeallocationOpInterface implementations and the |
104 | /// BufferDeallocation pass. |
105 | class DeallocationState { |
106 | public: |
107 | DeallocationState(Operation *op); |
108 | |
109 | // The state should always be passed by reference. |
110 | DeallocationState(const DeallocationState &) = delete; |
111 | |
112 | /// Small helper function to update the ownership map by taking the current |
113 | /// ownership ('Uninitialized' state if not yet present), computing the join |
114 | /// with the passed ownership and storing this new value in the map. By |
115 | /// default, it will be performed for the block where 'owned' is defined. If |
116 | /// the ownership of the given value should be updated for another block, the |
117 | /// 'block' argument can be explicitly passed. |
118 | void updateOwnership(Value memref, Ownership ownership, |
119 | Block *block = nullptr); |
120 | |
121 | /// Removes ownerships associated with all values in the passed range for |
122 | /// 'block'. |
123 | void resetOwnerships(ValueRange memrefs, Block *block); |
124 | |
125 | /// Returns the ownership of 'memref' for the given basic block. |
126 | Ownership getOwnership(Value memref, Block *block) const; |
127 | |
128 | /// Remember the given 'memref' to deallocate it at the end of the 'block'. |
129 | void addMemrefToDeallocate(Value memref, Block *block); |
130 | |
131 | /// Forget about a MemRef that we originally wanted to deallocate at the end |
132 | /// of 'block', possibly because it already gets deallocated before the end of |
133 | /// the block. |
134 | void dropMemrefToDeallocate(Value memref, Block *block); |
135 | |
136 | /// Return a sorted list of MemRef values which are live at the start of the |
137 | /// given block. |
138 | void getLiveMemrefsIn(Block *block, SmallVectorImpl<Value> &memrefs); |
139 | |
140 | /// Given an SSA value of MemRef type, this function queries the ownership and |
141 | /// if it is not already in the 'Unique' state, potentially inserts IR to get |
142 | /// a new SSA value, returned as the first element of the pair, which has |
143 | /// 'Unique' ownership and can be used instead of the passed Value with the |
144 | /// the ownership indicator returned as the second element of the pair. |
145 | std::pair<Value, Value> |
146 | getMemrefWithUniqueOwnership(OpBuilder &builder, Value memref, Block *block); |
147 | |
148 | /// Given two basic blocks and the values passed via block arguments to the |
149 | /// destination block, compute the list of MemRefs that have to be retained in |
150 | /// the 'fromBlock' to not run into a use-after-free situation. |
151 | /// This list consists of the MemRefs in the successor operand list of the |
152 | /// terminator and the MemRefs in the 'out' set of the liveness analysis |
153 | /// intersected with the 'in' set of the destination block. |
154 | /// |
155 | /// toRetain = filter(successorOperands + (liveOut(fromBlock) insersect |
156 | /// liveIn(toBlock)), isMemRef) |
157 | void getMemrefsToRetain(Block *fromBlock, Block *toBlock, |
158 | ValueRange destOperands, |
159 | SmallVectorImpl<Value> &toRetain) const; |
160 | |
161 | /// For a given block, computes the list of MemRefs that potentially need to |
162 | /// be deallocated at the end of that block. This list also contains values |
163 | /// that have to be retained (and are thus part of the list returned by |
164 | /// `getMemrefsToRetain`) and is computed by taking the MemRefs in the 'in' |
165 | /// set of the liveness analysis of 'block' appended by the set of MemRefs |
166 | /// allocated in 'block' itself and subtracted by the set of MemRefs |
167 | /// deallocated in 'block'. |
168 | /// Note that we don't have to take the intersection of the liveness 'in' set |
169 | /// with the 'out' set of the predecessor block because a value that is in the |
170 | /// 'in' set must be defined in an ancestor block that dominates all direct |
171 | /// predecessors and thus the 'in' set of this block is a subset of the 'out' |
172 | /// sets of each predecessor. |
173 | /// |
174 | /// memrefs = filter((liveIn(block) U |
175 | /// allocated(block) U arguments(block)) \ deallocated(block), isMemRef) |
176 | /// |
177 | /// The list of conditions is then populated by querying the internal |
178 | /// datastructures for the ownership value of that MemRef. |
179 | LogicalResult |
180 | getMemrefsAndConditionsToDeallocate(OpBuilder &builder, Location loc, |
181 | Block *block, |
182 | SmallVectorImpl<Value> &memrefs, |
183 | SmallVectorImpl<Value> &conditions) const; |
184 | |
185 | /// Returns the symbol cache to lookup functions from call operations to check |
186 | /// attributes on the function operation. |
187 | SymbolTableCollection *getSymbolTable() { return &symbolTable; } |
188 | |
189 | private: |
190 | // Symbol cache to lookup functions from call operations to check attributes |
191 | // on the function operation. |
192 | SymbolTableCollection symbolTable; |
193 | |
194 | // Mapping from each SSA value with MemRef type to the associated ownership in |
195 | // each block. |
196 | DenseMap<std::pair<Value, Block *>, Ownership> ownershipMap; |
197 | |
198 | // Collects the list of MemRef values that potentially need to be deallocated |
199 | // per block. It is also fine (albeit not efficient) to add MemRef values that |
200 | // don't have to be deallocated, but only when the ownership is not 'Unknown'. |
201 | DenseMap<Block *, SmallVector<Value>> memrefsToDeallocatePerBlock; |
202 | |
203 | // The underlying liveness analysis to compute fine grained information about |
204 | // alloc and dealloc positions. |
205 | Liveness liveness; |
206 | }; |
207 | |
208 | namespace deallocation_impl { |
209 | /// Insert a `bufferization.dealloc` operation right before `op` which has to be |
210 | /// a terminator without any successors. Note that it is not required to have |
211 | /// the ReturnLike trait attached. The MemRef values in the `operands` argument |
212 | /// will be added to the list of retained values and their updated ownership |
213 | /// values will be appended to the `updatedOperandOwnerships` list. `op` is not |
214 | /// modified in any way. Returns failure if at least one of the MemRefs to |
215 | /// deallocate does not have 'Unique' ownership (likely as a result of an |
216 | /// incorrect implementation of the `process` or |
217 | /// `materializeUniqueOwnershipForMemref` interface method) or the original |
218 | /// `op`. |
219 | FailureOr<Operation *> |
220 | insertDeallocOpForReturnLike(DeallocationState &state, Operation *op, |
221 | ValueRange operands, |
222 | SmallVectorImpl<Value> &updatedOperandOwnerships); |
223 | } // namespace deallocation_impl |
224 | |
225 | } // namespace bufferization |
226 | } // namespace mlir |
227 | |
228 | //===----------------------------------------------------------------------===// |
229 | // Buffer Deallocation Interface |
230 | //===----------------------------------------------------------------------===// |
231 | |
232 | #include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h.inc" |
233 | |
234 | #endif // MLIR_DIALECT_BUFFERIZATION_IR_BUFFERDEALLOCATIONOPINTERFACE_H_ |
235 | |