llvm-project

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

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.
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.
46
static std::optional<MemorySlotDestructuringInfo>
47
computeDestructuringInfo(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.getOrInsertDefault(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

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(&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
    if (!info.userToBlockingUses.contains(user))
104
      continue;
105

106
    SmallPtrSet<OpOperand *, 4> &blockingUses = info.userToBlockingUses[user];
107
    auto promotable = dyn_cast<PromotableOpInterface>(user);
108

109
    // An operation that has blocking uses must be promoted. If it is not
110
    // promotable, destructuring must fail.
111
    if (!promotable)
112
      return {};
113

114
    SmallVector<OpOperand *> newBlockingUses;
115
    // If the operation decides it cannot deal with removing the blocking uses,
116
    // destructuring must fail.
117
    if (!promotable.canUsesBeRemoved(blockingUses, newBlockingUses, dataLayout))
118
      return {};
119

120
    // Then, register any new blocking uses for coming operations.
121
    for (OpOperand *blockingUse : newBlockingUses) {
122
      assert(llvm::is_contained(user->getResults(), blockingUse->get()));
123

124
      SmallPtrSetImpl<OpOperand *> &newUserBlockingUseSet =
125
          info.userToBlockingUses.getOrInsertDefault(blockingUse->getOwner());
126
      newUserBlockingUseSet.insert(blockingUse);
127
    }
128
  }
129

130
  return info;
131
}
132

133
/// Performs the destructuring of a destructible slot given associated
134
/// destructuring information. The provided slot will be destructured in
135
/// subslots as specified by its allocator.
136
static void destructureSlot(
137
    DestructurableMemorySlot &slot,
138
    DestructurableAllocationOpInterface allocator, OpBuilder &builder,
139
    const DataLayout &dataLayout, MemorySlotDestructuringInfo &info,
140
    SmallVectorImpl<DestructurableAllocationOpInterface> &newAllocators,
141
    const SROAStatistics &statistics) {
142
  OpBuilder::InsertionGuard guard(builder);
143

144
  builder.setInsertionPointToStart(slot.ptr.getParentBlock());
145
  DenseMap<Attribute, MemorySlot> subslots =
146
      allocator.destructure(slot, info.usedIndices, builder, newAllocators);
147

148
  if (statistics.slotsWithMemoryBenefit &&
149
      slot.subelementTypes.size() != info.usedIndices.size())
150
    (*statistics.slotsWithMemoryBenefit)++;
151

152
  if (statistics.maxSubelementAmount)
153
    statistics.maxSubelementAmount->updateMax(slot.subelementTypes.size());
154

155
  SetVector<Operation *> usersToRewire;
156
  for (Operation *user : llvm::make_first_range(info.userToBlockingUses))
157
    usersToRewire.insert(user);
158
  for (DestructurableAccessorOpInterface accessor : info.accessors)
159
    usersToRewire.insert(accessor);
160
  usersToRewire = mlir::topologicalSort(usersToRewire);
161

162
  llvm::SmallVector<Operation *> toErase;
163
  for (Operation *toRewire : llvm::reverse(usersToRewire)) {
164
    builder.setInsertionPointAfter(toRewire);
165
    if (auto accessor = dyn_cast<DestructurableAccessorOpInterface>(toRewire)) {
166
      if (accessor.rewire(slot, subslots, builder, dataLayout) ==
167
          DeletionKind::Delete)
168
        toErase.push_back(accessor);
169
      continue;
170
    }
171

172
    auto promotable = cast<PromotableOpInterface>(toRewire);
173
    if (promotable.removeBlockingUses(info.userToBlockingUses[promotable],
174
                                      builder) == DeletionKind::Delete)
175
      toErase.push_back(promotable);
176
  }
177

178
  for (Operation *toEraseOp : toErase)
179
    toEraseOp->erase();
180

181
  assert(slot.ptr.use_empty() && "after destructuring, the original slot "
182
                                 "pointer should no longer be used");
183

184
  LLVM_DEBUG(llvm::dbgs() << "[sroa] Destructured memory slot: " << slot.ptr
185
                          << "\n");
186

187
  if (statistics.destructuredAmount)
188
    (*statistics.destructuredAmount)++;
189

190
  std::optional<DestructurableAllocationOpInterface> newAllocator =
191
      allocator.handleDestructuringComplete(slot, builder);
192
  // Add newly created allocators to the worklist for further processing.
193
  if (newAllocator)
194
    newAllocators.push_back(*newAllocator);
195
}
196

197
LogicalResult mlir::tryToDestructureMemorySlots(
198
    ArrayRef<DestructurableAllocationOpInterface> allocators,
199
    OpBuilder &builder, const DataLayout &dataLayout,
200
    SROAStatistics statistics) {
201
  bool destructuredAny = false;
202

203
  SmallVector<DestructurableAllocationOpInterface> workList(allocators.begin(),
204
                                                            allocators.end());
205
  SmallVector<DestructurableAllocationOpInterface> newWorkList;
206
  newWorkList.reserve(allocators.size());
207
  // Destructuring a slot can allow for further destructuring of other
208
  // slots, destructuring is tried until no destructuring succeeds.
209
  while (true) {
210
    bool changesInThisRound = false;
211

212
    for (DestructurableAllocationOpInterface allocator : workList) {
213
      bool destructuredAnySlot = false;
214
      for (DestructurableMemorySlot slot : allocator.getDestructurableSlots()) {
215
        std::optional<MemorySlotDestructuringInfo> info =
216
            computeDestructuringInfo(slot, dataLayout);
217
        if (!info)
218
          continue;
219

220
        destructureSlot(slot, allocator, builder, dataLayout, *info,
221
                        newWorkList, statistics);
222
        destructuredAnySlot = true;
223

224
        // A break is required, since destructuring a slot may invalidate the
225
        // remaning slots of an allocator.
226
        break;
227
      }
228
      if (!destructuredAnySlot)
229
        newWorkList.push_back(allocator);
230
      changesInThisRound |= destructuredAnySlot;
231
    }
232

233
    if (!changesInThisRound)
234
      break;
235
    destructuredAny |= changesInThisRound;
236

237
    // Swap the vector's backing memory and clear the entries in newWorkList
238
    // afterwards. This ensures that additional heap allocations can be avoided.
239
    workList.swap(newWorkList);
240
    newWorkList.clear();
241
  }
242

243
  return success(destructuredAny);
244
}
245

246
namespace {
247

248
struct SROA : public impl::SROABase<SROA> {
249
  using impl::SROABase<SROA>::SROABase;
250

251
  void runOnOperation() override {
252
    Operation *scopeOp = getOperation();
253

254
    SROAStatistics statistics{&destructuredAmount, &slotsWithMemoryBenefit,
255
                              &maxSubelementAmount};
256

257
    auto &dataLayoutAnalysis = getAnalysis<DataLayoutAnalysis>();
258
    const DataLayout &dataLayout = dataLayoutAnalysis.getAtOrAbove(scopeOp);
259
    bool changed = false;
260

261
    for (Region &region : scopeOp->getRegions()) {
262
      if (region.getBlocks().empty())
263
        continue;
264

265
      OpBuilder builder(&region.front(), region.front().begin());
266

267
      SmallVector<DestructurableAllocationOpInterface> allocators;
268
      // Build a list of allocators to attempt to destructure the slots of.
269
      region.walk([&](DestructurableAllocationOpInterface allocator) {
270
        allocators.emplace_back(allocator);
271
      });
272

273
      // Attempt to destructure as many slots as possible.
274
      if (succeeded(tryToDestructureMemorySlots(allocators, builder, dataLayout,
275
                                                statistics)))
276
        changed = true;
277
    }
278
    if (!changed)
279
      markAllAnalysesPreserved();
280
  }
281
};
282

283
} // namespace
284