[mlir][vector] Drop incorrect startRootUpdate calls in vector distribution

[qedawkins]

Fixes asan failures in https://lab.llvm.org/buildbot/#/builders/5/builds/38191 introduced by #71964.

[llvmbot]

@llvm/pr-subscribers-mlir-vector

Author: Quinn Dawkins (qedawkins)

Changes

Fixes asan failures in https://lab.llvm.org/buildbot/#/builders/5/builds/38191 introduced by #71964.

---

Full diff: https://github.com/llvm/llvm-project/pull/71988.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp (-9)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
index ac73cf07004ded8..049ebcd748781d4 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
@@ -646,10 +646,6 @@ struct WarpOpElementwise : public OpRewritePattern<WarpExecuteOnLane0Op> {
     if (!yieldOperand)
       return failure();
 
-    // Notify the rewriter that the warp op is changing (see the comment on
-    // the WarpOpTransferRead pattern).
-    rewriter.startRootUpdate(warpOp);
-
     Operation *elementWise = yieldOperand->get().getDefiningOp();
     unsigned operandIndex = yieldOperand->getOperandNumber();
     Value distributedVal = warpOp.getResult(operandIndex);
@@ -688,7 +684,6 @@ struct WarpOpElementwise : public OpRewritePattern<WarpExecuteOnLane0Op> {
         {newWarpOp.getResult(operandIndex).getType()});
     rewriter.replaceAllUsesWith(newWarpOp.getResult(operandIndex),
                                 newOp->getResult(0));
-    rewriter.finalizeRootUpdate(warpOp);
     return success();
   }
 };
@@ -1058,9 +1053,6 @@ struct WarpOpBroadcast : public OpRewritePattern<WarpExecuteOnLane0Op> {
     if (vector::isBroadcastableTo(broadcastSrcType, destVecType) !=
         vector::BroadcastableToResult::Success)
       return failure();
-    // Notify the rewriter that the warp op is changing (see the comment on
-    // the WarpOpTransferRead pattern).
-    rewriter.startRootUpdate(warpOp);
     SmallVector<size_t> newRetIndices;
     WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(
         rewriter, warpOp, {broadcastSrc}, {broadcastSrcType}, newRetIndices);
@@ -1069,7 +1061,6 @@ struct WarpOpBroadcast : public OpRewritePattern<WarpExecuteOnLane0Op> {
         loc, destVecType, newWarpOp->getResult(newRetIndices[0]));
     rewriter.replaceAllUsesWith(newWarpOp->getResult(operandNumber),
                                 broadcasted);
-    rewriter.finalizeRootUpdate(warpOp);
     return success();
   }
 };

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Quinn Dawkins (qedawkins)

Changes

Fixes asan failures in https://lab.llvm.org/buildbot/#/builders/5/builds/38191 introduced by #71964.

---

Full diff: https://github.com/llvm/llvm-project/pull/71988.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp (-9)

diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
index ac73cf07004ded8..049ebcd748781d4 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
@@ -646,10 +646,6 @@ struct WarpOpElementwise : public OpRewritePattern<WarpExecuteOnLane0Op> {
     if (!yieldOperand)
       return failure();
 
-    // Notify the rewriter that the warp op is changing (see the comment on
-    // the WarpOpTransferRead pattern).
-    rewriter.startRootUpdate(warpOp);
-
     Operation *elementWise = yieldOperand->get().getDefiningOp();
     unsigned operandIndex = yieldOperand->getOperandNumber();
     Value distributedVal = warpOp.getResult(operandIndex);
@@ -688,7 +684,6 @@ struct WarpOpElementwise : public OpRewritePattern<WarpExecuteOnLane0Op> {
         {newWarpOp.getResult(operandIndex).getType()});
     rewriter.replaceAllUsesWith(newWarpOp.getResult(operandIndex),
                                 newOp->getResult(0));
-    rewriter.finalizeRootUpdate(warpOp);
     return success();
   }
 };
@@ -1058,9 +1053,6 @@ struct WarpOpBroadcast : public OpRewritePattern<WarpExecuteOnLane0Op> {
     if (vector::isBroadcastableTo(broadcastSrcType, destVecType) !=
         vector::BroadcastableToResult::Success)
       return failure();
-    // Notify the rewriter that the warp op is changing (see the comment on
-    // the WarpOpTransferRead pattern).
-    rewriter.startRootUpdate(warpOp);
     SmallVector<size_t> newRetIndices;
     WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(
         rewriter, warpOp, {broadcastSrc}, {broadcastSrcType}, newRetIndices);
@@ -1069,7 +1061,6 @@ struct WarpOpBroadcast : public OpRewritePattern<WarpExecuteOnLane0Op> {
         loc, destVecType, newWarpOp->getResult(newRetIndices[0]));
     rewriter.replaceAllUsesWith(newWarpOp->getResult(operandNumber),
                                 broadcasted);
-    rewriter.finalizeRootUpdate(warpOp);
     return success();
   }
 };