[mlir][bufferize][NFC] Rename ModuleAnalysisState to FuncAnalysisState

[springerm]

Summary

This is for consistency reasons. *AnalysisState always starts with the name of the dialect.

Depends On D119120

[springerm]

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

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

 // implementations for FuncOp, CallOp and ReturnOp.
 //
 // Module Bufferization is run via `runModuleBufferize(ModuleOp, ...)`. This
 // function analyzes the given module and determines the order of analysis and
 // bufferization: Functions that are called are processed before their
 // respective callers.
 //
 // After analyzing a FuncOp, additional information about its bbArgs is
-// gathered through PostAnalysisStepFns and stored in
-// `ModuleAnalysisState`.
+// gathered through PostAnalysisStepFns and stored in `FuncAnalysisState`.
 //
 // * `aliasingFuncOpBBArgsAnalysis` determines the equivalent/aliasing bbArgs
 // for
 //   each tensor return value (if any).
 // * `funcOpBbArgReadWriteAnalysis` determines whether or not a tensor bbArg is
 //   read/written.
 //
 // Only tensors that are equivalent to some FuncOp bbArg may be returned.

 using namespace mlir::bufferization;
 
 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 ModuleAnalysisState : public DialectAnalysisState {
+struct FuncAnalysisState : public DialectAnalysisState {
   // Note: Function arguments and/or function return values may disappear during
   // bufferization. Functions and their CallOps are analyzed and bufferized
   // separately. To ensure that a CallOp analysis/bufferization can access an
   // already bufferized function's analysis results, we store bbArg/return value
   // indices instead of BlockArguments/OpOperand pointers.
 
   /// A set of block argument indices.
   using BbArgIndexSet = DenseSet<int64_t>;
#ifndef NDEBUG
     assert(createdAliasingResults.second && "aliasing info exists already");
     assert(createdRead.second && "bbarg access info exists already");
     assert(createdWritten.second && "bbarg access info exists already");
 #endif // NDEBUG
   }
 };
 } // namespace
 
