[Mlir-commits] [mlir] [mlir][spirv] Fix FuncOpVectorUnroll to process placeholder values in all blocks (PR #142339)

[Darren Wihandi]

================
@@ -1020,35 +1020,39 @@ 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. We assume there is
-    // only one block for now.
+    // Replace the placeholder values with the new arguments.
     size_t unrolledInputIdx = 0;
-    for (auto [count, op] : enumerate(entryBlock.getOperations())) {
+    newFuncOp.walk([&](Operation *op) {
       // We first look for operands that are placeholders for initially legal
       // arguments.
-      Operation &curOp = op;
-      for (auto [operandIdx, operandVal] : llvm::enumerate(op.getOperands())) {
+      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(&curOp, [&curOp, &newFuncOp, it, idx] {
-            curOp.setOperand(idx, newFuncOp.getArgument(it->second));
----------------
fairywreath wrote:

This code originally only processes one block, but the placeholder values for the function arguments were applied with `replaceAllUsesWith` and may replace the function arguments across multiple blocks. The current implementation leaves the placeholder values in the other blocks (except for the first one) as is instead of the original function arguments, which is incorrect.

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