[MLIR] Add method to invalidate cached symbol table (#138014)
This PR adds a method to the `SymbolTableCollection` class to 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.
diff --git a/mlir/include/mlir/IR/SymbolTable.h b/mlir/include/mlir/IR/SymbolTable.h
index 597c6a9..a22d6f3 100644
--- a/mlir/include/mlir/IR/SymbolTable.h
+++ b/mlir/include/mlir/IR/SymbolTable.h
@@ -316,6 +316,13 @@
/// 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 76c9b7b..075a0ba 100644
--- a/mlir/lib/IR/SymbolTable.cpp
+++ b/mlir/lib/IR/SymbolTable.cpp
@@ -998,6 +998,10 @@
return *it.first->second;
}
+void SymbolTableCollection::invalidateSymbolTable(Operation *op) {
+ symbolTables.erase(op);
+}
+
//===----------------------------------------------------------------------===//
// LockedSymbolTableCollection
//===----------------------------------------------------------------------===//
