[Mlir-commits] [mlir] [MLIR] Cache symbol tables during OneShotBufferization analyses (PR #138125)

[Matthias Springer]

================
@@ -76,23 +76,34 @@ getBufferizedFunctionArgType(FuncOp funcOp, int64_t index,
 }
 
 /// Return the FuncOp called by `callOp`.
-static FuncOp getCalledFunction(CallOpInterface callOp) {
+static FuncOp getCalledFunction(CallOpInterface callOp,
+                                mlir::SymbolTableCollection &symbolTable) {
   SymbolRefAttr sym =
       llvm::dyn_cast_if_present<SymbolRefAttr>(callOp.getCallableForCallee());
   if (!sym)
     return nullptr;
   return dyn_cast_or_null<FuncOp>(
-      SymbolTable::lookupNearestSymbolFrom(callOp, sym));
+      symbolTable.lookupNearestSymbolFrom(callOp, sym));
 }
 
-/// Get FuncAnalysisState.
+/// Get or create FuncAnalysisState.
 static const FuncAnalysisState &
-getFuncAnalysisState(const AnalysisState &state) {
+getOrCreateFuncAnalysisState(const AnalysisState &state) {
   assert(isa<OneShotAnalysisState>(state) && "expected OneShotAnalysisState");
-  auto *result = static_cast<const OneShotAnalysisState &>(state)
-                     .getExtension<FuncAnalysisState>();
-  assert(result && "FuncAnalysisState does not exist");
-  return *result;
+
+  // Unfortunately, at the moment the BufferizableOpInterface methods do provide
----------------
matthias-springer wrote:

This looks a bit hacky. And I think it's not really necessary. When running through `analyzeModuleOp` (`OneShotModuleBufferize.cpp`), there will always be a `FuncAnalysisState`.

When running through the normal One-Shot Bufferize (`bufferize-function-boundaries=false`), there is no `FuncAnalysisState`. (But this entry point should only be used for functions, not modules. And I think it will ignore function calls anyway.) Implementations like `getFuncOpAnalysisState` fall back to the conservative path. We should do the same for `getCalledFunction`: If there is no `FuncAnalysisState`, use the previous "slow" lookup instead of the symbol table.


https://github.com/llvm/llvm-project/pull/138125