[mlir][func][bufferization][NFC] Simplify implementation

[springerm]

Summary

The bufferization implementation of func.func and func.call can be simplified. It still contained code that was necessary when One-Shot Bufferize removed return values. This functionality has been extracted into a separate pass a while ago.

Depends On: D157863

[springerm]

more simplifications

[mehdi_amini]

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td

546

There is no way for a static method to be overridden, the body is supposed to use information from the concrete op class here otherwise it does not make sense.
Shall we make it a non-static method?

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td

let methods = [
         /*retType=*/"bool",
         /*methodName=*/"isRepetitiveRegion",
         /*args=*/(ins "unsigned":$index),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           return ::mlir::bufferization::detail::defaultIsRepetitiveRegion(
               ::llvm::cast<BufferizableOpInterface>($_op.getOperation()), index);
         }]
-      >
+      >,
+      StaticInterfaceMethod<
+        /*desc=*/[{
+          Return `true` if the op and this interface implementation supports
+          unstructured control flow. I.e., regions with multiple blocks. This is
+          not supported in most ops, so the default answer is `false`.
+        }],
+        /*retType=*/"bool",
+        /*methodName=*/"supportsUnstructuredControlFlow",
+        /*args=*/(ins),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          return false;
+        }]
+      >,
   ];
 
   let extraClassDeclaration = [{
     /// Resolve out-of-place tensor OpOperands with explicit allocations in the
     /// form of `bufferization.alloc_tensor` ops.
     ::mlir::LogicalResult resolveTensorOpOperandConflicts(
         ::mlir::RewriterBase &rewriter,
         const ::mlir::bufferization::AnalysisState &state);

mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp

for (unsigned i = 0; i < worklist.size(); ++i) {
     auto bufferizableOp = options.dynCastBufferizableOp(nextOp);
     if (!bufferizableOp)
       continue;
     if (opFilter && !opFilter->isOpAllowed(nextOp))
       continue;
     // Skip ops that no longer have tensor semantics.
     if (!hasTensorSemantics(nextOp))
       continue;
+    // Check for unsupported unstructured control flow.
+    if (!bufferizableOp.supportsUnstructuredControlFlow())
+      for (Region &r : nextOp->getRegions())
+        if (r.getBlocks().size() > 1)
+          return nextOp->emitOpError(
+              "op or BufferizableOpInterface implementation does not support "
+              "unstructured control flow, but at least one region has multiple "
+              "blocks");
+
     // Bufferize the op.
     LLVM_DEBUG(llvm::dbgs()
                << "//===-------------------------------------------===//\n"
                << "IR after bufferizing: " << nextOp->getName() << "\n");
     rewriter.setInsertionPoint(nextOp);
     if (failed(bufferizableOp.bufferize(rewriter, options))) {
       LLVM_DEBUG(llvm::dbgs()
                  << "failed to bufferize\n"

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

LogicalResult OneShotAnalysisState::analyzeOp(Operation *op,
   } else {
     llvm_unreachable("unsupported heuristic");
   }
 
   equivalenceAnalysis(op, *this);
   return success();
 }
 
-/// Assert that the current bufferization decisions are consistent.
-static LogicalResult checkAliasInfoConsistency(Operation *op,
-                                               const DominanceInfo &domInfo,
-                                               OneShotAnalysisState &state) {
+/// Perform various checks on the input IR to see if it contains IR constructs
+/// that are unsupported by One-Shot Bufferize.
+static LogicalResult
+checkPreBufferizationAssumptions(Operation *op, const DominanceInfo &domInfo,
   const BufferizationOptions &options = state.getOptions();
 
+  // Note: This walk cannot be combined with the one below because interface
+  // methods of invalid/unsupported ops may be called during the second walk.
+  // (On ops different from `op`.)
   WalkResult walkResult = op->walk([&](BufferizableOpInterface op) {
     // Skip ops that are not in the filter.
     if (!options.isOpAllowed(op.getOperation()))
       return WalkResult::advance();
 
+    // Check for unsupported unstructured control flow.
+    if (!op.supportsUnstructuredControlFlow()) {
+      for (Region &r : op->getRegions()) {
+        if (r.getBlocks().size() > 1) {
+          op->emitOpError("op or BufferizableOpInterface implementation does "
+                          "not support unstructured control flow, but at least "
+                          "one region has multiple blocks");
+          return WalkResult::interrupt();
+        }
+      }
+    }
+
+    return WalkResult::advance();
+  });
+  if (walkResult.wasInterrupted())
+    return failure();
+
+  walkResult = op->walk([&](BufferizableOpInterface op) {
+    // Skip ops that are not in the filter.
+    if (!options.isOpAllowed(op.getOperation()))
+      return WalkResult::advance();
+
     // Input IR may not contain any ToTensorOps without the "restrict"
     // attribute. Such tensors may alias any other tensor, which is currently
     // not handled in the analysis.
     if (auto toTensorOp = dyn_cast<ToTensorOp>(op.getOperation())) {
       if (!toTensorOp.getRestrict() && !toTensorOp->getUses().empty()) {
         op->emitError("to_tensor ops without `restrict` are not supported by "
                       "One-Shot Analysis");
         return WalkResult::interrupt();

 }
 
 LogicalResult bufferization::analyzeOp(Operation *op,
                                        OneShotAnalysisState &state,
                                        BufferizationStatistics *statistics) {
   DominanceInfo domInfo(op);
   const OneShotBufferizationOptions &options = state.getOptions();
 
-  if (failed(checkAliasInfoConsistency(op, domInfo, state)))
+  if (failed(checkPreBufferizationAssumptions(op, domInfo, state)))
     return failure();
 
   // If the analysis fails, just return.
   if (failed(state.analyzeOp(op, domInfo)))
     return failure();
 
   if (statistics) {
     statistics->numTensorInPlace = state.getStatNumTensorInPlace();

mlir/test/Dialect/SCF/one-shot-bufferize-invalid.mlir

-// RUN: mlir-opt %s -one-shot-bufferize -split-input-file -verify-diagnostics
+// RUN: mlir-opt %s -one-shot-bufferize="allow-unknown-ops allow-return-allocs" -split-input-file -verify-diagnostics
 
 func.func @inconsistent_memory_space_scf_if(%c: i1) -> tensor<10xf32> {
   // Yielding tensors with different memory spaces. Such IR cannot be
   // bufferized.
   %0 = bufferization.alloc_tensor() {memory_space = 0 : ui64} : tensor<10xf32>
   %1 = bufferization.alloc_tensor() {memory_space = 1 : ui64} : tensor<10xf32>
   // expected-error @+1 {{inconsistent memory space on then/else branches}}
   %r = scf.if %c -> tensor<10xf32> {
     // expected-error @+1 {{failed to bufferize op}}
     scf.yield %0 : tensor<10xf32>
   } else {
     scf.yield %1 : tensor<10xf32>
   }
   func.return %r : tensor<10xf32>
 }
+
+// -----
+
+func.func @execute_region_multiple_blocks(%t: tensor<5xf32>) -> tensor<5xf32> {
+  // expected-error @below{{op or BufferizableOpInterface implementation does not support unstructured control flow, but at least one region has multiple blocks}}
+  %0 = scf.execute_region -> tensor<5xf32> {
+    cf.br ^bb1(%t : tensor<5xf32>)
+  ^bb1(%arg1 : tensor<5xf32>):
+    scf.yield %arg1 : tensor<5xf32>
+  }
+  func.return %0 : tensor<5xf32>
+}