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 @@ -1062,10 +1062,19 @@ }); } -/// Assert that IR is in destination-passing style. I.e., every value that is -/// returned or yielded from a block is: -/// * aliasing a bbArg of that block or a parent block, or -/// * aliasing an OpResult of a op in a parent block. +/// Assert that every allocation can be deallocated in the same block. I.e., +/// every value that is returned or yielded from a block is: +/// * guaranteed to be aliasing a bbArg of that block or a parent block, or +/// * guaranteed to be aliasing an OpResult of a op in a parent block. +/// +/// In that case, buffer deallocation is simple: Every allocated buffer can be +/// deallocated in the same block. Otherwise, the buffer deallocation pass must +/// be run. +/// +/// Note: The current implementation checks for equivalent values instead of +/// aliasing values, which is stricter than needed. We can currently not check +/// for aliasing values because the analysis is a maybe-alias analysis and we +/// need a must-alias analysis here. /// /// Example: /// ``` @@ -1077,23 +1086,19 @@ /// scf.yield %t : tensor /// } /// ``` -/// In the above example, the first scf.yield op satifies destination-passing -/// style because the yielded value %0 is defined in the parent block. The -/// second scf.yield op does not satisfy destination-passing style because the -/// yielded value %t is defined in the same block as the scf.yield op. -// TODO: The current implementation checks for equivalent values instead of -// aliasing values, which is stricter than needed. We can currently not check -// for aliasing values because the analysis is a maybe-alias analysis and we -// need a must-alias analysis here. -static LogicalResult -assertDestinationPassingStyle(Operation *op, AnalysisState &state, - BufferizationAliasInfo &aliasInfo, - SmallVector &newOps) { +/// +/// In the above example, the second scf.yield op is problematic because the +/// yielded value %t is defined in the same block as the scf.yield op and +/// and bufferizes to a new allocation. +// TODO: Remove buffer deallocation from One-Shot Bufferize and fix the buffer +// deallocation pass. +static LogicalResult assertNoAllocsReturned(Operation *op, + const BufferizationOptions &options, + BufferizationAliasInfo &aliasInfo) { LogicalResult status = success(); DominanceInfo domInfo(op); op->walk([&](Operation *returnOp) { - if (!isRegionReturnLike(returnOp) || - !state.getOptions().isOpAllowed(returnOp)) + if (!isRegionReturnLike(returnOp) || !options.isOpAllowed(returnOp)) return WalkResult::advance(); for (OpOperand &returnValOperand : returnOp->getOpOperands()) { @@ -1119,11 +1124,12 @@ foundEquivValue = true; }); + // Note: Returning/yielding buffer allocations is allowed only if + // `allowReturnAllocs` is set. if (!foundEquivValue) - status = - returnOp->emitError() - << "operand #" << returnValOperand.getOperandNumber() - << " of ReturnLike op does not satisfy destination passing style"; + status = returnOp->emitError() + << "operand #" << returnValOperand.getOperandNumber() + << " may return/yield a new buffer allocation"; } return WalkResult::advance(); @@ -1148,11 +1154,8 @@ equivalenceAnalysis(op, aliasInfo, state); bool failedAnalysis = false; - if (!options.allowReturnAllocs) { - SmallVector newOps; - failedAnalysis |= - failed(assertDestinationPassingStyle(op, state, aliasInfo, newOps)); - } + if (!options.allowReturnAllocs) + failedAnalysis |= failed(assertNoAllocsReturned(op, options, aliasInfo)); // Gather some extra analysis data. state.gatherYieldedTensors(op); diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir --- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir @@ -44,10 +44,10 @@ } else { // This buffer aliases, but it is not equivalent. %t2 = tensor.extract_slice %t1 [%idx] [%idx] [1] : tensor to tensor - // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}} + // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}} scf.yield %t2 : tensor } - // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}} + // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}} return %r : tensor } @@ -61,7 +61,7 @@ } else { // This buffer aliases. %t2 = bufferization.alloc_tensor(%idx) : tensor - // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}} + // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}} scf.yield %t2 : tensor } %f = tensor.extract %r[%idx] : tensor @@ -190,7 +190,7 @@ // argument aliasing). %r0 = tensor.extract_slice %A[0][4][1] : tensor to tensor<4xf32> - // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}} + // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}} return %r0: tensor<4xf32> } @@ -203,7 +203,7 @@ } else { scf.yield %B : tensor<4xf32> } - // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}} + // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}} return %r: tensor<4xf32> } @@ -213,7 +213,7 @@ { // expected-error: @+1 {{op was not bufferized}} %r = "marklar"(%A) : (tensor<4xf32>) -> (tensor<4xf32>) - // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}} + // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}} return %r: tensor<4xf32> } @@ -223,7 +223,7 @@ %f0 = arith.constant 0.0 : f32 %t = bufferization.alloc_tensor() : tensor<10x20xf32> %r = linalg.fill ins(%f0 : f32) outs(%t : tensor<10x20xf32>) -> tensor<10x20xf32> - // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}} + // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}} return %r : tensor<10x20xf32> } @@ -232,11 +232,11 @@ func.func @main() -> tensor<4xi32> { %r = scf.execute_region -> tensor<4xi32> { %A = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32> - // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}} + // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}} scf.yield %A: tensor<4xi32> } - // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}} + // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}} return %r: tensor<4xi32> } @@ -275,7 +275,7 @@ func.func @foo(%t : tensor<5xf32>) -> (tensor<5xf32>) { %0 = bufferization.alloc_tensor() : tensor<5xf32> - // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}} + // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}} return %0 : tensor<5xf32> } @@ -288,11 +288,11 @@ // ----- -func.func @destination_passing_style_dominance_test_1(%cst : f32, %idx : index, - %idx2 : index) -> f32 { +func.func @yield_alloc_dominance_test_1(%cst : f32, %idx : index, + %idx2 : index) -> f32 { %0 = scf.execute_region -> tensor { %1 = bufferization.alloc_tensor(%idx) : tensor - // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}} + // expected-error @+1 {{operand #0 may return/yield a new buffer allocation}} scf.yield %1 : tensor } %2 = tensor.insert %cst into %0[%idx] : tensor @@ -302,12 +302,13 @@ // ----- -func.func @destination_passing_style_dominance_test_2(%cst : f32, %idx : index, - %idx2 : index) -> f32 { +func.func @yield_alloc_dominance_test_2(%cst : f32, %idx : index, + %idx2 : index) -> f32 { %1 = bufferization.alloc_tensor(%idx) : tensor %0 = scf.execute_region -> tensor { - // This YieldOp is in destination-passing style, thus no error. + // This YieldOp returns a value that is defined in a parent block, thus + // no error. scf.yield %1 : tensor } %2 = tensor.insert %cst into %0[%idx] : tensor