From 9a7ad6fe1231291fb139f353b7a596a431c1b1b9 Mon Sep 17 00:00:00 2001
From: Xinan Jiang <xinan.jxn@gmail.com>
Date: Wed, 18 Mar 2020 18:35:17 +0800
Subject: [PATCH] [MLIR][XLA] Fix ops erase bug in DeadTempBufferRemoval

---
 .../compiler/xla/service/mlir_gpu/kernel_lowering.cc     | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
index 151d82fd2a1bdf..45a9470808b5bc 100644
--- a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
+++ b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
@@ -244,14 +244,19 @@ struct DeadTempBufferRemoval : mlir::FunctionPass<DeadTempBufferRemoval> {
   }
 
   void runOnFunction() override {
+    llvm::SmallVector<mlir::Operation *, 8> opsToErase;
     getFunction().walk([&](mlir::AllocOp allocOp) {
       if (!operationConsideredDead(allocOp)) {
         return;
       }
 
-      // TODO(herhut): There should be a generic helper for this.
-      recursiveErase(allocOp);
+      opsToErase.push_back(allocOp);
     });
+
+    for (auto *op : opsToErase) {
+      // TODO(herhut): There should be a generic helper for this.
+      recursiveErase(op);
+    }
   }
 };