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

[Amir Bishara]

================
@@ -37,12 +37,49 @@ func::replaceFuncWithNewOrder(RewriterBase &rewriter, func::FuncOp funcOp,
   ArrayRef<Type> origOutputTypes = funcOp.getFunctionType().getResults();
   SmallVector<Type> newInputTypes, newOutputTypes;
   SmallVector<Location> locs;
-  for (unsigned int idx : newArgsOrder) {
-    newInputTypes.push_back(origInputTypes[idx]);
-    locs.push_back(funcOp.getArgument(newArgsOrder[idx]).getLoc());
+
+  std::map<unsigned, SmallVector<unsigned>> newArgToOldArg;
+  for (auto [oldArgIdx, newArgIdx] : oldArgToNewArg)
+    newArgToOldArg[newArgIdx].push_back(oldArgIdx);
+
+  for (auto [newArgIdx, oldArgIdx] : newArgToOldArg) {
+    std::ignore = newArgIdx;
+    assert(llvm::all_of(oldArgIdx,
+                        [&funcOp](unsigned idx) -> bool {
+                          return idx < funcOp.getNumArguments();
+                        }) &&
+           "idx must be less than the number of arguments in the function");
+    assert(!oldArgIdx.empty() && "oldArgIdx must not be empty");
+    Type origInputTypeToCheck = origInputTypes[oldArgIdx.front()];
+    assert(llvm::all_of(oldArgIdx,
+                        [&](unsigned idx) -> bool {
+                          return origInputTypes[idx] == origInputTypeToCheck;
+                        }) &&
+           "all oldArgIdx must have the same type");
+    newInputTypes.push_back(origInputTypeToCheck);
+    locs.push_back(funcOp.getArgument(oldArgIdx.front()).getLoc());
+  }
+
+  std::map<unsigned, SmallVector<unsigned>> newResToOldRes;
+  for (auto [oldResIdx, newResIdx] : oldResToNewRes)
+    newResToOldRes[newResIdx].push_back(oldResIdx);
----------------
amirBish wrote:

Answered above :)

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