[Mlir-commits] [mlir] cb64c54 - [mlir][linalg][bufferize][NFC] Clean up bufferization entry point

[Matthias Springer]

Author: Matthias Springer
Date: 2022-01-10T21:50:36+09:00
New Revision: cb64c541f92f2bbec8d8cb6db80490d126e96c76

URL: https://github.com/llvm/llvm-project/commit/cb64c541f92f2bbec8d8cb6db80490d126e96c76
DIFF: https://github.com/llvm/llvm-project/commit/cb64c541f92f2bbec8d8cb6db80490d126e96c76.diff

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

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h
    mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h
    mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
    mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp
    mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h
index 75725e26cda81..a6b34f8fa643a 100644
--- a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h
+++ b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h
@@ -24,15 +24,11 @@ class BufferizationState;
 /// `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 the given operation.
+/// Bufferize `op` and its nested ops. Bufferization decisions are stored in
+/// `state`.
+LogicalResult bufferizeOp(Operation *op, BufferizationState &state);
+
+/// Run Comprehensive Bufferize on the given op: Analysis + Bufferization
 LogicalResult
 runComprehensiveBufferize(Operation *op,
                           std::unique_ptr<BufferizationOptions> options);

diff  --git a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h
index 88ccd0ea7727a..7d155512ef87d 100644
--- a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h
+++ b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h
@@ -22,8 +22,9 @@ 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);

diff  --git a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
index f625bc702939c..69a55b2e49d52 100644
--- a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
+++ b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
@@ -555,12 +555,6 @@ annotateOpsWithBufferizationMarkers(Operation *op,
   });
 }
 
-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> {
@@ -652,18 +646,26 @@ mlir::linalg::comprehensive_bufferize::analyzeOp(Operation *op,
   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();
 }

diff  --git a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp
index be49608c0bf92..95ecb21cf8e96 100644
--- a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp
+++ b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp
@@ -886,8 +886,7 @@ LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
     if (funcOp.body().empty())
       continue;
 
-    if (failed(runComprehensiveBufferize(funcOp, *options, state,
-                                         /*runAnalysis=*/false)))
+    if (failed(bufferizeOp(funcOp, state)))
       return failure();
   }
 

diff  --git a/mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp b/mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp
index 59e53eaba56a9..ae5252b7c3c1a 100644
--- a/mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp
+++ b/mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp
@@ -117,6 +117,9 @@ void TestComprehensiveFunctionBufferize::runOnFunction() {
   if (failed(runComprehensiveBufferize(op, std::move(options))))
     return;
 
+  if (testAnalysisOnly)
+    return;
+
   OpPassManager cleanupPipeline("builtin.func");
   cleanupPipeline.addPass(createCanonicalizerPass());
   cleanupPipeline.addPass(createCSEPass());