[Mlir-commits] [mlir] [mlir] Fix block merging (PR #97697)

[Jeff Niu]

================
@@ -674,6 +679,74 @@ static bool ableToUpdatePredOperands(Block *block) {
   return true;
 }
 
+/// Prunes the redundant list of arguments. E.g., if we are passing an argument
+/// list like [x, y, z, x] this would return [x, y, z] and it would update the
+/// `block` (to whom the argument are passed to) accordingly.
+static SmallVector<SmallVector<Value, 8>, 2> pruneRedundantArguments(
+    const SmallVector<SmallVector<Value, 8>, 2> &newArguments,
+    RewriterBase &rewriter, Block *block) {
+
+  SmallVector<SmallVector<Value, 8>, 2> newArgumentsPruned(
+      newArguments.size(), SmallVector<Value, 8>());
+
+  if (newArguments.empty())
+    return newArguments;
+
+  // `newArguments` is a 2D array of size `numLists` x `numArgs`
+  unsigned numLists = newArguments.size();
+  unsigned numArgs = newArguments[0].size();
+
+  // Map that for each arg index contains the index that we can use in place of
+  // the original index. E.g., if we have newArgs = [x, y, z, x], we will have
+  // idxToReplacement[3] = 0
+  llvm::DenseMap<unsigned, unsigned> idxToReplacement;
+
+  // Go through the first list of arguments (list 0).
+  for (unsigned j = 0; j < numArgs; ++j) {
+    bool shouldReplaceJ = false;
+    unsigned replacement = 0;
+    // Look back to see if there are possible redundancies in
+    // list 0.
+    for (unsigned k = 0; k < j; k++) {
+      if (newArguments[0][k] == newArguments[0][j]) {
+        shouldReplaceJ = true;
+        replacement = k;
+        // If a possible redundancy is found, then scan the other lists: we
+        // can prune the arguments if and only if they are redundant in every
+        // list.
+        for (unsigned i = 1; i < numLists; ++i)
----------------
Mogball wrote:

>  and numArgs is rarely bigger than 10
This is a great catch on @krzysz00's part! I actually find that the number of arguments in practice can be unbound. We've had this same algorithmic issue before with the number of call arguments in MLIR

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