[mlir][bufferize][NFC] Clean up ModuleBufferizationState

[springerm]

Summary

• Store bbArg indices instead of BlockArguments, so that args can be changed during bufferizationn.
• Use type aliases for better readability.

[springerm]

This is factored out from D119120, landing directly as discussed.

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

 
 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 {
+  /// A set of block argument indices.
+  using BbArgIndexSet = DenseSet<int64_t>;
+
+  /// A mapping of indices to indices.
+  using IndexMapping = DenseMap<int64_t, int64_t>;
+
   /// A mapping of ReturnOp OpOperand indices to equivalent FuncOp BBArg
   /// indices.
-  DenseMap<FuncOp, DenseMap<int64_t, int64_t>> equivalentFuncArgs;
+  DenseMap<FuncOp, IndexMapping> equivalentFuncArgs;
 
   /// A set of all read BlockArguments of FuncOps.
-  // Note: BlockArgument knows about its owner, so we do not need to store
-  // FuncOps here.
-  DenseSet<BlockArgument> readBbArgs;
+  DenseMap<FuncOp, BbArgIndexSet> readBbArgs;
 
   /// A set of all written-to BlockArguments of FuncOps.
-  DenseSet<BlockArgument> writtenBbArgs;
+  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.
+  /// A list of functions in the order in which they are analyzed + bufferized.
   SmallVector<FuncOp> orderedFuncOps;
 
-  // A mapping of FuncOps to their callers.
+  /// 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 createdRead = readBbArgs.try_emplace(funcOp, BbArgIndexSet());
+    auto createdWritten = writtenBbArgs.try_emplace(funcOp, BbArgIndexSet());
+    (void)createdEquiv;
+    (void)createdRead;
+    (void)createdWritten;
+#ifndef NDEBUG
+    assert(createdEquiv.second && "equivalence 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>(
funcOpBbArgReadWriteAnalysis(Operation *op, AnalysisState &state,
                              SmallVector<Operation *> &newOps) {
   ModuleAnalysisState &moduleState = getModuleAnalysisState(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.insert(bbArg);
-      moduleState.writtenBbArgs.insert(bbArg);
+      moduleState.readBbArgs[funcOp].insert(bbArg.getArgNumber());
+      moduleState.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.insert(bbArg);
+      moduleState.readBbArgs[funcOp].insert(bbArg.getArgNumber());
     if (isWritten)
-      moduleState.writtenBbArgs.insert(bbArg);
+      moduleState.writtenBbArgs[funcOp].insert(bbArg.getArgNumber());
   }
 
   return success();
 }
 } // namespace
 
 static bool isaTensor(Type t) { return t.isa<TensorType>(); }
 
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
     FuncOp funcOp = getCalledFunction(callOp);
     assert(funcOp && "expected CallOp to a FuncOp");
 
     const ModuleAnalysisState &moduleState = getModuleAnalysisState(state);
     if (getFuncOpAnalysisState(state, funcOp) != FuncOpAnalysisState::Analyzed)
       // FuncOp not analyzed yet. Assume that OpOperand is read.
       return true;
 
-    return moduleState.readBbArgs.contains(
-        funcOp.getArgument(opOperand.getOperandNumber()));
+    return moduleState.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);
     if (getFuncOpAnalysisState(state, funcOp) != FuncOpAnalysisState::Analyzed)
       // FuncOp not analyzed yet. Assume that OpOperand is written.
       return true;
 
-    return moduleState.writtenBbArgs.contains(
-        funcOp.getArgument(opOperand.getOperandNumber()));
+    return moduleState.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);
LogicalResult mlir::linalg::comprehensive_bufferize::runModuleBufferize(
 
   // Analyze ops.
   for (FuncOp funcOp : moduleState.orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.getBody().empty())
       continue;
 
     // Now analyzing function.
-    moduleState.analyzedFuncOps[funcOp] = FuncOpAnalysisState::InProgress;
+    moduleState.startFunctionAnalysis(funcOp);
 
     // Analyze funcOp.
     if (failed(analyzeOp(funcOp, analysisState)))
       return failure();
 
     // Gather equivalence info for CallOps.
     // TODO: Make this a post-analysis step.
     equivalenceAnalysis(funcOp, aliasInfo, moduleState);