From 089b910abc49043d3972b911057411ef059c1a55 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm@google.com>
Date: Sat, 8 Jan 2022 01:26:14 +0900
Subject: [PATCH] [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
---
 .../ComprehensiveBufferize.h                  |  9 +++++++-
 .../ComprehensiveBufferize.cpp                | 23 ++++++++++++-------
 .../ModuleBufferization.cpp                   | 18 +++++++++++++--
 3 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h
index a984f5749380..75725e26cda8 100644
--- a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h
+++ b/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
diff --git a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
index 312e7421b5e1..f625bc702939 100644
--- a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
+++ b/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());
diff --git a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp
index bb9a06c25150..8c6b32d73317 100644
--- a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp
+++ b/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.