diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td --- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td +++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td @@ -531,7 +531,21 @@ return ::mlir::bufferization::detail::defaultIsRepetitiveRegion( ::llvm::cast($_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 = [{ diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp --- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp @@ -471,6 +471,15 @@ // 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" diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp --- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp @@ -948,17 +948,43 @@ 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, + OneShotAnalysisState &state) { 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. @@ -1107,7 +1133,7 @@ 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. diff --git a/mlir/test/Dialect/SCF/one-shot-bufferize-invalid.mlir b/mlir/test/Dialect/SCF/one-shot-bufferize-invalid.mlir --- a/mlir/test/Dialect/SCF/one-shot-bufferize-invalid.mlir +++ b/mlir/test/Dialect/SCF/one-shot-bufferize-invalid.mlir @@ -1,4 +1,4 @@ -// 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 @@ -14,3 +14,15 @@ } 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> +}