Address review feedbacks - 3

Author: giuseros

mlir/lib/Transforms/Utils/RegionUtils.cpp

@@ -715,30 +715,29 @@ static SmallVector<SmallVector<Value, 8>, 2> pruneRedundantArguments(
 
   // 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. Please
     // note that we are using a map to annotate when an argument was seen first
     // to avoid a O(N^2) algorithm. This has the drawback that if we have two
     // lists like:
     // list0: [%a, %a, %a]
     // list1: [%c, %b, %b]
-    // We cannot simplify it, because firstVlaueToIdx[%a] = 0, but we cannot
+    // We cannot simplify it, because firstValueToIdx[%a] = 0, but we cannot
     // point list1[1](==%b) or list1[2](==%b) to list1[0](==%c).  However, since
     // the number of arguments can be potentially unbounded we cannot afford a
     // O(N^2) algorithm (to search to all the possible pairs) and we need to
     // accept the trade-off.
     unsigned k = firstValueToIdx[newArguments[0][j]];
-    if (k != 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)
-        shouldReplaceJ =
-            shouldReplaceJ && (newArguments[i][k] == newArguments[i][j]);
-    }
+    if (k == j)
+      continue;
+
+    bool shouldReplaceJ = true;
+    unsigned 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)
+      shouldReplaceJ =
+          shouldReplaceJ && (newArguments[i][k] == newArguments[i][j]);
     // Save the replacement.
     if (shouldReplaceJ)
       idxToReplacement[j] = replacement;
@@ -918,6 +917,8 @@ static LogicalResult mergeIdenticalBlocks(RewriterBase &rewriter,
   return success(anyChanged);
 }
 
+/// If a block's argument is always the same across different invocations, then
+/// drop the argument and use the value directly inside the block
 static LogicalResult dropRedundantArguments(RewriterBase &rewriter,
                                             Block &block) {
   SmallVector<size_t> argsToErase;
@@ -929,7 +930,7 @@ static LogicalResult dropRedundantArguments(RewriterBase &rewriter,
 
     // Go through the block predecessor and flag if they pass to the block
     // different values for the same argument.
-    for (auto predIt = block.pred_begin(), predE = block.pred_end();
+    for (Block::pred_iterator predIt = block.pred_begin(), predE = block.pred_end();
          predIt != predE; ++predIt) {
       auto branch = dyn_cast<BranchOpInterface>((*predIt)->getTerminator());
       if (!branch) {
@@ -941,11 +942,11 @@ static LogicalResult dropRedundantArguments(RewriterBase &rewriter,
       auto branchOperands = succOperands.getForwardedOperands();
       if (!commonValue) {
         commonValue = branchOperands[argIdx];
-      } else {
-        if (branchOperands[argIdx] != commonValue) {
-          sameArg = false;
-          break;
-        }
+        continue;
+      }
+      if (branchOperands[argIdx] != commonValue) {
+        sameArg = false;
+        break;
       }
     }
 
@@ -959,7 +960,7 @@ static LogicalResult dropRedundantArguments(RewriterBase &rewriter,
   }
 
   // Remove the arguments.
-  for (auto argIdx : llvm::reverse(argsToErase)) {
+  for (size_t argIdx : llvm::reverse(argsToErase)) {
     block.eraseArgument(argIdx);
 
     // Remove the argument from the branch ops.