1 //===- OneShotAnalysis.h - One-Shot (Single Pass) Analysis ------*- 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_TRANSFORMS_ONESHOTANALYSIS_H 10 #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTANALYSIS_H 11 12 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h" 13 #include "llvm/ADT/EquivalenceClasses.h" 14 15 namespace mlir { 16 namespace bufferization { 17 18 struct OneShotBufferizationOptions; 19 class BufferizationAliasInfo; 20 class OneShotAnalysisState; 21 22 /// Options for analysis-enabled bufferization. 23 struct OneShotBufferizationOptions : public BufferizationOptions { 24 OneShotBufferizationOptions() = default; 25 26 /// Specifies whether returning newly allocated memrefs should be allowed. 27 /// Otherwise, a pass failure is triggered. 28 bool allowReturnAllocs = false; 29 }; 30 31 /// The BufferizationAliasInfo class maintains a list of buffer aliases and 32 /// equivalence classes to support bufferization. 33 class BufferizationAliasInfo { 34 public: 35 explicit BufferizationAliasInfo(Operation *rootOp); 36 37 // BufferizationAliasInfo should be passed as a reference. 38 BufferizationAliasInfo(const BufferizationAliasInfo &) = delete; 39 40 /// Add a new entry for `v` in the `aliasInfo` and `equivalentInfo`. In the 41 /// beginning the alias and equivalence sets only contain `v` itself. 42 void createAliasInfoEntry(Value v); 43 44 /// Insert an info entry for `newValue` and merge its alias set with that of 45 /// `alias`. 46 void insertNewBufferAlias(Value newValue, Value alias); 47 48 /// Insert an info entry for `newValue` and merge its alias set with that of 49 /// `alias`. Additionally, merge their equivalence classes. 50 void insertNewBufferEquivalence(Value newValue, Value alias); 51 52 /// Set the inPlace bufferization spec to true. 53 /// Merge result's and operand's aliasing sets and iterate to a fixed point. 54 void bufferizeInPlace(OpOperand &operand, AnalysisState &state); 55 56 /// Set the inPlace bufferization spec to false. 57 void bufferizeOutOfPlace(OpOperand &operand); 58 59 /// Return true if `v1` and `v2` may bufferize to aliasing buffers. areAliasingBufferizedValues(Value v1,Value v2)60 bool areAliasingBufferizedValues(Value v1, Value v2) const { 61 return aliasInfo.isEquivalent(v1, v2); 62 } 63 64 /// Return true if `v1` and `v2` bufferize to equivalent buffers. areEquivalentBufferizedValues(Value v1,Value v2)65 bool areEquivalentBufferizedValues(Value v1, Value v2) const { 66 return equivalentInfo.isEquivalent(v1, v2); 67 } 68 69 /// Union the alias sets of `v1` and `v2`. unionAliasSets(Value v1,Value v2)70 void unionAliasSets(Value v1, Value v2) { aliasInfo.unionSets(v1, v2); } 71 72 /// Union the equivalence classes of `v1` and `v2`. unionEquivalenceClasses(Value v1,Value v2)73 void unionEquivalenceClasses(Value v1, Value v2) { 74 equivalentInfo.unionSets(v1, v2); 75 } 76 77 /// Apply `fun` to all the members of the equivalence class of `v`. 78 void applyOnEquivalenceClass(Value v, function_ref<void(Value)> fun) const; 79 80 /// Apply `fun` to all aliases of `v`. 81 void applyOnAliases(Value v, function_ref<void(Value)> fun) const; 82 83 /// Mark a value as in-place bufferized. markInPlace(OpOperand & o)84 void markInPlace(OpOperand &o) { inplaceBufferized.insert(&o); } 85 86 /// Return `true` if a value was marked as in-place bufferized. 87 bool isInPlace(OpOperand &opOperand) const; 88 89 private: 90 /// llvm::EquivalenceClasses wants comparable elements. This comparator uses 91 /// uses pointer comparison on the defining op. This is a poor man's 92 /// comparison but it's not like UnionFind needs ordering anyway. 93 struct ValueComparator { operatorValueComparator94 bool operator()(const Value &lhs, const Value &rhs) const { 95 return lhs.getImpl() < rhs.getImpl(); 96 } 97 }; 98 99 using EquivalenceClassRangeType = llvm::iterator_range< 100 llvm::EquivalenceClasses<Value, ValueComparator>::member_iterator>; 101 /// Check that aliasInfo for `v` exists and return a reference to it. 102 EquivalenceClassRangeType getAliases(Value v) const; 103 104 /// Set of all OpResults that were decided to bufferize in-place. 105 llvm::DenseSet<OpOperand *> inplaceBufferized; 106 107 /// Auxiliary structure to store all the values a given value may alias with. 108 /// Alias information is "may be" conservative: In the presence of branches, a 109 /// value may alias with one of multiple other values. The concrete aliasing 110 /// value may not even be known at compile time. All such values are 111 /// considered to be aliases. 112 llvm::EquivalenceClasses<Value, ValueComparator> aliasInfo; 113 114 /// Auxiliary structure to store all the equivalent buffer classes. Equivalent 115 /// buffer information is "must be" conservative: Only if two values are 116 /// guaranteed to be equivalent at runtime, they said to be equivalent. It is 117 /// possible that, in the presence of branches, it cannot be determined 118 /// statically if two values are equivalent. In that case, the values are 119 /// considered to be not equivalent. 120 llvm::EquivalenceClasses<Value, ValueComparator> equivalentInfo; 121 }; 122 123 /// State for analysis-enabled bufferization. This class keeps track of alias 124 /// (via BufferizationAliasInfo) to decide if tensor OpOperands should bufferize 125 /// in-place. 126 class OneShotAnalysisState : public AnalysisState { 127 public: 128 OneShotAnalysisState(Operation *op, 129 const OneShotBufferizationOptions &options); 130 131 OneShotAnalysisState(const OneShotAnalysisState &) = delete; 132 133 virtual ~OneShotAnalysisState() = default; 134 135 /// Return a reference to the BufferizationAliasInfo. getAliasInfo()136 BufferizationAliasInfo &getAliasInfo() { return aliasInfo; } 137 138 /// Return `true` if the given OpResult has been decided to bufferize inplace. 139 bool isInPlace(OpOperand &opOperand) const override; 140 141 /// Return true if `v1` and `v2` bufferize to equivalent buffers. 142 bool areEquivalentBufferizedValues(Value v1, Value v2) const override; 143 144 /// Return true if `v1` and `v2` may bufferize to aliasing buffers. 145 bool areAliasingBufferizedValues(Value v1, Value v2) const override; 146 147 /// Return `true` if the given tensor has undefined contents. 148 bool hasUndefinedContents(OpOperand *opOperand) const override; 149 150 /// Return true if the given tensor (or an aliasing tensor) is yielded from 151 /// the containing block. Also include all aliasing tensors in the same block. 152 bool isTensorYielded(Value tensor) const override; 153 154 /// Find all tensor values in the given operation that have undefined contents 155 /// and store them in `undefinedTensorUses`. 156 void gatherUndefinedTensorUses(Operation *op); 157 158 /// Find all tensors that are yielded/returned from a block and store them in 159 /// `yieldedTensors`. Also include all aliasing tensors in the same block. 160 void gatherYieldedTensors(Operation *op); 161 162 /// Return true if the buffer of the given tensor value is written to. Must 163 /// not be called for values inside not yet analyzed functions. 164 bool isValueWritten(Value value) const; 165 166 /// Return true if the buffer of the given tensor value is writable. 167 bool isWritable(Value value) const; 168 169 private: 170 /// `aliasInfo` keeps track of aliasing and equivalent values. Only internal 171 /// functions and `runOneShotBufferize` may access this object. 172 BufferizationAliasInfo aliasInfo; 173 174 /// A set of all tensors (and maybe aliasing tensors) that yielded from a 175 /// block. 176 DenseSet<Value> yieldedTensors; 177 178 /// A set of uses of tensors that have undefined contents. 179 DenseSet<OpOperand *> undefinedTensorUses; 180 }; 181 182 /// Analyze `op` and its nested ops. Bufferization decisions are stored in 183 /// `state`. 184 LogicalResult analyzeOp(Operation *op, OneShotAnalysisState &state); 185 186 /// Run One-Shot Bufferize on the given op: Analysis + Bufferization 187 LogicalResult runOneShotBufferize(Operation *op, 188 const OneShotBufferizationOptions &options); 189 190 } // namespace bufferization 191 } // namespace mlir 192 193 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTANALYSIS_H 194