diff --git a/mlir/include/mlir/IR/SymbolTable.h b/mlir/include/mlir/IR/SymbolTable.h index 597c6a9a1d891..a22d6f3a0426b 100644 --- a/mlir/include/mlir/IR/SymbolTable.h +++ b/mlir/include/mlir/IR/SymbolTable.h @@ -316,6 +316,13 @@ class SymbolTableCollection { /// Lookup, or create, a symbol table for an operation. SymbolTable &getSymbolTable(Operation *op); + /// Invalidate the cached symbol table for an operation. + /// This is important when doing IR modifications that erase and also create + /// operations having the 'OpTrait::SymbolTable' trait. If a symbol table of + /// an erased operation is not invalidated, a new operation sharing the same + /// address would be associated with outdated, and wrong, information. + void invalidateSymbolTable(Operation *op); + private: friend class LockedSymbolTableCollection; diff --git a/mlir/lib/IR/SymbolTable.cpp b/mlir/lib/IR/SymbolTable.cpp index 76c9b7b1e8afa..075a0ba15d7cd 100644 --- a/mlir/lib/IR/SymbolTable.cpp +++ b/mlir/lib/IR/SymbolTable.cpp @@ -998,6 +998,10 @@ SymbolTable &SymbolTableCollection::getSymbolTable(Operation *op) { return *it.first->second; } +void SymbolTableCollection::invalidateSymbolTable(Operation *op) { + symbolTables.erase(op); +} + //===----------------------------------------------------------------------===// // LockedSymbolTableCollection //===----------------------------------------------------------------------===//