[Mlir-commits] [mlir] [MLIR] [Mem2Reg] Fix unused block argument removal logic (PR #188484)

[Tobias Gysi]

=?utf-8?q?Théo?= Degioanni <tdegioanni at nvidia.com>,
=?utf-8?q?Théo?= Degioanni <tdegioanni at nvidia.com>,
=?utf-8?q?Théo?= Degioanni <tdegioanni at nvidia.com>,
=?utf-8?q?Théo?= Degioanni <tdegioanni at nvidia.com>
Message-ID:
In-Reply-To: <llvm.org/llvm/llvm-project/pull/188484 at github.com>


================
@@ -757,46 +769,96 @@ void MemorySlotPromoter::removeBlockingUses(Region *region) {
     builder.setInsertionPointAfter(toPromote);
     if (toPromoteBasic.removeBlockingUses(blockingUsesMap[toPromote],
                                           builder) == DeletionKind::Delete)
-      toErase.push_back(toPromote);
+      toErase.insert(toPromote);
     if (toPromoteBasic.requiresReplacedValues())
       toVisitReplacedValues.push_back(toPromoteBasic);
   }
 }
 
-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`.
+void MemorySlotPromoter::removeUnusedItems() {
+  // We want to eliminate unused block arguments. Because block arguments can be
+  // used to populate other block arguments, there might be cycles of arguments
+  // that are only used to populate each-other. We therefore need a small
+  // dataflow analysis to identify which block arguments are truly used.
 
   SmallPtrSet<BlockArgument, 8> mergePointArgsUnused;
-  SmallVector<BlockArgument> mergePointArgsToProcess;
+  SmallVector<BlockArgument> usedMergePointArgsToProcess;
+
+  // First, separate the block arguments that are not used or only used for the
+  // purpose of populating a merge point block argument from the others. These
+  // block arguments are potentially unused. Meanwhile, arguments that are
+  // definitely used will be the starting point of the propagation of the
+  // analysis.
+  auto isDefinitelyUsed = [&](BlockArgument arg) {
+    for (auto &use : arg.getUses()) {
+      if (llvm::is_contained(toErase, use.getOwner()))
+        continue;
+
+      // We now want to detect whether the use is to populate a merge point
+      // block argument. If it is not, the argument is definitely used.
+
+      auto branchOp = dyn_cast<BranchOpInterface>(use.getOwner());
+      if (!branchOp)
+        return true;
+
+      auto successorArgument =
+          branchOp.getSuccessorBlockArgument(use.getOperandNumber());
+      if (!successorArgument.has_value())
----------------
gysit wrote:

```suggestion
      if (!successorArgument)
```
nit: I believe the `has_value` is redundant assuming this is an optional? Also maybe we can spell out the type on line 804?

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