[Mlir-commits] [mlir] [MLIR] Prevent invalid IR from being passed outside of RemoveDeadValues (PR #121079)

[Andrzej Warzyński]

================
@@ -538,48 +607,48 @@ static void cleanRegionBranchOp(RegionBranchOpInterface regionBranchOp,
                    terminatorOperandsToKeep);
 
   // Do (1).
-  regionBranchOp->eraseOperands(operandsToKeep.flip());
+  cl.operands.push_back({regionBranchOp, operandsToKeep.flip()});
 
   // Do (2.a) and (2.b).
   for (Region &region : regionBranchOp->getRegions()) {
     assert(!region.empty() && "expected a non-empty region in an op "
                               "implementing `RegionBranchOpInterface`");
-    for (auto [index, arg] : llvm::enumerate(region.front().getArguments())) {
-      if (argsToKeep[&region][index])
-        continue;
-      if (arg)
-        arg.dropAllUses();
-    }
-    region.front().eraseArguments(argsToKeep[&region].flip());
+    BitVector argsToRemove = argsToKeep[&region].flip();
+    cl.blocks.push_back({&region.front(), argsToRemove});
+    collectNonLiveValues(nonLiveSet, region.front().getArguments(),
+                         argsToRemove);
   }
 
   // Do (2.c).
   for (Region &region : regionBranchOp->getRegions()) {
     Operation *terminator = region.front().getTerminator();
-    terminator->eraseOperands(terminatorOperandsToKeep[terminator].flip());
+    cl.operands.push_back(
+        {terminator, terminatorOperandsToKeep[terminator].flip()});
   }
 
   // Do (3) and (4).
-  dropUsesAndEraseResults(regionBranchOp.getOperation(), resultsToKeep.flip());
+  BitVector resultsToRemove = resultsToKeep.flip();
+  collectNonLiveValues(nonLiveSet, regionBranchOp.getOperation()->getResults(),
+                       resultsToRemove);
+  cl.results.push_back({regionBranchOp.getOperation(), resultsToRemove});
 }
 
-// 1. Iterate over each successor block of the given BranchOpInterface
-//    operation.
-// 2. For each successor block:
-//    a. Retrieve the operands passed to the successor.
-//    b. Use the provided liveness analysis (`RunLivenessAnalysis`) to determine
-//       which operands are live in the successor block.
-//    c. Mark each operand as live or dead based on the analysis.
-// 3. Remove dead operands from the branch operation and arguments accordingly
+/// Steps to process a `BranchOpInterface` operation:
+/// Iterate through each successor block of the operation.
+/// (1) For each successor block, gather all operands from all successors.
+/// (2) Fetch their associated liveness analysis data and collect for future
+/// removal. (3) Identify and collect the dead operands from the successor block
+///     as well as their corresponding arguments.
----------------
banach-space wrote:

[nit] Missing new-line.
```suggestion
/// removal. 
/// (3) Identify and collect the dead operands from the successor block
///      as well as their corresponding arguments.
```

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