[Mlir-commits] [mlir] Allowing RDV to call `getArgOperandsMutable()` (PR #160415)

[Francisco Geiman Thiesen]

================
@@ -754,17 +756,51 @@ static void cleanUpDeadVals(RDVFinalCleanupList &list) {
     // Some functions may not allow erasing arguments or results. These calls
     // return failure in such cases without modifying the function, so it's okay
     // to proceed.
-    (void)f.funcOp.eraseArguments(f.nonLiveArgs);
+    if (succeeded(f.funcOp.eraseArguments(f.nonLiveArgs))) {
+      // Record only if we actually erased something.
+      if (f.nonLiveArgs.any())
+        erasedFuncArgs.try_emplace(f.funcOp.getOperation(), f.nonLiveArgs);
+    }
     (void)f.funcOp.eraseResults(f.nonLiveRets);
   }
 
   // 4. Operands
   LDBG() << "Cleaning up " << list.operands.size() << " operand lists";
   for (OperationToCleanup &o : list.operands) {
-    if (o.op->getNumOperands() > 0) {
-      LDBG() << "Erasing " << o.nonLive.count()
-             << " non-live operands from operation: "
-             << OpWithFlags(o.op, OpPrintingFlags().skipRegions());
+    if (auto call = dyn_cast<CallOpInterface>(o.op)) {
+      if (SymbolRefAttr sym = call.getCallableForCallee().dyn_cast<SymbolRefAttr>()) {
+        Operation *callee = SymbolTable::lookupNearestSymbolFrom(o.op, sym);
+        auto it = erasedFuncArgs.find(callee);
+        if (it != erasedFuncArgs.end()) {
+          const BitVector &deadArgIdxs = it->second;
+          MutableOperandRange args = call.getArgOperandsMutable();
+          // First, erase the call arguments corresponding to erased callee args.
+          for (int i = static_cast<int>(args.size()) - 1; i >= 0; --i) {
+            if (i < static_cast<int>(deadArgIdxs.size()) && deadArgIdxs.test(i))
----------------
FranciscoThiesen wrote:

Agree that the syntax would be better, but I've tried this but found out that:

`llvm::const_set_bits_iterator_impl`seems to only have a forward iterator, so `llvm_reverse` doesn't work

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