Add convenience methods to set an OpBuilder insertion point after an Operation (NFC)

PiperOrigin-RevId: 270727180

mlir/examples/toy/Ch5/mlir/ShapeInferencePass.cpp

@@ -320,8 +320,7 @@ public:
         // Found a specialized callee! Let's turn this into a normal call
         // operation.
         SmallVector<mlir::Value *, 8> operands(op->getOperands());
-        mlir::OpBuilder builder(f.getBody());
-        builder.setInsertionPoint(op);
+        mlir::OpBuilder builder(op);
         auto newCall =
             builder.create<mlir::CallOp>(op->getLoc(), mangledCallee, operands);
         if (newCall.getNumResults()) {

mlir/include/mlir/IR/Builders.h

@@ -280,6 +280,12 @@ public:
     setInsertionPoint(op->getBlock(), Block::iterator(op));
   }
 
+  /// Sets the insertion point to the node after the specified operation, which
+  /// will cause subsequent insertions to go right after it.
+  void setInsertionPointAfter(Operation *op) {
+    setInsertionPoint(op->getBlock(), ++Block::iterator(op));
+  }
+
   /// Sets the insertion point to the start of the specified block.
   void setInsertionPointToStart(Block *block) {
     setInsertionPoint(block, block->begin());

mlir/lib/Conversion/GPUToCUDA/GenerateCubinAccessors.cpp

@@ -80,7 +80,7 @@ private:
 
     // Insert the getter function just after the original function.
     OpBuilder moduleBuilder(module.getBody(), module.getBody()->begin());
-    moduleBuilder.setInsertionPoint(stubFunc.getOperation()->getNextNode());
+    moduleBuilder.setInsertionPointAfter(stubFunc.getOperation());
     auto getterType = moduleBuilder.getFunctionType(
         llvm::None, LLVM::LLVMType::getInt8PtrTy(llvmDialect));
     nameBuffer.append(kCubinGetterSuffix);

mlir/lib/Transforms/PipelineDataTransfer.cpp

@@ -73,7 +73,6 @@ static unsigned getTagMemRefPos(Operation &dmaInst) {
 static bool doubleBuffer(Value *oldMemRef, AffineForOp forOp) {
   auto *forBody = forOp.getBody();
   OpBuilder bInner(forBody, forBody->begin());
-  bInner.setInsertionPoint(forBody, forBody->begin());
 
   // Doubles the shape with a leading dimension extent of 2.
   auto doubleShape = [&](MemRefType oldMemRefType) -> MemRefType {
@@ -130,8 +129,7 @@ static bool doubleBuffer(Value *oldMemRef, AffineForOp forOp) {
     return false;
   }
   // Insert the dealloc op right after the for loop.
-  bOuter.setInsertionPoint(forInst->getBlock(),
-                           std::next(Block::iterator(forInst)));
+  bOuter.setInsertionPointAfter(forInst);
   bOuter.create<DeallocOp>(forInst->getLoc(), newMemRef);
 
   return true;