[mlir][Vector]Fix bug where vector::WarpExecuteOnLane0Op are created with 2 blocks in the region

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

rewriteWarpOpToScfFor(RewriterBase &rewriter, WarpExecuteOnLane0Op warpOp,
   // Store vectors that are defined outside of warpOp into the scratch pad
   // buffer.
   SmallVector<Value> bbArgReplacements;
   for (const auto &it : llvm::enumerate(warpOp.getArgs())) {
     Value val = it.value();
     Value bbArg = warpOpBody->getArgument(it.index());
 
     rewriter.setInsertionPoint(ifOp);
-    Value buffer = options.warpAllocationFn(warpOp->getLoc(), rewriter, warpOp,
-                                            bbArg.getType());
+    Value buffer =
+        options.warpAllocationFn(loc, rewriter, warpOp, bbArg.getType());
 
     // Store arg vector into buffer.
     rewriter.setInsertionPoint(ifOp);
     auto vectorType = val.getType().cast<VectorType>();
     int64_t storeSize = vectorType.getShape()[0];
     Value storeOffset = rewriter.create<arith::MulIOp>(
         loc, warpOp.getLaneid(),
         rewriter.create<arith::ConstantIndexOp>(loc, storeSize));
     rewriter.create<vector::StoreOp>(loc, val, buffer, storeOffset);
 
     // Load bbArg vector from buffer.
     rewriter.setInsertionPointToStart(ifOp.thenBlock());
     auto bbArgType = bbArg.getType().cast<VectorType>();
     Value loadOp = rewriter.create<vector::LoadOp>(loc, bbArgType, buffer, c0);
     bbArgReplacements.push_back(loadOp);
   }
 
   // Insert sync after all the stores and before all the loads.
   if (!warpOp.getArgs().empty()) {
     rewriter.setInsertionPoint(ifOp);
-    options.warpSyncronizationFn(warpOp->getLoc(), rewriter, warpOp);
+    options.warpSyncronizationFn(loc, rewriter, warpOp);
   }
 
   // Move body of warpOp to ifOp.
   rewriter.mergeBlocks(warpOpBody, ifOp.thenBlock(), bbArgReplacements);
 
   // Rewrite terminator and compute replacements of WarpOp results.
   SmallVector<Value> replacements;
   auto yieldOp = cast<vector::YieldOp>(ifOp.thenBlock()->getTerminator());
   Location yieldLoc = yieldOp.getLoc();
   for (const auto &it : llvm::enumerate(yieldOp.operands())) {
     Value val = it.value();
     Type resultType = warpOp->getResultTypes()[it.index()];
     rewriter.setInsertionPoint(ifOp);
-    Value buffer = options.warpAllocationFn(warpOp->getLoc(), rewriter, warpOp,
-                                            val.getType());
+    Value buffer =
+        options.warpAllocationFn(loc, rewriter, warpOp, val.getType());
 
     // Store yielded value into buffer.
     rewriter.setInsertionPoint(yieldOp);
     if (val.getType().isa<VectorType>())
       rewriter.create<vector::StoreOp>(yieldLoc, val, buffer, c0);
     else
       rewriter.create<memref::StoreOp>(yieldLoc, val, buffer, c0);
 
if (resultType == val.getType()) {
                                                      buffer, loadOffset);
       replacements.push_back(loadOp);
     }
   }
 
   // Insert sync after all the stores and before all the loads.
   if (!yieldOp.operands().empty()) {
     rewriter.setInsertionPointAfter(ifOp);
-    options.warpSyncronizationFn(warpOp->getLoc(), rewriter, warpOp);
+    options.warpSyncronizationFn(loc, rewriter, warpOp);
   }
 
   // Delete terminator and add empty scf.yield.
   rewriter.eraseOp(yieldOp);
   rewriter.setInsertionPointToEnd(ifOp.thenBlock());
   rewriter.create<scf::YieldOp>(yieldLoc);
 
   // Compute replacements for WarpOp results.
static WarpExecuteOnLane0Op moveRegionToNewWarpOpAndReplaceReturns(
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(warpOp);
   auto newWarpOp = rewriter.create<WarpExecuteOnLane0Op>(
       warpOp.getLoc(), newReturnTypes, warpOp.getLaneid(), warpOp.getWarpSize(),
       warpOp.getArgs(), warpOp.getBody()->getArgumentTypes());
 
   Region &opBody = warpOp.getBodyRegion();
   Region &newOpBody = newWarpOp.getBodyRegion();
+  Block &newOpFirstBlock = newOpBody.front();
   rewriter.inlineRegionBefore(opBody, newOpBody, newOpBody.begin());
+  rewriter.eraseBlock(&newOpFirstBlock);
+  assert(newWarpOp.getWarpRegion().hasOneBlock() &&
+         "expected WarpOp with single block");
+
   auto yield =
       cast<vector::YieldOp>(newOpBody.getBlocks().begin()->getTerminator());
 
   rewriter.updateRootInPlace(
       yield, [&]() { yield.operandsMutable().assign(newYieldedValues); });
   return newWarpOp;
 }