[mlir][bufferize][NFC] Remove caller map and ordered func list from FuncAnalysisState

[springerm]

Summary

These can be local variables. No need to store them in the struct.

Depends On D123209

[springerm]

Landing this NFC change without further review as discussed in D119120.

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp

 #include "mlir/IR/Operation.h"
 
 using namespace mlir;
 using namespace linalg;
 using namespace tensor;
 using namespace comprehensive_bufferize;
 using namespace mlir::bufferization;
 
+/// A mapping of FuncOps to their callers.
+using FuncCallerMap = DenseMap<FuncOp, DenseSet<Operation *>>;
+
 namespace {
 /// The state of analysis of a FuncOp.
 enum class FuncOpAnalysisState { NotAnalyzed, InProgress, Analyzed };
 
 /// Extra analysis state that is required for bufferization of function
 /// boundaries.
 struct FuncAnalysisState : public DialectAnalysisState {
   // Note: Function arguments and/or function return values may disappear during
struct FuncAnalysisState : public DialectAnalysisState {
 
   /// A set of all written-to BlockArguments of FuncOps.
   DenseMap<FuncOp, BbArgIndexSet> writtenBbArgs;
 
   /// Keep track of which FuncOps are fully analyzed or currently being
   /// analyzed.
   DenseMap<FuncOp, FuncOpAnalysisState> analyzedFuncOps;
 
-  /// A list of functions in the order in which they are analyzed + bufferized.
-  SmallVector<FuncOp> orderedFuncOps;
-
-  /// A mapping of FuncOps to their callers.
-  DenseMap<FuncOp, DenseSet<Operation *>> callerMap;
-
   /// This function is called right before analyzing the given FuncOp. It
   /// initializes the data structures for the FuncOp in this state object.
   void startFunctionAnalysis(FuncOp funcOp) {
     analyzedFuncOps[funcOp] = FuncOpAnalysisState::InProgress;
     auto createdEquiv = equivalentFuncArgs.try_emplace(funcOp, IndexMapping());
     auto createdAliasingOperands =
         aliasingFuncArgs.try_emplace(funcOp, IndexToIndexListMapping());
     auto createdAliasingResults =

 /// Store all functions of the `moduleOp` in `orderedFuncOps`, sorted by
 /// callee-caller order (i.e. callees without callers first).
 /// Store the map of FuncOp to all its callers in `callerMap`.
 /// Return `failure()` if a cycle of calls is detected or if we are unable to
 /// retrieve the called FuncOp from any CallOpInterface.
 static LogicalResult
 getFuncOpsOrderedByCalls(ModuleOp moduleOp,
                          SmallVectorImpl<FuncOp> &orderedFuncOps,
-                         DenseMap<FuncOp, DenseSet<Operation *>> &callerMap) {
+                         FuncCallerMap &callerMap) {
   // For each FuncOp, the set of functions called by it (i.e. the union of
   // symbols of all nested CallOpInterfaceOp).
   DenseMap<FuncOp, DenseSet<FuncOp>> calledBy;
   // For each FuncOp, the number of CallOpInterface it contains.
   DenseMap<FuncOp, unsigned> numberCallOpsContainedInFuncOp;
   WalkResult res = moduleOp.walk([&](FuncOp funcOp) -> WalkResult {
     if (!funcOp.getBody().empty()) {
       func::ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
while (!numberCallOpsContainedInFuncOp.empty()) {
     orderedFuncOps.push_back(it->getFirst());
     for (auto callee : calledBy[it->getFirst()])
       numberCallOpsContainedInFuncOp[callee]--;
     numberCallOpsContainedInFuncOp.erase(it);
   }
   return success();
 }
 
-static void
-foreachCaller(const DenseMap<FuncOp, DenseSet<Operation *>> &callerMap,
-              FuncOp callee, llvm::function_ref<void(Operation *)> doit) {
+static void foreachCaller(const FuncCallerMap &callerMap, FuncOp callee,
+                          llvm::function_ref<void(Operation *)> doit) {
   auto itCallers = callerMap.find(callee);
   if (itCallers == callerMap.end())
     return;
   for (Operation *caller : itCallers->second)
     doit(caller);
 }
 
 /// Postprocess the linalg.buffer_layout annotation across function boundaries.

 LogicalResult mlir::linalg::comprehensive_bufferize::runModuleBufferize(
     ModuleOp moduleOp, OneShotBufferizationOptions options) {
   IRRewriter rewriter(moduleOp.getContext());
   OneShotAnalysisState analysisState(moduleOp, options);
   BufferizationState bufferizationState(analysisState);
   FuncAnalysisState &funcState = getFuncAnalysisState(analysisState);
   BufferizationAliasInfo &aliasInfo = analysisState.getAliasInfo();
 
-  if (failed(getFuncOpsOrderedByCalls(moduleOp, funcState.orderedFuncOps,
-                                      funcState.callerMap)))
+  // A list of functions in the order in which they are analyzed + bufferized.
+  SmallVector<FuncOp> orderedFuncOps;
+
+  // A mapping of FuncOps to their callers.
+  FuncCallerMap callerMap;
+
+  if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
     return failure();
 
   // Collect bbArg/return value information after the analysis.
   options.addPostAnalysisStep(aliasingFuncOpBBArgsAnalysis);
   options.addPostAnalysisStep(funcOpBbArgReadWriteAnalysis);
 
   // Analyze ops.
-  for (FuncOp funcOp : funcState.orderedFuncOps) {
+  for (FuncOp funcOp : orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.getBody().empty())
       continue;
 
     // Now analyzing function.
     funcState.startFunctionAnalysis(funcOp);
 
     // Gather equivalence info for CallOps.
for (FuncOp funcOp : orderedFuncOps) {
     if (options.testAnalysisOnly)
       annotateOpsWithBufferizationMarkers(funcOp, analysisState);
   }
 
   if (options.testAnalysisOnly)
     return success();
 
   // Bufferize functions.
-  for (FuncOp funcOp : funcState.orderedFuncOps) {
+  for (FuncOp funcOp : orderedFuncOps) {
     // No body => no analysis.
     if (!funcOp.getBody().empty())
       if (failed(bufferizeOp(funcOp, bufferizationState)))
         return failure();
 
     // Note: It would be good to apply cleanups here but we cannot as aliasInfo
     // would be invalidated.
     if (failed(bufferizeFuncOpBoundary(funcOp, rewriter, bufferizationState)))
       return failure();
   }
 
   // Check result.
-  for (FuncOp funcOp : funcState.orderedFuncOps) {
+  for (FuncOp funcOp : orderedFuncOps) {
     if (!options.allowReturnAllocs &&
         llvm::any_of(funcOp.getFunctionType().getResults(), [](Type t) {
           return t.isa<MemRefType, UnrankedMemRefType>();
         })) {
       funcOp->emitError("memref return type is unsupported");
       return failure();
     }
   }