[mlir][linalg][bufferize][NFC] Add `analyzeOp` helper function

This function runs just the analysis of Comprehensive Bufferize, but does not bufferize the IR yet.

This is in preparation of fixing CallOp bufferization. Also needed for unifying Comprehensive Bufferize and core bufferization; the new partial bufferization can simply run bufferization without an analysis.

Differential Revision: https://reviews.llvm.org/D116456

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

@@ -16,14 +16,21 @@ namespace mlir {
 namespace linalg {
 namespace comprehensive_bufferize {
 
+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);
+                                        BufferizationState &state,
+                                        bool runAnalysis = true);
 
 /// Bufferize the given operation.
 LogicalResult

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

@@ -618,13 +618,12 @@ checkBufferizationResult(Operation *op, const BufferizationOptions &options) {
   return success();
 }
 
-LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
-    Operation *op, const BufferizationOptions &options,
-    BufferizationState &state) {
-
-  IRRewriter rewriter(op->getContext());
+LogicalResult
+mlir::linalg::comprehensive_bufferize::analyzeOp(Operation *op,
+                                                 BufferizationState &state) {
   DominanceInfo domInfo(op);
   BufferizationAliasInfo &aliasInfo = state.getAliasInfo();
+  const BufferizationOptions &options = state.getOptions();
 
   if (failed(checkAliasInfoConsistency(op, domInfo, state, aliasInfo)))
     return failure();
@@ -647,10 +646,18 @@ LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
   }
 
   // Annotate operations if we only want to report the analysis.
-  if (options.testAnalysisOnly) {
+  if (options.testAnalysisOnly)
     annotateOpsWithBufferizationMarkers(op, aliasInfo, state);
-    return success();
-  }
+
+  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();
 
   // Bufferize the op and its nested ops.
   OwningRewritePatternList patterns(op->getContext());

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

@@ -842,6 +842,8 @@ LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
   // inplace. Therefore, we just bufferize funcOp as if none of its results were
   // inplaceable, detect which operands are cloned internally and decide what to
   // do at call sites.
+
+  // Analyze ops.
   for (FuncOp funcOp : moduleState.orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.body().empty())
@@ -854,8 +856,8 @@ LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
     // Gather equivalence info for CallOps.
     equivalenceAnalysis(funcOp, aliasInfo, moduleState);
 
-    // Analyze and bufferize funcOp.
-    if (failed(runComprehensiveBufferize(funcOp, *options, state)))
+    // Analyze funcOp.
+    if (failed(analyzeOp(funcOp, state)))
       return failure();
 
     // Add annotations to function arguments.
@@ -866,6 +868,18 @@ LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
   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)))
+      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.