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
17namespace mlir {
18namespace 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.
24struct 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
46class Ownership {
47public:
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
81private:
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.
94struct 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.
105class DeallocationState {
106public:
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
189private:
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
208namespace 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`.
219FailureOr<Operation *>
220insertDeallocOpForReturnLike(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

source code of mlir/include/mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h