Merge pull request #37687 from xinan-jiang:pr/mlir/fix-dead-temp-buff…

…er-removal PiperOrigin-RevId: 302504799 Change-Id: I508140041e5abff92f433460e22e8af770b1c963
tensorflow · Mar 23, 2020 · 4286b65 · 4286b65
2 parents 25ae1e1 + 9a7ad6f
commit 4286b65
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc b/tensorflow/compiler/xla/service/mlir_gpu/kernel_lowering.cc
@@ -243,14 +243,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);
+    }
   }
 };