[Mlir-commits] [mlir] cd7de44 - [mlir][bufferize] Simplify ModuleBufferization driver

[Matthias Springer]

Author: Matthias Springer
Date: 2022-04-06T23:53:07+09:00
New Revision: cd7de446fd27ba7f9865087ce1b97a3e01d42cee

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

LOG: [mlir][bufferize] Simplify ModuleBufferization driver

* Bufferize FuncOp bodies and boundaries in the same loop. This is in preparation of moving FuncOp bufferization into an external model implementation.
* As a side effect, stop bufferization earlier if there was an error. (Do not continue bufferization, fewer error messages.)
* Run equivalence analysis of CallOps before the main analysis. This is needed so that equialvence info is propagated properly.

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

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp
    mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp
index 03a754d3e8443..3ef695f414631 100644
--- a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp
+++ b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp
@@ -1033,14 +1033,13 @@ LogicalResult mlir::linalg::comprehensive_bufferize::runModuleBufferize(
     // Now analyzing function.
     moduleState.startFunctionAnalysis(funcOp);
 
+    // Gather equivalence info for CallOps.
+    equivalenceAnalysis(funcOp, aliasInfo, moduleState);
+
     // Analyze funcOp.
     if (failed(analyzeOp(funcOp, analysisState)))
       return failure();
 
-    // Gather equivalence info for CallOps.
-    // TODO: Make this a post-analysis step.
-    equivalenceAnalysis(funcOp, aliasInfo, moduleState);
-
     // Mark op as fully analyzed.
     moduleState.analyzedFuncOps[funcOp] = FuncOpAnalysisState::Analyzed;
 
@@ -1052,23 +1051,21 @@ LogicalResult mlir::linalg::comprehensive_bufferize::runModuleBufferize(
   if (options.testAnalysisOnly)
     return success();
 
-  // Bufferize function bodies.
+  // Bufferize functions.
   for (FuncOp funcOp : moduleState.orderedFuncOps) {
     // No body => no analysis.
-    if (funcOp.getBody().empty())
-      continue;
-
-    if (failed(bufferizeOp(funcOp, bufferizationState)))
-      return failure();
-  }
+    if (!funcOp.getBody().empty())
+      if (failed(bufferizeOp(funcOp, bufferizationState)))
+        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, bufferizationState)))
       return failure();
+  }
 
+  // Check result.
+  for (FuncOp funcOp : moduleState.orderedFuncOps) {
     if (!options.allowReturnAllocs &&
         llvm::any_of(funcOp.getFunctionType().getResults(), [](Type t) {
           return t.isa<MemRefType, UnrankedMemRefType>();

diff  --git a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir
index 02410c811e933..85abc8883d208 100644
--- a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir
+++ b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir
@@ -212,11 +212,10 @@ func @to_memref_op_is_writing(
 
 // -----
 
+// expected-error @+1 {{cannot bufferize bodiless function that returns a tensor}}
 func private @foo(%t : tensor<?xf32>) -> (f32, tensor<?xf32>, f32)
 
 func @call_to_unknown_tensor_returning_func(%t : tensor<?xf32>) {
-  // expected-error @+2 {{call to FuncOp that returns non-equivalent tensors not supported}}
-  // expected-error @+1 {{op was not bufferized}}
   call @foo(%t) : (tensor<?xf32>) -> (f32, tensor<?xf32>, f32)
   return
 }