Improve implementation to note use walk

Author: fairywreath

mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp

@@ -1020,39 +1020,33 @@ struct FuncOpVectorUnroll final : OpRewritePattern<func::FuncOp> {
     SmallVector<Location> locs(convertedTypes.size(), newFuncOp.getLoc());
     entryBlock.addArguments(convertedTypes, locs);
 
-    // Replace the placeholder values with the new arguments.
-    size_t unrolledInputIdx = 0;
-    newFuncOp.walk([&](Operation *op) {
-      // We first look for operands that are placeholders for initially legal
-      // arguments.
-      for (auto [operandIdx, operandVal] : llvm::enumerate(op->getOperands())) {
-        Operation *operandOp = operandVal.getDefiningOp();
-        if (auto it = tmpOps.find(operandOp); it != tmpOps.end()) {
-          size_t idx = operandIdx;
-          rewriter.modifyOpInPlace(op, [&] {
-            op->setOperand(idx, newFuncOp.getArgument(it->second));
-          });
-        }
-      }
-
-      // Only consider `vector.insert_strided_slice` ops that were newly created
-      // at the beginning of the entry block. Once we encounter operations
-      // outside the entry block or past the `newOpCount`-th operation in the
-      // entry block, we skip and leave exisintg `vector.insert_strided_slice`
-      // ops as is.
-      if (op->getBlock() != &entryBlock ||
-          static_cast<size_t>(std::distance(entryBlock.begin(),
-                                            op->getIterator())) >= newOpCount)
-        return;
+    // Replace all uses of placeholders for initially legal arguments with their
+    // original function arguments (that were added to `newFuncOp`).
+    for (auto &[placeholderOp, argIdx] : tmpOps) {
+      if (!placeholderOp)
+        continue;
+      Value replacement = newFuncOp.getArgument(argIdx);
+      rewriter.replaceAllUsesWith(placeholderOp->getResult(0), replacement);
+    }
 
+    // Replace dummy operands of new `vector.insert_strided_slice` ops with
+    // their corresponding new function arguments.
+    size_t unrolledInputIdx = 0;
+    for (auto [count, op] : enumerate(entryBlock.getOperations())) {
+      Operation &curOp = op;
+      // Since all newly created operations are in the beginning, reaching the
+      // end of them means that any later `vector.insert_strided_slice` should
+      // not be touched.
+      if (count >= newOpCount)
+        continue;
       if (auto vecOp = dyn_cast<vector::InsertStridedSliceOp>(op)) {
         size_t unrolledInputNo = unrolledInputNums[unrolledInputIdx];
-        rewriter.modifyOpInPlace(op, [&] {
-          op->setOperand(0, newFuncOp.getArgument(unrolledInputNo));
+        rewriter.modifyOpInPlace(&curOp, [&] {
+          curOp.setOperand(0, newFuncOp.getArgument(unrolledInputNo));
         });
         ++unrolledInputIdx;
       }
-    });
+    }
 
     // Erase the original funcOp. The `tmpOps` do not need to be erased since
     // they have no uses and will be handled by dead-code elimination.