[mlir][bufferize] Simplify ModuleBufferization driver

[springerm]

Summary

• 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.

Depends On D123192

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

LogicalResult mlir::linalg::comprehensive_bufferize::runModuleBufferize(
   for (FuncOp funcOp : moduleState.orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.getBody().empty())
       continue;
 
     // 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;
 
     // Add annotations to function arguments.
     if (options.testAnalysisOnly)
       annotateOpsWithBufferizationMarkers(funcOp, analysisState);
   }
 
   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 (!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>();
         })) {
       funcOp->emitError("memref return type is unsupported");
       return failure();
     }
   }

mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir

func @to_memref_op_is_writing(
   %r1 = vector.transfer_read %t1[%idx3], %cst : tensor<?xf32>, vector<5xf32>
   %r2 = vector.transfer_read %0[%idx3], %cst : memref<?xf32>, vector<5xf32>
 
   return %r1, %r2 : vector<5xf32>, vector<5xf32>
 }
 
 // -----
 
+// 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
 }
 
 // -----
 
 func @foo(%t : tensor<5xf32>) -> (tensor<5xf32>) {
   %0 = linalg.init_tensor [5] : tensor<5xf32>