-/// Get ModuleAnalysisState.
-static const ModuleAnalysisState &
-getModuleAnalysisState(const AnalysisState &state) {
-  Optional<const ModuleAnalysisState *> maybeState =
-      state.getDialectState<ModuleAnalysisState>(
+/// Get FuncAnalysisState.
+static const FuncAnalysisState &
+getFuncAnalysisState(const AnalysisState &state) {
+  Optional<const FuncAnalysisState *> maybeState =
+      state.getDialectState<FuncAnalysisState>(
           func::FuncDialect::getDialectNamespace());
-  assert(maybeState.hasValue() && "ModuleAnalysisState does not exist");
+  assert(maybeState.hasValue() && "FuncAnalysisState does not exist");
   return **maybeState;
 }
 
-/// Get or create ModuleAnalysisState.
-static ModuleAnalysisState &getModuleAnalysisState(AnalysisState &state) {
-  return state.getOrCreateDialectState<ModuleAnalysisState>(
+/// Get or create FuncAnalysisState.
+static FuncAnalysisState &getFuncAnalysisState(AnalysisState &state) {
+  return state.getOrCreateDialectState<FuncAnalysisState>(
       func::FuncDialect::getDialectNamespace());
 }
 
 /// Return the state (phase) of analysis of the FuncOp.
 static FuncOpAnalysisState getFuncOpAnalysisState(const AnalysisState &state,
                                                   FuncOp funcOp) {
-  const ModuleAnalysisState &moduleState = getModuleAnalysisState(state);
+  const FuncAnalysisState &moduleState = getFuncAnalysisState(state);
   auto it = moduleState.analyzedFuncOps.find(funcOp);
   if (it == moduleState.analyzedFuncOps.end())
     return FuncOpAnalysisState::NotAnalyzed;
   return it->second;
 }
 
 /// Return the unique ReturnOp that terminates `funcOp`.
 /// Return nullptr if there is no such unique ReturnOp.
static void annotateEquivalentReturnBbArg(OpOperand &returnVal,
   }
   equivBbArgs[returnVal.getOperandNumber()] = bbArg.getArgNumber();
 
   OpBuilder b(op->getContext());
   op->setAttr(kEquivalentArgsAttr, b.getI64ArrayAttr(equivBbArgs));
 }
 
 /// Store function BlockArguments that are equivalent to/aliasing a returned
-/// value in ModuleAnalysisState.
+/// value in FuncAnalysisState.
 static LogicalResult
 aliasingFuncOpBBArgsAnalysis(Operation *op, AnalysisState &state,
                              BufferizationAliasInfo &aliasInfo,
                              SmallVector<Operation *> &newOps) {
-  ModuleAnalysisState &moduleState = getModuleAnalysisState(state);
+  FuncAnalysisState &funcState = getFuncAnalysisState(state);
 
   // Support only single return-terminated block in the function.
   auto funcOp = cast<FuncOp>(op);
   func::ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
   assert(returnOp && "expected func with single return op");
 
   for (OpOperand &returnVal : returnOp->getOpOperands())
     if (returnVal.get().getType().isa<RankedTensorType>())
       for (BlockArgument bbArg : funcOp.getArguments())
         if (bbArg.getType().isa<RankedTensorType>()) {
           int64_t returnIdx = returnVal.getOperandNumber();
           int64_t bbArgIdx = bbArg.getArgNumber();
           if (aliasInfo.areEquivalentBufferizedValues(returnVal.get(), bbArg)) {
-            moduleState.equivalentFuncArgs[funcOp][returnIdx] = bbArgIdx;
+            funcState.equivalentFuncArgs[funcOp][returnIdx] = bbArgIdx;
             if (state.getOptions().testAnalysisOnly)
               annotateEquivalentReturnBbArg(returnVal, bbArg);
           }
           if (aliasInfo.areAliasingBufferizedValues(returnVal.get(), bbArg)) {
-            moduleState.aliasingFuncArgs[funcOp][returnIdx].push_back(bbArgIdx);
-            moduleState.aliasingReturnVals[funcOp][bbArgIdx].push_back(
-                returnIdx);
+            funcState.aliasingFuncArgs[funcOp][returnIdx].push_back(bbArgIdx);
+            funcState.aliasingReturnVals[funcOp][bbArgIdx].push_back(returnIdx);
           }
         }
 
   return success();
 }
 
 /// Return true if the buffer of the given tensor value is written to. Must not
 /// be called for values inside not yet analyzed functions. (Post-analysis

 
 /// Determine which FuncOp bbArgs are read and which are written. If this
 /// PostAnalysisStepFn is run on a function with unknown ops, it will
 /// conservatively assume that such ops bufferize to a read + write.
 static LogicalResult
 funcOpBbArgReadWriteAnalysis(Operation *op, AnalysisState &state,
                              BufferizationAliasInfo &aliasInfo,
                              SmallVector<Operation *> &newOps) {
-  ModuleAnalysisState &moduleState = getModuleAnalysisState(state);
+  FuncAnalysisState &funcState = getFuncAnalysisState(state);
   auto funcOp = cast<FuncOp>(op);
 
   // If the function has no body, conservatively assume that all args are
   // read + written.
   if (funcOp.getBody().empty()) {
     for (BlockArgument bbArg : funcOp.getArguments()) {
-      moduleState.readBbArgs[funcOp].insert(bbArg.getArgNumber());
-      moduleState.writtenBbArgs[funcOp].insert(bbArg.getArgNumber());
+      funcState.readBbArgs[funcOp].insert(bbArg.getArgNumber());
+      funcState.writtenBbArgs[funcOp].insert(bbArg.getArgNumber());
     }
 
     return success();
   }
 
   for (BlockArgument bbArg : funcOp.getArguments()) {
     if (!bbArg.getType().isa<TensorType>())
       continue;
     bool isRead = state.isValueRead(bbArg);
     bool isWritten = isValueWritten(bbArg, state, aliasInfo);
     if (state.getOptions().testAnalysisOnly)
       annotateFuncArgAccess(funcOp, bbArg, isRead, isWritten);
     if (isRead)
-      moduleState.readBbArgs[funcOp].insert(bbArg.getArgNumber());
+      funcState.readBbArgs[funcOp].insert(bbArg.getArgNumber());
     if (isWritten)
-      moduleState.writtenBbArgs[funcOp].insert(bbArg.getArgNumber());
+      funcState.writtenBbArgs[funcOp].insert(bbArg.getArgNumber());
   }
 
   return success();
 }
 } // namespace
 
 static bool isaTensor(Type t) { return t.isa<TensorType>(); }
 

 }
 
 /// Gather equivalence info of CallOps.
 /// Note: This only adds new equivalence info if the called function was already
 /// analyzed.
 // TODO: This does not handle cyclic function call graphs etc.
 static void equivalenceAnalysis(FuncOp funcOp,
                                 BufferizationAliasInfo &aliasInfo,
-                                ModuleAnalysisState &moduleState) {
+                                FuncAnalysisState &funcState) {
   funcOp->walk([&](func::CallOp callOp) {
     FuncOp calledFunction = getCalledFunction(callOp);
     assert(calledFunction && "could not retrieved called FuncOp");
 
     // No equivalence info available for the called function.
-    if (!moduleState.equivalentFuncArgs.count(calledFunction))
+    if (!funcState.equivalentFuncArgs.count(calledFunction))
       return WalkResult::skip();
 
-    for (auto it : moduleState.equivalentFuncArgs[calledFunction]) {
+    for (auto it : funcState.equivalentFuncArgs[calledFunction]) {
       int64_t returnIdx = it.first;
       int64_t bbargIdx = it.second;
       Value returnVal = callOp.getResult(returnIdx);
       Value argVal = callOp->getOperand(bbargIdx);
       aliasInfo.unionEquivalenceClasses(returnVal, argVal);
     }
 
     return WalkResult::advance();

 /// future.
 ///
 /// Note: Returning a memref currently fails bufferization. If such memrefs
 /// originate from an op with an Alloc effect, they could be hoisted in the
 /// future.
 static LogicalResult bufferizeFuncOpBoundary(FuncOp funcOp,
                                              RewriterBase &rewriter,
                                              BufferizationState &state) {
-  const ModuleAnalysisState &moduleState =
-      getModuleAnalysisState(state.getAnalysisState());
+  const FuncAnalysisState &funcState =
+      getFuncAnalysisState(state.getAnalysisState());
 
   // If nothing to do then we are done.
   if (!llvm::any_of(funcOp.getFunctionType().getInputs(), isaTensor) &&
       !llvm::any_of(funcOp.getFunctionType().getResults(), isaTensor))
     return success();
 
   // Get the bufferized FunctionType for funcOp or construct it if not yet
   // available.
for (OpOperand &returnOperand : returnOp->getOpOperands()) {
 
     // If not a renturn tensor type just forward it.
     if (!returnVal.getType().isa<RankedTensorType>()) {
       returnValues.push_back(returnVal);
       continue;
     }
 
     // If return operand is equivalent to some bbArg, no need to return it.
-    auto funcOpIt = moduleState.equivalentFuncArgs.find(funcOp);
-    if (funcOpIt != moduleState.equivalentFuncArgs.end() &&
+    auto funcOpIt = funcState.equivalentFuncArgs.find(funcOp);
+    if (funcOpIt != funcState.equivalentFuncArgs.end() &&
         funcOpIt->second.count(returnOperand.getOperandNumber()))
       continue;
 
     // Cast values at the call site if necessary.
     returnValues.push_back(
         getNonCastedValue(*state.getBuffer(rewriter, returnOperand)));
   }
 

 
 namespace mlir {
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace std_ext {
 
 /// Return the index of the bbArg in the given FuncOp that is equivalent to the
 /// specified return value (if any).
-static Optional<int64_t>
-getEquivalentFuncArgIdx(FuncOp funcOp, const ModuleAnalysisState &state,
+static Optional<int64_t> getEquivalentFuncArgIdx(FuncOp funcOp,
+                                                 const FuncAnalysisState &state,
                                                  int64_t returnValIdx) {
   auto funcOpIt = state.equivalentFuncArgs.find(funcOp);
   if (funcOpIt == state.equivalentFuncArgs.end())
     // No equivalence info stores for funcOp.
     return None;
 
   auto retValIt = funcOpIt->getSecond().find(returnValIdx);
   if (retValIt == funcOpIt->getSecond().end())
     // Return value has no equivalent bbArg.
     return None;
 
   return retValIt->getSecond();
 }
 
 struct CallOpInterface
     : public BufferizableOpInterface::ExternalModel<CallOpInterface,
                                                     func::CallOp> {
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
     func::CallOp callOp = cast<func::CallOp>(op);
     FuncOp funcOp = getCalledFunction(callOp);
     assert(funcOp && "expected CallOp to a FuncOp");
 
-    const ModuleAnalysisState &moduleState = getModuleAnalysisState(state);
+    const FuncAnalysisState &funcState = getFuncAnalysisState(state);
     if (getFuncOpAnalysisState(state, funcOp) != FuncOpAnalysisState::Analyzed)
       // FuncOp not analyzed yet. Assume that OpOperand is read.
       return true;
 
-    return moduleState.readBbArgs.lookup(funcOp).contains(
+    return funcState.readBbArgs.lookup(funcOp).contains(
         opOperand.getOperandNumber());
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
     func::CallOp callOp = cast<func::CallOp>(op);
     FuncOp funcOp = getCalledFunction(callOp);
     assert(funcOp && "expected CallOp to a FuncOp");
 
-    const ModuleAnalysisState &moduleState = getModuleAnalysisState(state);
+    const FuncAnalysisState &funcState = getFuncAnalysisState(state);
     if (getFuncOpAnalysisState(state, funcOp) != FuncOpAnalysisState::Analyzed)
       // FuncOp not analyzed yet. Assume that OpOperand is written.
       return true;
 
-    return moduleState.writtenBbArgs.lookup(funcOp).contains(
+    return funcState.writtenBbArgs.lookup(funcOp).contains(
         opOperand.getOperandNumber());
   }
 
   SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
                                             const AnalysisState &state) const {
     func::CallOp callOp = cast<func::CallOp>(op);
     FuncOp funcOp = getCalledFunction(callOp);
     assert(funcOp && "expected CallOp to a FuncOp");
-    const ModuleAnalysisState &moduleState = getModuleAnalysisState(state);
+    const FuncAnalysisState &funcState = getFuncAnalysisState(state);
     if (getFuncOpAnalysisState(state, funcOp) !=
         FuncOpAnalysisState::Analyzed) {
       // FuncOp not analyzed yet. Any OpResult may be aliasing.
       SmallVector<OpResult> result;
       for (OpResult opResult : op->getOpResults())
         if (opResult.getType().isa<TensorType>())
           result.push_back(opResult);
       return result;
     }
 
     // Get aliasing results from state.
     auto aliasingReturnVals =
-        moduleState.aliasingReturnVals.lookup(funcOp).lookup(
+        funcState.aliasingReturnVals.lookup(funcOp).lookup(
             opOperand.getOperandNumber());
     SmallVector<OpResult> result;
     for (int64_t resultIdx : aliasingReturnVals)
       result.push_back(callOp->getOpResult(resultIdx));
     return result;
   }
 
   SmallVector<OpOperand *>
   getAliasingOpOperand(Operation *op, OpResult opResult,
                        const AnalysisState &state) const {
     func::CallOp callOp = cast<func::CallOp>(op);
     FuncOp funcOp = getCalledFunction(callOp);
     assert(funcOp && "expected CallOp to a FuncOp");
-    const ModuleAnalysisState &moduleState = getModuleAnalysisState(state);
+    const FuncAnalysisState &funcState = getFuncAnalysisState(state);
     if (getFuncOpAnalysisState(state, funcOp) !=
         FuncOpAnalysisState::Analyzed) {
       // FuncOp not analyzed yet. Any OpOperand may be aliasing.
       SmallVector<OpOperand *> result;
       for (OpOperand &opOperand : op->getOpOperands())
         if (opOperand.get().getType().isa<TensorType>())
           result.push_back(&opOperand);
       return result;
     }
 
     // Get aliasing bbArgs from state.
-    auto aliasingFuncArgs = moduleState.aliasingFuncArgs.lookup(funcOp).lookup(
+    auto aliasingFuncArgs = funcState.aliasingFuncArgs.lookup(funcOp).lookup(
         opResult.getResultNumber());
     SmallVector<OpOperand *> result;
     for (int64_t bbArgIdx : aliasingFuncArgs)
       result.push_back(&callOp->getOpOperand(bbArgIdx));
     return result;
   }
 
   BufferRelation bufferRelation(Operation *op, OpResult opResult,
                                 const AnalysisState &state) const {
     return BufferRelation::Equivalent;
   }
 
   /// In a first approximation, all the function arguments of a FuncOp are
   /// marked inplaceable. For now, it is the responsibility of the `callOp`
   /// bufferization to allow FuncOp that are inplaceable to write inPlace.
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           BufferizationState &state) const {
     func::CallOp callOp = cast<func::CallOp>(op);
     unsigned numResults = callOp.getNumResults();
     unsigned numOperands = callOp->getNumOperands();
     FuncOp funcOp = getCalledFunction(callOp);
     assert(funcOp && "expected CallOp to a FuncOp");
-    const ModuleAnalysisState &moduleState =
-        getModuleAnalysisState(state.getAnalysisState());
+    const FuncAnalysisState &funcState =
+        getFuncAnalysisState(state.getAnalysisState());
     const OneShotBufferizationOptions &options =
         static_cast<const OneShotBufferizationOptions &>(state.getOptions());
 
     // Result types of the bufferized CallOp.
     SmallVector<Type> resultTypes;
     // Replacement values for the existing CallOp. These are usually the results
     // of the bufferized CallOp, unless a tensor result folds onto an operand.
     SmallVector<Value> replacementValues(numResults, Value());
for (const auto &it : llvm::enumerate(callOp.getResultTypes())) {
       if (!isaTensor(returnType)) {
         // Non-tensor values are returned.
         retValMapping[returnValIdx] = resultTypes.size();
         resultTypes.push_back(returnType);
         continue;
       }
 
       if (Optional<int64_t> bbArgIdx =
-              getEquivalentFuncArgIdx(funcOp, moduleState, returnValIdx)) {
+              getEquivalentFuncArgIdx(funcOp, funcState, returnValIdx)) {
         // Return operands that are equivalent to some bbArg, are not
         // returned.
         FailureOr<Value> bufferOrFailure =
             state.getBuffer(rewriter, callOp->getOpOperand(*bbArgIdx));
         if (failed(bufferOrFailure))
           return failure();
         replacementValues[returnValIdx] = *bufferOrFailure;
         newOperands[*bbArgIdx] = *bufferOrFailure;
if (bbArg.getType().isa<TensorType>())
       setInPlaceFuncArgument(bbArg, bufferizableOp.isWritable(bbArg, state));
 }
 
 LogicalResult mlir::linalg::comprehensive_bufferize::runModuleBufferize(
     ModuleOp moduleOp, OneShotBufferizationOptions options) {
   IRRewriter rewriter(moduleOp.getContext());
   OneShotAnalysisState analysisState(moduleOp, options);
   BufferizationState bufferizationState(analysisState);
-  ModuleAnalysisState &moduleState = getModuleAnalysisState(analysisState);
+  FuncAnalysisState &funcState = getFuncAnalysisState(analysisState);
   BufferizationAliasInfo &aliasInfo = analysisState.getAliasInfo();
 
-  if (failed(getFuncOpsOrderedByCalls(moduleOp, moduleState.orderedFuncOps,
-                                      moduleState.callerMap)))
+  if (failed(getFuncOpsOrderedByCalls(moduleOp, funcState.orderedFuncOps,
+                                      funcState.callerMap)))
     return failure();
 
   // Collect bbArg/return value information after the analysis.
   options.addPostAnalysisStep(aliasingFuncOpBBArgsAnalysis);
   options.addPostAnalysisStep(funcOpBbArgReadWriteAnalysis);
 
   // Analyze ops.
-  for (FuncOp funcOp : moduleState.orderedFuncOps) {
+  for (FuncOp funcOp : funcState.orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.getBody().empty())
       continue;
 
     // Now analyzing function.
-    moduleState.startFunctionAnalysis(funcOp);
+    funcState.startFunctionAnalysis(funcOp);
 
     // Gather equivalence info for CallOps.
-    equivalenceAnalysis(funcOp, aliasInfo, moduleState);
+    equivalenceAnalysis(funcOp, aliasInfo, funcState);
 
     // Analyze funcOp.
     if (failed(analyzeOp(funcOp, analysisState)))
       return failure();
 
     // Mark op as fully analyzed.
-    moduleState.analyzedFuncOps[funcOp] = FuncOpAnalysisState::Analyzed;
+    funcState.analyzedFuncOps[funcOp] = FuncOpAnalysisState::Analyzed;
 
     // Add annotations to function arguments.
     if (options.testAnalysisOnly)
       annotateOpsWithBufferizationMarkers(funcOp, analysisState);
   }
 
   if (options.testAnalysisOnly)
     return success();
 
   // Bufferize functions.
-  for (FuncOp funcOp : moduleState.orderedFuncOps) {
+  for (FuncOp funcOp : funcState.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 : moduleState.orderedFuncOps) {
+  for (FuncOp funcOp : funcState.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();
     }
   }