[mlir][linalg][bufferize][NFC] Clean up bufferization entry point

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h

 class BufferizationAliasInfo;
 struct BufferizationOptions;
 class BufferizationState;
 
 /// Analyze `op` and its nested ops. Bufferization decisions are stored in
 /// `state`.
 LogicalResult analyzeOp(Operation *op, BufferizationState &state);
 
-/// Bufferize the given operation. Reuses an existing BufferizationState object.
-/// If `runAnalysis` is set to false, all OpOperands bufferize out-of-place.
-/// This function overload is for internal usage only.
-LogicalResult runComprehensiveBufferize(Operation *op,
-                                        const BufferizationOptions &options,
-                                        BufferizationState &state,
-                                        bool runAnalysis = true);
+/// Bufferize `op` and its nested ops. Bufferization decisions are stored in
+/// `state`.
+LogicalResult bufferizeOp(Operation *op, BufferizationState &state);
 
-/// Bufferize the given operation.
+/// Run Comprehensive Bufferize on the given op: Analysis + Bufferization
 LogicalResult
 runComprehensiveBufferize(Operation *op,
                           std::unique_ptr<BufferizationOptions> options);
 
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_COMPREHENSIVE_BUFFERIZE_H

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h

 struct LogicalResult;
 class ModuleOp;
 
 namespace linalg {
 namespace comprehensive_bufferize {
 
 struct BufferizationOptions;
 
-/// Bufferize the given module. This bufferizations performs a simple function
-/// call analysis to determine which function arguments are inplaceable.
+/// Run 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 Comprehensive Bufferization.
 LogicalResult
 runComprehensiveBufferize(ModuleOp moduleOp,
                           std::unique_ptr<BufferizationOptions> options);
 
 namespace std_ext {
 
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
 } // namespace std_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_MODULE_BUFFERIZATION_H

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

annotateOpsWithBufferizationMarkers(Operation *op,
   op->walk([&](Operation *op) {
     if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(op))
       for (OpOperand &opOperand : op->getOpOperands())
         if (opOperand.get().getType().isa<TensorType>())
           setInPlaceOpOperand(opOperand, aliasInfo.isInPlace(opOperand));
   });
 }
 
-LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
-    Operation *op, std::unique_ptr<BufferizationOptions> options) {
-  BufferizationState state(op, *options);
-  return runComprehensiveBufferize(op, *options, state);
-}
-
 /// Rewrite pattern that bufferizes bufferizable ops.
 struct BufferizationPattern
     : public OpInterfaceRewritePattern<BufferizableOpInterface> {
   BufferizationPattern(MLIRContext *context, BufferizationState &state,
                        PatternBenefit benefit = 1)
       : OpInterfaceRewritePattern<BufferizableOpInterface>(context, benefit),
         state(state) {}
 
mlir::linalg::comprehensive_bufferize::analyzeOp(Operation *op,
 
   // Annotate operations if we only want to report the analysis.
   if (options.testAnalysisOnly)
     annotateOpsWithBufferizationMarkers(op, aliasInfo, state);
 
   return success();
 }
 
-LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
-    Operation *op, const BufferizationOptions &options,
-    BufferizationState &state, bool runAnalysis) {
-  if (runAnalysis)
-    if (failed(analyzeOp(op, state)))
-      return failure();
-
+LogicalResult
+mlir::linalg::comprehensive_bufferize::bufferizeOp(Operation *op,
+                                                   BufferizationState &state) {
   // Bufferize the op and its nested ops.
   OwningRewritePatternList patterns(op->getContext());
   patterns.add<BufferizationPattern>(op->getContext(), state);
   if (failed(applyPatternsAndFoldGreedily(op, std::move(patterns))))
     return failure();
 
-  return checkBufferizationResult(op, options);
+  return checkBufferizationResult(op, state.getOptions());
+}
+
+LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
+    Operation *op, std::unique_ptr<BufferizationOptions> options) {
+  BufferizationState state(op, *options);
+  if (failed(analyzeOp(op, state)))
+    return failure();
+  if (options->testAnalysisOnly)
+    return success();
+  if (failed(bufferizeOp(op, state)))
+    return failure();
+  return success();
 }

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

if (options->testAnalysisOnly)
     return success();
 
   // Bufferize function bodies.
   for (FuncOp funcOp : moduleState.orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.body().empty())
       continue;
 
-    if (failed(runComprehensiveBufferize(funcOp, *options, state,
-                                         /*runAnalysis=*/false)))
+    if (failed(bufferizeOp(funcOp, state)))
       return failure();
   }
 
   // Bufferize function boundaries.
   for (FuncOp funcOp : moduleState.orderedFuncOps) {
     // Note: It would be good to apply cleanups here but we cannot as aliasInfo
     // would be invalidated.
     if (failed(bufferizeFuncOpBoundary(funcOp, rewriter, state)))

mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp

if (dialectFilter.hasValue()) {
     for (const std::string &dialectNamespace : dialectFilter)
       options->dialectFilter->insert(dialectNamespace);
   }
 
   Operation *op = getFunction().getOperation();
   if (failed(runComprehensiveBufferize(op, std::move(options))))
     return;
 
+  if (testAnalysisOnly)
+    return;
+
   OpPassManager cleanupPipeline("builtin.func");
   cleanupPipeline.addPass(createCanonicalizerPass());
   cleanupPipeline.addPass(createCSEPass());
   cleanupPipeline.addPass(createLoopInvariantCodeMotionPass());
   (void)this->runPipeline(cleanupPipeline, op);
 }
 
 namespace mlir {
 namespace test {
 void registerTestComprehensiveFunctionBufferize() {
   PassRegistration<TestComprehensiveFunctionBufferize>();
 }
 } // namespace test
 } // namespace mlir