[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:

sounds good, comment added, thank you!

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