[Mlir-commits] [mlir] [MLIR][SROA] Reuse allocators to avoid rewalking the IR (PR #91971)

[Tobias Gysi]

================
@@ -1199,35 +1199,90 @@ void TestMultiSlotAlloca::handleBlockArgument(const MemorySlot &slot,
   // Not relevant for testing.
 }
 
-std::optional<PromotableAllocationOpInterface>
-TestMultiSlotAlloca::handlePromotionComplete(const MemorySlot &slot,
-                                             Value defaultValue,
-                                             OpBuilder &builder) {
-  if (defaultValue && defaultValue.use_empty())
-    defaultValue.getDefiningOp()->erase();
+/// Creates a new TestMultiSlotAlloca operation, just without the `slot`.
+static std::optional<TestMultiSlotAlloca>
+createNewMultiAllocaWithoutSlot(const MemorySlot &slot, OpBuilder &builder,
+                                TestMultiSlotAlloca oldOp) {
 
-  if (getNumResults() == 1) {
-    erase();
+  if (oldOp.getNumResults() == 1) {
+    oldOp.erase();
     return std::nullopt;
   }
 
   SmallVector<Type> newTypes;
   SmallVector<Value> remainingValues;
 
-  for (Value oldResult : getResults()) {
+  for (Value oldResult : oldOp.getResults()) {
     if (oldResult == slot.ptr)
       continue;
     remainingValues.push_back(oldResult);
     newTypes.push_back(oldResult.getType());
   }
 
   OpBuilder::InsertionGuard guard(builder);
-  builder.setInsertionPoint(*this);
-  auto replacement = builder.create<TestMultiSlotAlloca>(getLoc(), newTypes);
+  builder.setInsertionPoint(oldOp);
+  auto replacement =
+      builder.create<TestMultiSlotAlloca>(oldOp->getLoc(), newTypes);
   for (auto [oldResult, newResult] :
        llvm::zip_equal(remainingValues, replacement.getResults()))
     oldResult.replaceAllUsesWith(newResult);
 
-  erase();
+  oldOp.erase();
   return replacement;
 }
+
+std::optional<PromotableAllocationOpInterface>
+TestMultiSlotAlloca::handlePromotionComplete(const MemorySlot &slot,
+                                             Value defaultValue,
+                                             OpBuilder &builder) {
+  if (defaultValue && defaultValue.use_empty())
+    defaultValue.getDefiningOp()->erase();
+  return createNewMultiAllocaWithoutSlot(slot, builder, *this);
+}
+
+SmallVector<DestructurableMemorySlot>
+TestMultiSlotAlloca::getDestructurableSlots() {
+  SmallVector<DestructurableMemorySlot> slots;
+  for (Value result : getResults()) {
+    auto memrefType = cast<MemRefType>(result.getType());
+    auto destructurable =
+        llvm::dyn_cast<DestructurableTypeInterface>(memrefType);
----------------
gysit wrote:

```suggestion
        dyn_cast<DestructurableTypeInterface>(memrefType);
```
nit:

https://github.com/llvm/llvm-project/pull/91971