Address review comment from PR3616 (#3627)

iree/compiler/Conversion/CodegenUtils/ForOpCanonicalization.cpp

diff --git a/iree/compiler/Conversion/CodegenUtils/ForOpCanonicalization.cpp b/iree/compiler/Conversion/CodegenUtils/ForOpCanonicalization.cpp
index c7f2134..e0e96f5 100644
--- a/iree/compiler/Conversion/CodegenUtils/ForOpCanonicalization.cpp
+++ b/iree/compiler/Conversion/CodegenUtils/ForOpCanonicalization.cpp
@@ -80,8 +80,7 @@
     rewriter.mergeBlocks(source, dest, dest->getArguments());
     // Replace the yield op by one that returns only the used values.
     auto yieldOp = cast<scf::YieldOp>(dest->getTerminator());
-    rewriter.updateRootInPlace(
-        yieldOp, [&]() { yieldOp.getOperation()->setOperands(results); });
+    yieldOp.getOperation()->setOperands(results);
   }
 
   LogicalResult matchAndRewrite(scf::ForOp forOp,
@@ -103,7 +102,7 @@
       returnValues[it.index()] = newReturn;
 
       BlockAndValueMapping mapping;
-      mapping.map(op->getOperand(0), initArgs[it.index()]);
+      mapping.map(it.value(), initArgs[it.index()]);
       initArgs[it.index()] = rewriter.clone(*op, mapping)->getResult(0);
     }
     if (iteratorFolded.empty()) return success();