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

[Renat Idrisov]

================
@@ -588,22 +661,68 @@ static void cleanBranchOp(BranchOpInterface branchOp, RunLivenessAnalysis &la) {
       operandValues.push_back(successorOperands[operandIdx]);
     }
 
-    BitVector successorLiveOperands = markLives(operandValues, la);
-
     // Do (3)
-    for (int argIdx = successorLiveOperands.size() - 1; argIdx >= 0; --argIdx) {
-      if (!successorLiveOperands[argIdx]) {
-        if (successorBlock->getNumArguments() < successorOperands.size()) {
-          // if block was cleaned through a different code path
-          // we only need to remove operands from the invokation
-          successorOperands.erase(argIdx);
-          continue;
-        }
+    BitVector successorNonLive =
+        markLives(operandValues, nonLiveSet, la).flip();
+    collectNonLiveValues(nonLiveSet, successorBlock->getArguments(),
+                         successorNonLive);
+    cl.blocks.push_back({successorBlock, successorNonLive});
+    cl.successorOperands.push_back({branchOp, succIdx, successorNonLive});
+  }
+}
+
+static void cleanUpDeadVals(RDVFinalCleanupList &RDVFinalCleanupList) {
+  // 1. Operations
+  for (auto &op : RDVFinalCleanupList.operations) {
+    op->dropAllUses();
+    op->erase();
+  }
+
+  // 2. Values
+  for (auto &v : RDVFinalCleanupList.values) {
+    v.dropAllUses();
+  }
+
+  // 3. Functions
+  for (auto &f : RDVFinalCleanupList.functions) {
+    f.funcOp.eraseArguments(f.nonLiveArgs);
+    f.funcOp.eraseResults(f.nonLiveRets);
+  }
+
+  // 4. Operands
+  for (auto &o : RDVFinalCleanupList.operands) {
+    o.op->eraseOperands(o.nonLive);
+  }
 
-        successorBlock->getArgument(argIdx).dropAllUses();
-        successorOperands.erase(argIdx);
-        successorBlock->eraseArgument(argIdx);
-      }
+  // 5. Results
+  for (auto &r : RDVFinalCleanupList.results) {
+    dropUsesAndEraseResults(r.op, r.nonLive);
+  }
+
+  // 6. Blocks
+  for (auto &b : RDVFinalCleanupList.blocks) {
+    // blocks that are accessed via multiple codepaths processed once
+    if (b.b->getNumArguments() != b.nonLiveArgs.size())
+      continue;
+    for (int i = b.nonLiveArgs.size() - 1; i >= 0; --i) {
+      if (!b.nonLiveArgs[i])
+        continue;
+      b.b->getArgument(i).dropAllUses();
+      b.b->eraseArgument(i);
+    }
----------------
parsifal-47 wrote:

the collection of arguments is changed while iterating, if we iterate from start, it would be invalid because each deletion is invalidating all successor indexes, that is why iteration is reversed

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