[TransformInterpreter] Avoid modifying a region while walking it (#11788)
The current code results in a use-after-free. Fixes
https://github.com/iree-org/iree/pull/11788
diff --git a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgTransform/Passes/TransformInterpreter.cpp b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgTransform/Passes/TransformInterpreter.cpp
index ee8075c..476e914 100644
--- a/llvm-external-projects/iree-dialects/lib/Dialect/LinalgTransform/Passes/TransformInterpreter.cpp
+++ b/llvm-external-projects/iree-dialects/lib/Dialect/LinalgTransform/Passes/TransformInterpreter.cpp
@@ -242,23 +242,31 @@
}
void runOnOperation() override {
+ SmallVector<Operation *> toDelete;
getOperation()->walk<WalkOrder::PreOrder>([&](Operation *nestedOp) {
- if (isa<iree_compiler::IREE::LinalgExt::DoNotDCEOperandsOp>(nestedOp))
- nestedOp->erase();
- if (isa<::mlir::transform::TransformOpInterface>(nestedOp)) {
- nestedOp->erase();
+ if (isa<iree_compiler::IREE::LinalgExt::DoNotDCEOperandsOp>(nestedOp)) {
+ toDelete.push_back(nestedOp);
+ } else if (isa<::mlir::transform::TransformOpInterface>(nestedOp)) {
+ toDelete.push_back(nestedOp);
return WalkResult::skip();
}
return WalkResult::advance();
});
+ for (auto op : toDelete) {
+ op->erase();
+ }
+ SmallVector<ModuleOp> modulesToDelete;
// Remove potential empty module after cleanup.
getOperation()->walk([&](ModuleOp module) {
if (module.getBodyRegion().hasOneBlock() && module.getBody()->empty()) {
- module->erase();
+ modulesToDelete.push_back(module);
return WalkResult::skip();
}
return WalkResult::advance();
});
+ for (auto module : modulesToDelete) {
+ module->erase();
+ }
}
};
} // namespace
