[Mlir-commits] [mlir] [mlir][func]-Add deduplicate funcOp arguments transform (PR #158266)

[Amir Bishara]

================
@@ -90,32 +144,86 @@ func::replaceFuncWithNewOrder(RewriterBase &rewriter, func::FuncOp funcOp,
 }
 
 func::CallOp
-func::replaceCallOpWithNewOrder(RewriterBase &rewriter, func::CallOp callOp,
-                                ArrayRef<unsigned> newArgsOrder,
-                                ArrayRef<unsigned> newResultsOrder) {
+func::replaceCallOpWithNewMapping(RewriterBase &rewriter, func::CallOp callOp,
+                                  llvm::ArrayRef<unsigned> oldArgToNewArg,
+                                  llvm::ArrayRef<unsigned> oldResToNewRes) {
   assert(
-      callOp.getNumOperands() == newArgsOrder.size() &&
-      "newArgsOrder must match the number of operands in the call operation");
+      callOp.getNumOperands() == oldArgToNewArg.size() &&
+      "oldArgToNewArg must match the number of operands in the call operation");
   assert(
-      callOp.getNumResults() == newResultsOrder.size() &&
-      "newResultsOrder must match the number of results in the call operation");
+      callOp.getNumResults() == oldResToNewRes.size() &&
+      "oldResToNewRes must match the number of results in the call operation");
+
+  // Inverse mapping from new arguments to old arguments.
+  unsigned numOfNewArgs = 0;
+  auto maxNewArgIdx = llvm::max_element(oldArgToNewArg);
+  if (maxNewArgIdx != oldArgToNewArg.end())
+    numOfNewArgs = *maxNewArgIdx + 1;
+  llvm::SmallVector<llvm::SmallVector<unsigned>> newArgIdxToOldArgIdxs(
+      numOfNewArgs);
+  for (auto [oldArgIdx, newArgIdx] : llvm::enumerate(oldArgToNewArg))
+    newArgIdxToOldArgIdxs[newArgIdx].push_back(oldArgIdx);
+
   SmallVector<Value> newArgsOrderValues;
-  for (unsigned int argIdx : newArgsOrder)
-    newArgsOrderValues.push_back(callOp.getOperand(argIdx));
+  for (const auto &[newArgIdx, oldArgIdxs] :
+       llvm::enumerate(newArgIdxToOldArgIdxs)) {
+    std::ignore = newArgIdx;
----------------
amirBish wrote:

Right, fixed it.

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