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"
10#include "mlir/Analysis/DataLayoutAnalysis.h"
11#include "mlir/Analysis/SliceAnalysis.h"
12#include "mlir/Analysis/TopologicalSortUtils.h"
13#include "mlir/Interfaces/MemorySlotInterfaces.h"
14#include "mlir/Transforms/Passes.h"
15
16namespace mlir {
17#define GEN_PASS_DEF_SROA
18#include "mlir/Transforms/Passes.h.inc"
19} // namespace mlir
20
21#define DEBUG_TYPE "sroa"
22
23using namespace mlir;
24
25namespace {
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.
30struct 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.
34 DenseMap<Operation *, SmallPtrSet<OpOperand *, 4>> userToBlockingUses;
35 /// List of potentially indirect accessors of the memory slot that need
36 /// rewiring.
37 SmallVector<DestructurableAccessorOpInterface> accessors;
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.
46static std::optional<MemorySlotDestructuringInfo>
47computeDestructuringInfo(DestructurableMemorySlot &slot,
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(Ptr: &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
80 SmallPtrSet<OpOperand *, 16> visited;
81 while (!usedSafelyWorklist.empty()) {
82 MemorySlot mustBeUsedSafely = usedSafelyWorklist.pop_back_val();
83 for (OpOperand &subslotUse : mustBeUsedSafely.ptr.getUses()) {
84 if (!visited.insert(Ptr: &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(root: slot.ptr, forwardSlice: &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(Ptr: 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.
137static void destructureSlot(
138 DestructurableMemorySlot &slot,
139 DestructurableAllocationOpInterface allocator, OpBuilder &builder,
140 const DataLayout &dataLayout, MemorySlotDestructuringInfo &info,
141 SmallVectorImpl<DestructurableAllocationOpInterface> &newAllocators,
142 const SROAStatistics &statistics) {
143 OpBuilder::InsertionGuard guard(builder);
144
145 builder.setInsertionPointToStart(slot.ptr.getParentBlock());
146 DenseMap<Attribute, MemorySlot> subslots =
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(V: 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(toSort: usersToRewire);
160
161 llvm::SmallVector<Operation *> toErase;
162 for (Operation *toRewire : llvm::reverse(C&: usersToRewire)) {
163 builder.setInsertionPointAfter(toRewire);
164 if (auto accessor = dyn_cast<DestructurableAccessorOpInterface>(toRewire)) {
165 if (accessor.rewire(slot, subslots, builder, dataLayout) ==
166 DeletionKind::Delete)
167 toErase.push_back(Elt: 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(Elt: 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
196LogicalResult mlir::tryToDestructureMemorySlots(
197 ArrayRef<DestructurableAllocationOpInterface> allocators,
198 OpBuilder &builder, const DataLayout &dataLayout,
199 SROAStatistics statistics) {
200 bool destructuredAny = false;
201
202 SmallVector<DestructurableAllocationOpInterface> workList(allocators);
203 SmallVector<DestructurableAllocationOpInterface> newWorkList;
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(IsSuccess: destructuredAny);
242}
243
244namespace {
245
246struct 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
265 SmallVector<DestructurableAllocationOpInterface> allocators;
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
282

source code of mlir/lib/Transforms/SROA.cpp