From 9d80f1742294aeedfbe547349a5422c363b5fb1a Mon Sep 17 00:00:00 2001 From: Matthias Springer Date: Mon, 27 Mar 2023 08:57:46 +0200 Subject: [PATCH] [mlir][bufferize] Fix tracking of erased ops This is a workaround until D144193 has landed. Differential Revision: https://reviews.llvm.org/D146868 --- mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp index 0b10baf..e5e125f 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp @@ -358,9 +358,12 @@ public: protected: void notifyOperationRemoved(Operation *op) override { - erasedOps.insert(op); - // Erase if present. - toMemrefOps.erase(op); + // TODO: Walk can be removed when D144193 has landed. + op->walk([&](Operation *op) { + erasedOps.insert(op); + // Erase if present. + toMemrefOps.erase(op); + }); } void notifyOperationInserted(Operation *op) override { -- 2.7.4