[mlir][bufferization] Better error checking for ops with unstructured control flow

[springerm]

Summary

Report an error when trying to bufferize an op that contains unstructured control flow but for ops for which the bufferization implementation does not support unstructured control flow. At the moment, there are no ops for which unstructured control flow is supported.

[maerhart]

I have a few questions regarding the motivation of this change:
Is this intended to be a temporal solution to slowly add support for unstructured control-flow one op after another?
Shouldn't the goal be that every op that implements the bufferization interface and may have multiple blocks in a region (unstructured control-flow) should implement the interface in a way that the region can be properly bufferized?
Are there fundamental limitations to when unstructured contol flow can be bufferized? If so, is it fine-grained enough to check for multiple blocks in a region or should there be some method like isBufferizable such that ops can check whether they are bufferizable more dynamically and provide an adequate error message if not?

[springerm]

In D157863#4585269, @maerhart wrote:

Is this intended to be a temporal solution to slowly add support for unstructured control-flow one op after another?

Kind of. The blocks of a region are bufferized together with the owner of that region. E.g., the BufferizableOpInterface implementation of func.func is responsible for bufferizing all block signatures directly nested under that op. We do not support unstructured control flow yet, so no implementations actually support that at the moment. I added a new interface method so that interface implementations that support unstructured control flow can override that function with true. This gives us proper error messages instead of failed assertions in case it is not supported.

I expect this new interface method to stay indefinitely to give proper error messages for downstream ops that may support unstructured control flow and where the implementer of the BufferizableOpInterface did not account for that.

Shouldn't the goal be that every op that implements the bufferization interface and may have multiple blocks in a region (unstructured control-flow) should implement the interface in a way that the region can be properly bufferized?

Yes. In practice we will not have many ops. func.func and scf.execute_region should be the only ops in MLIR for now.

Are there fundamental limitations to when unstructured contol flow can be bufferized?

No, a simple bufferization strategy that always works is "make a copy of the buffer and pass that one to the cf.br."

[maerhart]

Thanks for providing more info! In that case it's not necessary to allow more fine-grained decisions.
I guess that, given all terminators in the region implement the BranchOpInterface, it is possible to provide the bufferization logic for unstructured control flow as a library function that the downstream users could just call and thus make it very easy to add support for them?
Nonetheless, I think it's better to present some nice error message if support was not added rather than throwing an assertion (even though you could argue that an assertion is good enough when an interface is implemented incorrectly).

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

535

Would it make sense to make this a static method since it probably does not need to access fields of the op, but only informs the user of the interface (i.e., one-shot bufferization) about a property of the interface implementation?

[springerm]

In D157863#4587562, @maerhart wrote:

I guess that, given all terminators in the region implement the BranchOpInterface, it is possible to provide the bufferization logic for unstructured control flow as a library function that the downstream users could just call and thus make it very easy to add support for them?

Yes, I have a WIP revision that does that.

[springerm]

This has landed as 061aa2e3ba27.

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);
         }]
-      >
+      >,
+      InterfaceMethod<
+        /*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>
+}