From b2b63d1b9122befaf44bf91a3ed1e8a71c6ffe06 Mon Sep 17 00:00:00 2001
From: Yi Zhang <cathyzhyi@google.com>
Date: Wed, 22 Sep 2021 16:02:05 -0400
Subject: [PATCH] Reset operation when canceling root update transaction

Should reset the operation to original state when canceling the updates.

Reviewed By: rriddle, ftynse

Differential Revision: https://reviews.llvm.org/D110176
---
 mlir/lib/Transforms/Utils/DialectConversion.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 4f1c8cfa70c5..a748eb9164f3 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -1485,6 +1485,7 @@ void ConversionPatternRewriter::cancelRootUpdate(Operation *op) {
   auto &rootUpdates = impl->rootUpdates;
   auto it = llvm::find_if(llvm::reverse(rootUpdates), stateHasOp);
   assert(it != rootUpdates.rend() && "no root update started on op");
+  (*it).resetOperation();
   int updateIdx = std::prev(rootUpdates.rend()) - it;
   rootUpdates.erase(rootUpdates.begin() + updateIdx);
 }