diff --git a/mlir/include/mlir/IR/PatternMatch.h b/mlir/include/mlir/IR/PatternMatch.h
--- a/mlir/include/mlir/IR/PatternMatch.h
+++ b/mlir/include/mlir/IR/PatternMatch.h
@@ -502,6 +502,11 @@
     finalizeRootUpdate(root);
   }
 
+  /// Find uses of `from` and replace it with `to`. It also marks every modified
+  /// uses and notifies the rewriter that an in-place operation modification is
+  /// about to happen.
+  void replaceAllUsesWith(Value from, Value to);
+
   /// Used to notify the rewriter that the IR failed to be rewritten because of
   /// a match failure, and provide a callback to populate a diagnostic with the
   /// reason why the failure occurred. This method allows for derived rewriters
diff --git a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
--- a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
+++ b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp
@@ -246,15 +246,9 @@
                                       sourceBlock.getOperations());
 
   // Step 5. RAUW thread indices to thread ops.
-  for (Value blockIdx : foreachThreadOp.getThreadIndices()) {
-    Value val = bvm.lookup(blockIdx);
-    SmallVector<OpOperand *> uses;
-    for (OpOperand &use : blockIdx.getUses())
-      uses.push_back(&use);
-    for (OpOperand *operand : uses) {
-      Operation *op = operand->getOwner();
-      rewriter.updateRootInPlace(op, [&]() { operand->set(val); });
-    }
+  for (Value loopIndex : foreachThreadOp.getThreadIndices()) {
+    Value blockIdx = bvm.lookup(loopIndex);
+    rewriter.replaceAllUsesWith(loopIndex, blockIdx);
   }
 
   // Step 6. Erase old op.
@@ -492,15 +486,9 @@
                                       sourceBlock.getOperations());
 
   // Step 6. RAUW thread indices to thread ops.
-  for (Value threadIdx : foreachThreadOp.getThreadIndices()) {
-    Value val = bvm.lookup(threadIdx);
-    SmallVector<OpOperand *> uses;
-    for (OpOperand &use : threadIdx.getUses())
-      uses.push_back(&use);
-    for (OpOperand *operand : uses) {
-      Operation *op = operand->getOwner();
-      rewriter.updateRootInPlace(op, [&]() { operand->set(val); });
-    }
+  for (Value loopIndex : foreachThreadOp.getThreadIndices()) {
+    Value threadIdx = bvm.lookup(loopIndex);
+    rewriter.replaceAllUsesWith(loopIndex, threadIdx);
   }
 
   // Step 7. syncthreads.
diff --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -309,6 +309,14 @@
   source->erase();
 }
 
+/// Find uses of `from` and replace it with `to`
+void RewriterBase::replaceAllUsesWith(Value from, Value to) {
+  for (OpOperand &operand : llvm::make_early_inc_range(from.getUses())) {
+    Operation *op = operand.getOwner();
+    updateRootInPlace(op, [&]() { operand.set(to); });
+  }
+}
+
 // Merge the operations of block 'source' before the operation 'op'. Source
 // block should not have existing predecessors or successors.
 void RewriterBase::mergeBlockBefore(Block *source, Operation *op,