[mlir] Introduce `replaceUsesOfWith` to `RewriterBase`

[guraypp]

Summary

Finding uses of a value and replacing them with a new one is a common method. I have not seen an safe and easy shortcut that does that. This revision attempts to address that by intoroducing replaceUsesOfWith to RewriterBase.

[rriddle]

mlir/lib/IR/PatternMatch.cpp

314–316

Why is this collecting all of the uses into a vector?

[guraypp]

mlir/lib/IR/PatternMatch.cpp

314–316

My motivation of this change comes from D138029. Shortly, the old code looked like this. getUsers iterates on actual users in this situation. It is incorrect because I am modifying uses while iterating. Ofc, I am happy to change if there is better way.

for (Operation *user : llvm::make_early_inc_range(from.getUsers())) 
  rewriter.updateRootInPlace(user, [&]() {
    user->replaceUsesOfWith(from, to);
});

[rriddle]

mlir/lib/IR/PatternMatch.cpp

314–316

Have you tried an early inc range of getUses?

[guraypp]

get rid of vector and use make_early_inc_range

[mehdi_amini]

mlir/include/mlir/IR/PatternMatch.h

505

Note that using this instead of from.replaceUsesOfWith(to) will mark every modified users and notify the listener.

[mehdi_amini]

mlir/include/mlir/IR/PatternMatch.h

505

Because otherwise, one may wonder "why reimplement it here?"

506

Please rename it to replaceAllUsesWith for consitency with Value::replaceAllUsesWith

[guraypp]

Address @mehdi_amini comments

[rriddle]

Nice!

mlir/include/mlir/IR/PatternMatch.h

public:
   /// callable.
   template <typename CallableT>
   void updateRootInPlace(Operation *root, CallableT &&callable) {
     startRootUpdate(root);
     callable();
     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
   /// to optionally hook into the reason why a rewrite failed, and display it to
   /// users.
   template <typename CallbackT>
   std::enable_if_t<!std::is_convertible<CallbackT, Twine>::value, LogicalResult>
   notifyMatchFailure(Location loc, CallbackT &&reasonCallback) {

mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp

DiagnosedSilenceableFailure mlir::transform::gpu::mapForeachToBlocksImpl(
   rewriter.eraseOp(foreachThreadOp.getTerminator());
   Block *targetBlock = foreachThreadOp->getBlock();
   Block::iterator insertionPoint = Block::iterator(foreachThreadOp);
   Block &sourceBlock = foreachThreadOp.getRegion().front();
   targetBlock->getOperations().splice(insertionPoint,
                                       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.
   rewriter.eraseOp(foreachThreadOp);
 
   return DiagnosedSilenceableFailure::success();
 }
 
if (predicate) {
     targetBlock = foreachThreadOp->getBlock();
     insertionPoint = Block::iterator(foreachThreadOp);
   }
   Block &sourceBlock = foreachThreadOp.getRegion().front();
   targetBlock->getOperations().splice(insertionPoint,
                                       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.
   // TODO: Need warpsync
   if (syncAfterDistribute)
     rewriter.create<BarrierOp>(loc);
 
   // Step 8. Erase old op.

mlir/lib/IR/PatternMatch.cpp

void RewriterBase::mergeBlocks(Block *source, Block *dest,
 
   // Splice the operations of the 'source' block into the 'dest' block and erase
   // it.
   dest->getOperations().splice(dest->end(), source->getOperations());
   source->dropAllUses();
   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,
                                     ValueRange argValues) {
   assert(source->hasNoPredecessors() &&
          "expected 'source' to have no predecessors");
   assert(source->hasNoSuccessors() &&
          "expected 'source' to have no successors");