[mlir][bufferize][NFC] Split analysis+bufferization of ModuleBufferization

[springerm]

Summary

Analysis and bufferization can now be run separately.

Depends On D126571

mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h

 #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTMODULEBUFFERIZE_H
 
 namespace mlir {
 
 struct LogicalResult;
 class ModuleOp;
 
 namespace bufferization {
+struct BufferizationState;
+class OneShotAnalysisState;
 struct OneShotBufferizationOptions;
 
+/// Analyze `moduleOp` and its nested ops. Bufferization decisions are stored in
+/// `state`.
+LogicalResult analyzeModuleOp(ModuleOp moduleOp, OneShotAnalysisState &state);
+
+/// Bufferize `op` and its nested ops that implement `BufferizableOpInterface`.
+/// Whether buffer copies are needed or not is queried from the given state.
+LogicalResult bufferizeModuleOp(ModuleOp moduleOp,
+                                const OneShotAnalysisState &analysisState);
+
 /// Run One-Shot Module Bufferization on the given module. Performs a simple
 /// function call analysis to determine which function arguments are
 /// inplaceable. Then analyzes and bufferizes FuncOps one-by-one with One-Shot
 /// Bufferize.
 LogicalResult
 runOneShotModuleBufferize(ModuleOp moduleOp,
                           bufferization::OneShotBufferizationOptions options);
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTMODULEBUFFERIZE_H

mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp

for (OpOperand &operand : returnOp->getOpOperands()) {
     }
   }
 
   auto newFuncType = FunctionType::get(
       funcOp.getContext(), funcOp.getFunctionType().getInputs(), resultTypes);
   funcOp.setType(newFuncType);
 }
 
-LogicalResult mlir::bufferization::runOneShotModuleBufferize(
-    ModuleOp moduleOp, OneShotBufferizationOptions options) {
+LogicalResult
+mlir::bufferization::analyzeModuleOp(ModuleOp moduleOp,
+                                     OneShotAnalysisState &state) {
+  OneShotBufferizationOptions options =
+      static_cast<const OneShotBufferizationOptions &>(state.getOptions());
   assert(options.bufferizeFunctionBoundaries &&
          "expected that function boundary bufferization is activated");
-  IRRewriter rewriter(moduleOp.getContext());
-  OneShotAnalysisState analysisState(moduleOp, options);
-  BufferizationState bufferizationState(analysisState);
-  FuncAnalysisState &funcState = getFuncAnalysisState(analysisState);
-  BufferizationAliasInfo &aliasInfo = analysisState.getAliasInfo();
+  FuncAnalysisState &funcState = getFuncAnalysisState(state);
+  BufferizationAliasInfo &aliasInfo = state.getAliasInfo();
 
   // A list of functions in the order in which they are analyzed + bufferized.
   SmallVector<func::FuncOp> orderedFuncOps;
 
   // A mapping of FuncOps to their callers.
   FuncCallerMap callerMap;
 
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
     return failure();
 
   // Analyze ops.
   for (func::FuncOp funcOp : orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.getBody().empty())
       continue;
 
     // Now analyzing function.
     funcState.startFunctionAnalysis(funcOp);
 
     // Gather equivalence info for CallOps.
     equivalenceAnalysis(funcOp, aliasInfo, funcState);
 
     // Analyze funcOp.
-    if (failed(analyzeOp(funcOp, analysisState)))
+    if (failed(analyzeOp(funcOp, state)))
       return failure();
 
     // Run some extra function analyses.
-    if (failed(aliasingFuncOpBBArgsAnalysis(funcOp, analysisState)) ||
-        failed(funcOpBbArgReadWriteAnalysis(funcOp, analysisState)))
+    if (failed(aliasingFuncOpBBArgsAnalysis(funcOp, state)) ||
+        failed(funcOpBbArgReadWriteAnalysis(funcOp, state)))
       return failure();
 
     // Mark op as fully analyzed.
     funcState.analyzedFuncOps[funcOp] = FuncOpAnalysisState::Analyzed;
 
     // Add annotations to function arguments.
     if (options.testAnalysisOnly)
-      annotateOpsWithBufferizationMarkers(funcOp, analysisState);
+      annotateOpsWithBufferizationMarkers(funcOp, state);
   }
 
-  if (options.testAnalysisOnly)
-    return success();
+}
+
+LogicalResult mlir::bufferization::bufferizeModuleOp(
+    ModuleOp moduleOp, const OneShotAnalysisState &analysisState) {
+  auto const &options = static_cast<const OneShotBufferizationOptions &>(
+      analysisState.getOptions());
+  assert(options.bufferizeFunctionBoundaries &&
+         "expected that function boundary bufferization is activated");
+  IRRewriter rewriter(moduleOp.getContext());
+  BufferizationState bufferizationState(analysisState);
+
+  // A list of functions in the order in which they are analyzed + bufferized.
+  SmallVector<func::FuncOp> orderedFuncOps;
+
+  // A mapping of FuncOps to their callers.
+  FuncCallerMap callerMap;
+
+  if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
+    return failure();
 
   // Bufferize functions.
   for (func::FuncOp funcOp : orderedFuncOps) {
     // Note: It would be good to apply cleanups here but we cannot as aliasInfo
     // would be invalidated.
     if (failed(bufferizeOp(funcOp, bufferizationState)))
       return failure();
     // Change buffer return types to more precise layout maps.
LogicalResult mlir::bufferization::bufferizeModuleOp(
   // Post-pass cleanup of function argument attributes.
   moduleOp.walk([&](func::FuncOp op) {
     for (BlockArgument bbArg : op.getArguments())
       removeBufferizationAttributes(bbArg);
   });
 
   return success();
 }
+
+LogicalResult mlir::bufferization::runOneShotModuleBufferize(
+    ModuleOp moduleOp, OneShotBufferizationOptions options) {
+  assert(options.bufferizeFunctionBoundaries &&
+         "expected that function boundary bufferization is activated");
+  OneShotAnalysisState analysisState(moduleOp, options);
+  if (failed(analyzeModuleOp(moduleOp, analysisState)))
+    return failure();
+  if (options.testAnalysisOnly)
+    return success();
+  if (failed(bufferizeModuleOp(moduleOp, analysisState)))
+    return failure();
+  return success();
+}