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

[llvmlistbot at llvm.org]

================
@@ -447,35 +514,96 @@ MemorySlotPromotionAnalyzer::computeInfo() {
   return info;
 }
 
-Value MemorySlotPromoter::computeReachingDefInBlock(Block *block,
-                                                    Value reachingDef) {
+Value MemorySlotPromoter::promoteInBlock(Block *block, Value reachingDef) {
+  llvm::SmallMapVector<Region *, Value, 2> regionsToProcess;
   SmallVector<Operation *> blockOps;
   for (Operation &op : block->getOperations())
     blockOps.push_back(&op);
   for (Operation *op : blockOps) {
+    // Promote operations that interact with the slot's memory.
     if (auto memOp = dyn_cast<PromotableMemOpInterface>(op)) {
-      if (info.userToBlockingUses.contains(memOp))
+      if (info.userToBlockingUses[memOp->getParentRegion()].contains(memOp))
         reachingDefs.insert({memOp, reachingDef});
 
       if (memOp.storesTo(slot)) {
         builder.setInsertionPointAfter(memOp);
+        // To not expose default value creation to the interfaces, if we have
+        // no reaching definition by now, we set it to the default value.
+        // This is slightly too eager as `getStored` may not need it.
+        if (!reachingDef)
+          reachingDef = getOrCreateDefaultValue();
         Value stored = memOp.getStored(slot, builder, reachingDef, dataLayout);
         assert(stored && "a memory operation storing to a slot must provide a "
                          "new definition of the slot");
         reachingDef = stored;
         replacedValuesMap[memOp] = stored;
       }
     }
+
+    // Promote regions that contain operations that interact with the slot's
+    // memory.
+    if (auto promotableRegionOp = dyn_cast<PromotableRegionOpInterface>(op)) {
+      bool needsPromotion = false;
+      bool hasValueStores = false;
+      for (Region &region : op->getRegions()) {
+        auto regionInfoIt = info.regionsToPromote.find(&region);
+        if (regionInfoIt == info.regionsToPromote.end())
+          continue;
+        needsPromotion = true;
+        if (!regionInfoIt->second.hasValueStores)
+          continue;
+
+        hasValueStores = true;
+        break;
+      }
+
+      if (needsPromotion) {
+        regionsToProcess.clear();
----------------
tdegioanni-nvidia wrote:

This is a silly optimization I did in the case where there would have many regions and there would be a need to allocate, to not have to reallocate the buffer for the next iteration if it were to happen again. But let's remove it, it's mega unlikely to be useful.

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