[Mlir-commits] [mlir] [MLIR][Mem2Reg] Add support for region control flow and SCF (PR #185036)

[Tobias Gysi]

================
@@ -559,79 +674,121 @@ static void dominanceSort(SmallVector<Operation *> &ops, Region &region,
   });
 }
 
-void MemorySlotPromoter::removeBlockingUses() {
+void MemorySlotPromoter::removeBlockingUses(Region *region) {
+  auto *blockingUsesMapIt = info.userToBlockingUses.find(region);
+  if (blockingUsesMapIt == info.userToBlockingUses.end())
+    return;
+  BlockingUsesMap &blockingUsesMap = blockingUsesMapIt->second;
+
   llvm::SmallVector<Operation *> usersToRemoveUses(
-      llvm::make_first_range(info.userToBlockingUses));
+      llvm::make_first_range(blockingUsesMap));
 
   // Sort according to dominance.
-  dominanceSort(usersToRemoveUses, *slot.ptr.getParentBlock()->getParent(),
-                blockIndexCache);
+  dominanceSort(usersToRemoveUses, *region, blockIndexCache);
 
-  llvm::SmallVector<Operation *> toErase;
-  // List of all replaced values in the slot.
-  llvm::SmallVector<std::pair<Operation *, Value>> replacedValuesList;
-  // Ops to visit with the `visitReplacedValues` method.
-  llvm::SmallVector<PromotableOpInterface> toVisit;
   for (Operation *toPromote : llvm::reverse(usersToRemoveUses)) {
     if (auto toPromoteMemOp = dyn_cast<PromotableMemOpInterface>(toPromote)) {
       Value reachingDef = reachingDefs.lookup(toPromoteMemOp);
       // If no reaching definition is known, this use is outside the reach of
       // the slot. The default value should thus be used.
+      // FIXME: This is too eager, and will generate default values even for
+      // pure stores. This cannot be removed easily as partial stores may
+      // still require a default value to complete.
       if (!reachingDef)
         reachingDef = getOrCreateDefaultValue();
 
       builder.setInsertionPointAfter(toPromote);
-      if (toPromoteMemOp.removeBlockingUses(
-              slot, info.userToBlockingUses[toPromote], builder, reachingDef,
-              dataLayout) == DeletionKind::Delete)
+      if (toPromoteMemOp.removeBlockingUses(slot, blockingUsesMap[toPromote],
+                                            builder, reachingDef,
+                                            dataLayout) == DeletionKind::Delete)
         toErase.push_back(toPromote);
       if (toPromoteMemOp.storesTo(slot))
         if (Value replacedValue = replacedValuesMap[toPromoteMemOp])
-          replacedValuesList.push_back({toPromoteMemOp, replacedValue});
+          replacedValues.push_back({toPromoteMemOp, replacedValue});
       continue;
     }
 
     auto toPromoteBasic = cast<PromotableOpInterface>(toPromote);
     builder.setInsertionPointAfter(toPromote);
-    if (toPromoteBasic.removeBlockingUses(info.userToBlockingUses[toPromote],
+    if (toPromoteBasic.removeBlockingUses(blockingUsesMap[toPromote],
                                           builder) == DeletionKind::Delete)
       toErase.push_back(toPromote);
     if (toPromoteBasic.requiresReplacedValues())
-      toVisit.push_back(toPromoteBasic);
+      toVisitReplacedValues.push_back(toPromoteBasic);
   }
-  for (PromotableOpInterface op : toVisit) {
-    builder.setInsertionPointAfter(op);
-    op.visitReplacedValues(replacedValuesList, builder);
+}
+
+void MemorySlotPromoter::linkMergePoints() {
+  // We want to eliminate unused block arguments. In case connecting a block
+  // argument to its predecessor would trigger the use of the predecessor's
+  // unused block argument, we need to process merge points in an expanding
+  // worklist, `mergePointArgsToProcess`.
+
+  SmallPtrSet<BlockArgument, 8> mergePointArgsUnused;
+  SmallVector<BlockArgument> mergePointArgsToProcess;
+  for (Block *mergePoint : info.mergePoints) {
+    BlockArgument arg = mergePoint->getArguments().back();
+    if (arg.use_empty())
+      mergePointArgsUnused.insert(arg);
+    else
+      mergePointArgsToProcess.push_back(arg);
   }
 
-  for (Operation *toEraseOp : toErase)
-    toEraseOp->erase();
+  while (!mergePointArgsToProcess.empty()) {
+    BlockArgument arg = mergePointArgsToProcess.pop_back_val();
+    Block *mergePoint = arg.getOwner();
 
-  assert(slot.ptr.use_empty() &&
-         "after promotion, the slot pointer should not be used anymore");
+    for (BlockOperand &use : mergePoint->getUses()) {
+      Value reachingDef = reachingAtBlockEnd[use.getOwner()->getBlock()];
+      if (!reachingDef)
+        reachingDef = getOrCreateDefaultValue();
+
+      // If the reaching definition is a block argument of an unused merge
+      // point, mark it as used and process it as such later.
+      auto reachingDefArgument = dyn_cast<BlockArgument>(reachingDef);
+      if (reachingDefArgument &&
+          mergePointArgsUnused.erase(reachingDefArgument))
----------------
gysit wrote:

Are this the region arguments introduced when promoting the region operations? 

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