MLIR  21.0.0git
SROA.cpp
Go to the documentation of this file.
1 //===-- SROA.cpp - Scalar Replacement Of Aggregates -------------*- 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 #include "mlir/Transforms/SROA.h"
14 #include "mlir/Transforms/Passes.h"
15 
16 namespace mlir {
17 #define GEN_PASS_DEF_SROA
18 #include "mlir/Transforms/Passes.h.inc"
19 } // namespace mlir
20 
21 #define DEBUG_TYPE "sroa"
22 
23 using namespace mlir;
24 
25 namespace {
26 
27 /// Information computed by destructurable memory slot analysis used to perform
28 /// actual destructuring of the slot. This struct is only constructed if
29 /// destructuring is possible, and contains the necessary data to perform it.
30 struct MemorySlotDestructuringInfo {
31  /// Set of the indices that are actually used when accessing the subelements.
32  SmallPtrSet<Attribute, 8> usedIndices;
33  /// Blocking uses of a given user of the memory slot that must be eliminated.
35  /// List of potentially indirect accessors of the memory slot that need
36  /// rewiring.
38 };
39 
40 } // namespace
41 
42 /// Computes information for slot destructuring. This will compute whether this
43 /// slot can be destructured and data to perform the destructuring. Returns
44 /// nothing if the slot cannot be destructured or if there is no useful work to
45 /// be done.
46 static std::optional<MemorySlotDestructuringInfo>
48  const DataLayout &dataLayout) {
49  assert(isa<DestructurableTypeInterface>(slot.elemType));
50 
51  if (slot.ptr.use_empty())
52  return {};
53 
54  MemorySlotDestructuringInfo info;
55 
56  SmallVector<MemorySlot> usedSafelyWorklist;
57 
58  auto scheduleAsBlockingUse = [&](OpOperand &use) {
59  SmallPtrSetImpl<OpOperand *> &blockingUses =
60  info.userToBlockingUses[use.getOwner()];
61  blockingUses.insert(&use);
62  };
63 
64  // Initialize the analysis with the immediate users of the slot.
65  for (OpOperand &use : slot.ptr.getUses()) {
66  if (auto accessor =
67  dyn_cast<DestructurableAccessorOpInterface>(use.getOwner())) {
68  if (accessor.canRewire(slot, info.usedIndices, usedSafelyWorklist,
69  dataLayout)) {
70  info.accessors.push_back(accessor);
71  continue;
72  }
73  }
74 
75  // If it cannot be shown that the operation uses the slot safely, maybe it
76  // can be promoted out of using the slot?
77  scheduleAsBlockingUse(use);
78  }
79 
81  while (!usedSafelyWorklist.empty()) {
82  MemorySlot mustBeUsedSafely = usedSafelyWorklist.pop_back_val();
83  for (OpOperand &subslotUse : mustBeUsedSafely.ptr.getUses()) {
84  if (!visited.insert(&subslotUse).second)
85  continue;
86  Operation *subslotUser = subslotUse.getOwner();
87 
88  if (auto memOp = dyn_cast<SafeMemorySlotAccessOpInterface>(subslotUser))
89  if (succeeded(memOp.ensureOnlySafeAccesses(
90  mustBeUsedSafely, usedSafelyWorklist, dataLayout)))
91  continue;
92 
93  // If it cannot be shown that the operation uses the slot safely, maybe it
94  // can be promoted out of using the slot?
95  scheduleAsBlockingUse(subslotUse);
96  }
97  }
98 
99  SetVector<Operation *> forwardSlice;
100  mlir::getForwardSlice(slot.ptr, &forwardSlice);
101  for (Operation *user : forwardSlice) {
102  // If the next operation has no blocking uses, everything is fine.
103  auto it = info.userToBlockingUses.find(user);
104  if (it == info.userToBlockingUses.end())
105  continue;
106 
107  SmallPtrSet<OpOperand *, 4> &blockingUses = it->second;
108  auto promotable = dyn_cast<PromotableOpInterface>(user);
109 
110  // An operation that has blocking uses must be promoted. If it is not
111  // promotable, destructuring must fail.
112  if (!promotable)
113  return {};
114 
115  SmallVector<OpOperand *> newBlockingUses;
116  // If the operation decides it cannot deal with removing the blocking uses,
117  // destructuring must fail.
118  if (!promotable.canUsesBeRemoved(blockingUses, newBlockingUses, dataLayout))
119  return {};
120 
121  // Then, register any new blocking uses for coming operations.
122  for (OpOperand *blockingUse : newBlockingUses) {
123  assert(llvm::is_contained(user->getResults(), blockingUse->get()));
124 
125  SmallPtrSetImpl<OpOperand *> &newUserBlockingUseSet =
126  info.userToBlockingUses[blockingUse->getOwner()];
127  newUserBlockingUseSet.insert(blockingUse);
128  }
129  }
130 
131  return info;
132 }
133 
134 /// Performs the destructuring of a destructible slot given associated
135 /// destructuring information. The provided slot will be destructured in
136 /// subslots as specified by its allocator.
137 static void destructureSlot(
139  DestructurableAllocationOpInterface allocator, OpBuilder &builder,
140  const DataLayout &dataLayout, MemorySlotDestructuringInfo &info,
142  const SROAStatistics &statistics) {
143  OpBuilder::InsertionGuard guard(builder);
144 
147  allocator.destructure(slot, info.usedIndices, builder, newAllocators);
148 
149  if (statistics.slotsWithMemoryBenefit &&
150  slot.subelementTypes.size() != info.usedIndices.size())
151  (*statistics.slotsWithMemoryBenefit)++;
152 
153  if (statistics.maxSubelementAmount)
154  statistics.maxSubelementAmount->updateMax(slot.subelementTypes.size());
155 
156  SetVector<Operation *> usersToRewire;
157  usersToRewire.insert_range(llvm::make_first_range(info.userToBlockingUses));
158  usersToRewire.insert_range(info.accessors);
159  usersToRewire = mlir::topologicalSort(usersToRewire);
160 
162  for (Operation *toRewire : llvm::reverse(usersToRewire)) {
163  builder.setInsertionPointAfter(toRewire);
164  if (auto accessor = dyn_cast<DestructurableAccessorOpInterface>(toRewire)) {
165  if (accessor.rewire(slot, subslots, builder, dataLayout) ==
167  toErase.push_back(accessor);
168  continue;
169  }
170 
171  auto promotable = cast<PromotableOpInterface>(toRewire);
172  if (promotable.removeBlockingUses(info.userToBlockingUses[promotable],
173  builder) == DeletionKind::Delete)
174  toErase.push_back(promotable);
175  }
176 
177  for (Operation *toEraseOp : toErase)
178  toEraseOp->erase();
179 
180  assert(slot.ptr.use_empty() && "after destructuring, the original slot "
181  "pointer should no longer be used");
182 
183  LLVM_DEBUG(llvm::dbgs() << "[sroa] Destructured memory slot: " << slot.ptr
184  << "\n");
185 
186  if (statistics.destructuredAmount)
187  (*statistics.destructuredAmount)++;
188 
189  std::optional<DestructurableAllocationOpInterface> newAllocator =
190  allocator.handleDestructuringComplete(slot, builder);
191  // Add newly created allocators to the worklist for further processing.
192  if (newAllocator)
193  newAllocators.push_back(*newAllocator);
194 }
195 
198  OpBuilder &builder, const DataLayout &dataLayout,
199  SROAStatistics statistics) {
200  bool destructuredAny = false;
201 
204  newWorkList.reserve(allocators.size());
205  // Destructuring a slot can allow for further destructuring of other
206  // slots, destructuring is tried until no destructuring succeeds.
207  while (true) {
208  bool changesInThisRound = false;
209 
210  for (DestructurableAllocationOpInterface allocator : workList) {
211  bool destructuredAnySlot = false;
212  for (DestructurableMemorySlot slot : allocator.getDestructurableSlots()) {
213  std::optional<MemorySlotDestructuringInfo> info =
214  computeDestructuringInfo(slot, dataLayout);
215  if (!info)
216  continue;
217 
218  destructureSlot(slot, allocator, builder, dataLayout, *info,
219  newWorkList, statistics);
220  destructuredAnySlot = true;
221 
222  // A break is required, since destructuring a slot may invalidate the
223  // remaning slots of an allocator.
224  break;
225  }
226  if (!destructuredAnySlot)
227  newWorkList.push_back(allocator);
228  changesInThisRound |= destructuredAnySlot;
229  }
230 
231  if (!changesInThisRound)
232  break;
233  destructuredAny |= changesInThisRound;
234 
235  // Swap the vector's backing memory and clear the entries in newWorkList
236  // afterwards. This ensures that additional heap allocations can be avoided.
237  workList.swap(newWorkList);
238  newWorkList.clear();
239  }
240 
241  return success(destructuredAny);
242 }
243 
244 namespace {
245 
246 struct SROA : public impl::SROABase<SROA> {
247  using impl::SROABase<SROA>::SROABase;
248 
249  void runOnOperation() override {
250  Operation *scopeOp = getOperation();
251 
252  SROAStatistics statistics{&destructuredAmount, &slotsWithMemoryBenefit,
253  &maxSubelementAmount};
254 
255  auto &dataLayoutAnalysis = getAnalysis<DataLayoutAnalysis>();
256  const DataLayout &dataLayout = dataLayoutAnalysis.getAtOrAbove(scopeOp);
257  bool changed = false;
258 
259  for (Region &region : scopeOp->getRegions()) {
260  if (region.getBlocks().empty())
261  continue;
262 
263  OpBuilder builder(&region.front(), region.front().begin());
264 
266  // Build a list of allocators to attempt to destructure the slots of.
267  region.walk([&](DestructurableAllocationOpInterface allocator) {
268  allocators.emplace_back(allocator);
269  });
270 
271  // Attempt to destructure as many slots as possible.
272  if (succeeded(tryToDestructureMemorySlots(allocators, builder, dataLayout,
273  statistics)))
274  changed = true;
275  }
276  if (!changed)
277  markAllAnalysesPreserved();
278  }
279 };
280 
281 } // namespace
static void destructureSlot(DestructurableMemorySlot &slot, DestructurableAllocationOpInterface allocator, OpBuilder &builder, const DataLayout &dataLayout, MemorySlotDestructuringInfo &info, SmallVectorImpl< DestructurableAllocationOpInterface > &newAllocators, const SROAStatistics &statistics)
Performs the destructuring of a destructible slot given associated destructuring information.
Definition: SROA.cpp:137
static std::optional< MemorySlotDestructuringInfo > computeDestructuringInfo(DestructurableMemorySlot &slot, const DataLayout &dataLayout)
Computes information for slot destructuring.
Definition: SROA.cpp:47
The main mechanism for performing data layout queries.
RAII guard to reset the insertion point of the builder when destroyed.
Definition: Builders.h:346
This class helps build Operations.
Definition: Builders.h:205
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
Definition: Builders.h:429
void setInsertionPointAfter(Operation *op)
Sets the insertion point to the node after the specified operation, which will cause subsequent inser...
Definition: Builders.h:410
This class represents an operand of an operation.
Definition: Value.h:243
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
MutableArrayRef< Region > getRegions()
Returns the regions held by this operation.
Definition: Operation.h:677
This class contains a list of basic blocks and a link to the parent operation it is attached to.
Definition: Region.h:26
bool use_empty() const
Returns true if this value has no uses.
Definition: Value.h:194
use_range getUses() const
Returns a range of all uses, which is useful for iterating over all uses.
Definition: Value.h:188
Block * getParentBlock()
Return the Block in which this Value is defined.
Definition: Value.cpp:48
Include the generated interface declarations.
const FrozenRewritePatternSet GreedyRewriteConfig bool * changed
@ Delete
Delete the operation after promotion.
LogicalResult tryToDestructureMemorySlots(ArrayRef< DestructurableAllocationOpInterface > allocators, OpBuilder &builder, const DataLayout &dataLayout, SROAStatistics statistics={})
Attempts to destructure the slots of destructurable allocators.
Definition: SROA.cpp:196
SetVector< Operation * > topologicalSort(const SetVector< Operation * > &toSort)
Sorts all operations in toSort topologically while also considering region semantics.
void getForwardSlice(Operation *op, SetVector< Operation * > *forwardSlice, const ForwardSliceOptions &options={})
Fills forwardSlice with the computed forward slice (i.e.
Memory slot attached with information about its destructuring procedure.
DenseMap< Attribute, Type > subelementTypes
Maps an index within the memory slot to the corresponding subelement type.
Represents a slot in memory.
Value ptr
Pointer to the memory slot, used by operations to refer to it.
Type elemType
Type of the value contained in the slot.
Statistics collected while applying SROA.
Definition: SROA.h:19
llvm::Statistic * maxSubelementAmount
Maximal number of sub-elements a successfully destructured slot initially had.
Definition: SROA.h:27
llvm::Statistic * slotsWithMemoryBenefit
Total amount of memory slots in which the destructured size was smaller than the total size after eli...
Definition: SROA.h:24
llvm::Statistic * destructuredAmount
Total amount of memory slots destructured.
Definition: SROA.h:21