[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();
+
+        // 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 `setupPromotion` may not need it.
+        if (!reachingDef)
+          reachingDef = getOrCreateDefaultValue();
+
+        promotableRegionOp.setupPromotion(slot, reachingDef, hasValueStores,
+                                          regionsToProcess);
+
+#ifndef NDEBUG
+        for (Region &region : op->getRegions())
+          if (info.regionsToPromote.contains(&region))
+            assert(
+                regionsToProcess.contains(&region) &&
+                "reaching definition must be provided for a required region");
+#endif // NDEBUG
+
+        for (auto &[region, reachingDef] : regionsToProcess) {
+#ifndef NDEBUG
+          Region *regionCapture = region;
+          assert(llvm::any_of(op->getRegions(),
+                              [&](Region &r) { return &r == regionCapture; }) &&
----------------
tdegioanni-nvidia wrote:

You are smarter than me Mr Gysi.

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