[Mlir-commits] [mlir] [mlir][bufferization] Never pass ownership to functions (PR #80655)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Feb 5 01:59:02 PST 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-bufferization
Author: Matthias Springer (matthias-springer)
<details>
<summary>Changes</summary>
Even when `private-function-dynamic-ownership` is set, ownership should never be passed to the callee. This can lead to double deallocs (#<!-- -->77096) or use-after-free in the caller because ownership is currently passed regardless of whether there are any further uses of the buffer in the caller or not.
Note: This is consistent with the fact that ownership is never passed to nested regions.
This commit fixes #<!-- -->77096.
---
Full diff: https://github.com/llvm/llvm-project/pull/80655.diff
5 Files Affected:
- (modified) mlir/docs/Bufferization.md (+22-30)
- (modified) mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h (+3-3)
- (modified) mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp (+13-32)
- (modified) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir (+2-2)
- (modified) mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir (+34-13)
``````````diff
diff --git a/mlir/docs/Bufferization.md b/mlir/docs/Bufferization.md
index 8329999162fb5..91cb35429433f 100644
--- a/mlir/docs/Bufferization.md
+++ b/mlir/docs/Bufferization.md
@@ -502,31 +502,26 @@ accordingly:
### Example
-The following example contains a few interesting cases:
-* Basic block arguments are modified to also pass along the ownership
- indicator, but not for entry bocks of non-private functions (assuming the
- `private-function-dynamic-ownership` pass option is disabled) where the
- function boundary ABI is applied instead. "Private" in this context refers
- to functions that cannot be called externally.
-* The result of `arith.select` initially has 'Unknown' assigned as ownership,
- but once the `bufferization.dealloc` operation is inserted it is put in the
- 'retained' list (since it has uses in a later basic block) and thus the
- 'Unknown' ownership can be replaced with a 'Unique' ownership using the
- corresponding result of the dealloc operation.
-* The `cf.cond_br` operation has more than one successor and thus has to
- insert two `bufferization.dealloc` operations (one for each successor).
- While they have the same list of MemRefs to deallocate (because they perform
- the deallocations for the same block), it must be taken into account that
- some MemRefs remain *live* for one branch but not the other (thus set
- intersection is performed on the *live-out* of the current block and the
- *live-in* of the target block). Also, `cf.cond_br` supports separate
- forwarding operands for each successor. To make sure that no MemRef is
- deallocated twice (because there are two `bufferization.dealloc` operations
- with the same MemRefs to deallocate), the condition operands are adjusted to
- take the branch condition into account. While a generic lowering for such
- terminator operations could be implemented, a specialized implementation can
- take all the semantics of this particular operation into account and thus
- generate a more efficient lowering.
+The following example contains a few interesting cases: * Basic block arguments
+are modified to also pass along the ownership indicator, but not for entry bocks
+where the function boundary ABI is applied instead. * The result of
+`arith.select` initially has 'Unknown' assigned as ownership, but once the
+`bufferization.dealloc` operation is inserted it is put in the 'retained' list
+(since it has uses in a later basic block) and thus the 'Unknown' ownership can
+be replaced with a 'Unique' ownership using the corresponding result of the
+dealloc operation. * The `cf.cond_br` operation has more than one successor and
+thus has to insert two `bufferization.dealloc` operations (one for each
+successor). While they have the same list of MemRefs to deallocate (because they
+perform the deallocations for the same block), it must be taken into account
+that some MemRefs remain *live* for one branch but not the other (thus set
+intersection is performed on the *live-out* of the current block and the
+*live-in* of the target block). Also, `cf.cond_br` supports separate forwarding
+operands for each successor. To make sure that no MemRef is deallocated twice
+(because there are two `bufferization.dealloc` operations with the same MemRefs
+to deallocate), the condition operands are adjusted to take the branch condition
+into account. While a generic lowering for such terminator operations could be
+implemented, a specialized implementation can take all the semantics of this
+particular operation into account and thus generate a more efficient lowering.
```mlir
func.func @example(%memref: memref<?xi8>, %select_cond: i1, %br_cond: i1) {
@@ -543,10 +538,7 @@ func.func @example(%memref: memref<?xi8>, %select_cond: i1, %br_cond: i1) {
After running `--ownership-based-buffer-deallocation`, it looks as follows:
```mlir
-// Since this is not a private function, the signature will not be modified even
-// when private-function-dynamic-ownership is enabled. Instead the function
-// boundary ABI has to be applied which means that ownership of `%memref` will
-// never be acquired.
+// Function boundary ABI: ownership of `%memref` will never be acquired.
func.func @example(%memref: memref<?xi8>, %select_cond: i1, %br_cond: i1) {
%false = arith.constant false
%true = arith.constant true
@@ -602,7 +594,7 @@ func.func @example(%memref: memref<?xi8>, %select_cond: i1, %br_cond: i1) {
: memref<i8>, memref<i8>, memref<i8>)
if (%false, %not_br_cond, %false)
retain (%memref, %select : memref<?xi8>, memref<?xi8>)
-
+
// Because %select is used in ^bb1 without passing it via block argument, we
// need to update it's ownership value here by merging the ownership values
// returned by the dealloc operations
diff --git a/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h b/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h
index 7acacb763cd2c..32565b61f5ff9 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h
@@ -24,9 +24,9 @@ struct BufferDeallocationPipelineOptions
PassOptions::Option<bool> privateFunctionDynamicOwnership{
*this, "private-function-dynamic-ownership",
llvm::cl::desc(
- "Allows to add additional arguments to private functions to "
- "dynamically pass ownership of memrefs to callees. This can enable "
- "earlier deallocations."),
+ "Allows to add additional results to private functions to return "
+ "ownership of returned memrefs to callers. This can avoid spurious "
+ "buffer clones in the callee."),
llvm::cl::init(false)};
};
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
index c535d6c130edb..fa4ea2c10842b 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp
@@ -292,11 +292,10 @@ class BufferDeallocation {
FailureOr<Operation *> handleInterface(RegionBranchOpInterface op);
/// If the private-function-dynamic-ownership pass option is enabled and the
- /// called function is private, additional arguments and results are added for
- /// each MemRef argument/result to pass the dynamic ownership indicator along.
- /// Otherwise, updates the ownership map and list of memrefs to be deallocated
- /// according to the function boundary ABI, i.e., assume ownership of all
- /// returned MemRefs.
+ /// called function is private, additional results are added for each MemRef
+ /// result to pass the dynamic ownership indicator along. Otherwise, updates
+ /// the ownership map and list of memrefs to be deallocated according to the
+ /// function boundary ABI, i.e., assume ownership of all returned MemRefs.
///
/// Example (assume `private-function-dynamic-ownership` is enabled):
/// ```
@@ -309,17 +308,15 @@ class BufferDeallocation {
/// becomes
/// ```
/// func.func @f(%arg0: memref<2xi32>) -> memref<2xi32> {...}
- /// func.func private @g(%arg0: memref<2xi32>) -> memref<2xi32> {...}
+ /// func.func private @g(%arg0: memref<2xi32>) -> (memref<2xi32>, i1) {...}
///
/// %ret_f = func.call @f(%memref) : (memref<2xi32>) -> memref<2xi32>
/// // set ownership(%ret_f) := true
/// // remember to deallocate %ret_f
///
- /// // (new_memref, own) = getmemrefWithUniqueOwnership(%memref)
- /// %ret_g:2 = func.call @g(new_memref, own) :
- /// (memref<2xi32>, i1) -> (memref<2xi32>, i1)
+ /// %ret_g:2 = func.call @g(%memref) : (memref<2xi32>) -> (memref<2xi32>, i1)
/// // set ownership(%ret_g#0) := %ret_g#1
- /// // remember to deallocate %ret_g
+ /// // remember to deallocate %ret_g if it comes with ownership
/// ```
FailureOr<Operation *> handleInterface(CallOpInterface op);
@@ -444,8 +441,8 @@ class BufferDeallocation {
static LogicalResult verifyOperationPreconditions(Operation *op);
/// When the 'private-function-dynamic-ownership' pass option is enabled,
- /// additional `i1` arguments and return values are added for each MemRef
- /// value in the function signature. This function takes care of updating the
+ /// additional `i1` return values are added for each MemRef result in the
+ /// function signature. This function takes care of updating the
/// `function_type` attribute of the function according to the actually
/// returned values from the terminators.
static LogicalResult updateFunctionSignature(FunctionOpInterface op);
@@ -650,7 +647,7 @@ LogicalResult BufferDeallocation::deallocate(Block *block) {
// Adhere to function boundary ABI: no ownership of function argument
// MemRefs is taken.
- if (isFunctionWithoutDynamicOwnership(block->getParentOp()) &&
+ if (isa<FunctionOpInterface>(block->getParentOp()) &&
block->isEntryBlock()) {
Value newArg = buildBoolValue(builder, arg.getLoc(), false);
state.updateOwnership(arg, newArg);
@@ -838,26 +835,10 @@ FailureOr<Operation *> BufferDeallocation::handleInterface(CallOpInterface op) {
isPrivate = symbol.isPrivate() && !symbol.isDeclaration();
// If the private-function-dynamic-ownership option is enabled and we are
- // calling a private function, we need to add an additional `i1`
- // argument/result for each MemRef argument/result to dynamically pass the
- // current ownership indicator rather than adhering to the function boundary
- // ABI.
+ // calling a private function, we need to add an additional `i1` result for
+ // each MemRef result to dynamically pass the current ownership indicator
+ // rather than adhering to the function boundary ABI.
if (options.privateFuncDynamicOwnership && isPrivate) {
- SmallVector<Value> newOperands, ownershipIndicatorsToAdd;
- for (Value operand : op.getArgOperands()) {
- if (!isMemref(operand)) {
- newOperands.push_back(operand);
- continue;
- }
- auto [memref, condition] =
- materializeUniqueOwnership(builder, operand, op->getBlock());
- newOperands.push_back(memref);
- ownershipIndicatorsToAdd.push_back(condition);
- }
- newOperands.append(ownershipIndicatorsToAdd.begin(),
- ownershipIndicatorsToAdd.end());
- op.getArgOperandsMutable().assign(newOperands);
-
unsigned numMemrefs = llvm::count_if(op->getResults(), isMemref);
SmallVector<Type> ownershipTypesToAppend(numMemrefs, builder.getI1Type());
unsigned ownershipCounter = op->getNumResults();
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir
index 02bf2d10e9e3f..63947150c23b3 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir
@@ -33,7 +33,7 @@ func.func @function_call() {
// CHECK-DYNAMIC-LABEL: func @function_call()
// CHECK-DYNAMIC: [[ALLOC0:%.+]] = memref.alloc(
// CHECK-DYNAMIC-NEXT: [[ALLOC1:%.+]] = memref.alloc(
-// CHECK-DYNAMIC-NEXT: [[RET:%.+]]:2 = call @f([[ALLOC0]], %true{{[0-9_]*}}) : (memref<f64>, i1) -> (memref<f64>, i1)
+// CHECK-DYNAMIC-NEXT: [[RET:%.+]]:2 = call @f([[ALLOC0]]) : (memref<f64>) -> (memref<f64>, i1)
// CHECK-DYNAMIC-NEXT: test.copy
// CHECK-DYNAMIC-NEXT: [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[RET]]#0
// CHECK-DYNAMIC-NEXT: bufferization.dealloc ([[ALLOC0]], [[ALLOC1]], [[BASE]] :{{.*}}) if (%true{{[0-9_]*}}, %true{{[0-9_]*}}, [[RET]]#1)
@@ -102,7 +102,7 @@ func.func @function_call_requries_merged_ownership_mid_block(%arg0: i1) {
// CHECK-DYNAMIC: [[ALLOC0:%.+]] = memref.alloc(
// CHECK-DYNAMIC-NEXT: [[ALLOC1:%.+]] = memref.alloca(
// CHECK-DYNAMIC-NEXT: [[SELECT:%.+]] = arith.select [[ARG0]], [[ALLOC0]], [[ALLOC1]]
-// CHECK-DYNAMIC-NEXT: [[RET:%.+]]:2 = call @f([[SELECT]], [[ARG0]])
+// CHECK-DYNAMIC-NEXT: [[RET:%.+]]:2 = call @f([[SELECT]])
// CHECK-DYNAMIC-NEXT: test.copy
// CHECK-DYNAMIC-NEXT: [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[RET]]#0
// CHECK-DYNAMIC-NEXT: bufferization.dealloc ([[ALLOC0]], [[BASE]] :
diff --git a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir
index de71f23cfad6f..a7735aa8ce11b 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir
@@ -24,11 +24,9 @@ func.func private @emptyUsesValue(%arg0: memref<4xf32>) {
// CHECK-NEXT: return
// CHECK-DYNAMIC-LABEL: func private @emptyUsesValue(
-// CHECK-DYNAMIC-SAME: [[ARG0:%.+]]: memref<4xf32>, [[ARG1:%.+]]: i1)
+// CHECK-DYNAMIC-SAME: [[ARG0:%.+]]: memref<4xf32>)
// CHECK-DYNAMIC: [[ALLOC:%.*]] = memref.alloc()
-// CHECK-DYNAMIC: [[BASE:%[a-zA-Z0-9_]+]], {{.*}} = memref.extract_strided_metadata [[ARG0]]
-// CHECK-DYNAMIC-NEXT: bufferization.dealloc ([[BASE]] :{{.*}}) if ([[ARG1]])
-// CHECK-DYNAMIC-NOT: retain
+// CHECK-DYNAMIC-NEXT: "test.read_buffer"
// CHECK-DYNAMIC-NEXT: bufferization.dealloc ([[ALLOC]] :{{.*}}) if (%true{{[0-9_]*}})
// CHECK-DYNAMIC-NOT: retain
// CHECK-DYNAMIC-NEXT: return
@@ -74,13 +72,11 @@ func.func private @redundantOperations(%arg0: memref<2xf32>) {
// CHECK-NEXT: return
// CHECK-DYNAMIC-LABEL: func private @redundantOperations
-// CHECK-DYNAMIC: (%[[ARG0:.*]]: memref{{.*}}, %[[ARG1:.*]]: i1)
+// CHECK-DYNAMIC: (%[[ARG0:.*]]: memref{{.*}})
// CHECK-DYNAMIC: %[[FIRST_ALLOC:.*]] = memref.alloc()
// CHECK-DYNAMIC-NEXT: test.buffer_based
// CHECK-DYNAMIC: %[[SECOND_ALLOC:.*]] = memref.alloc()
// CHECK-DYNAMIC-NEXT: test.buffer_based
-// CHECK-DYNAMIC-NEXT: %[[BASE:[a-zA-Z0-9_]+]], {{.*}} = memref.extract_strided_metadata %[[ARG0]]
-// CHECK-DYNAMIC-NEXT: bufferization.dealloc (%[[BASE]] : {{.*}}) if (%[[ARG1]])
// CHECK-DYNAMIC-NEXT: bufferization.dealloc (%[[FIRST_ALLOC]] : {{.*}}) if (%true{{[0-9_]*}})
// CHECK-DYNAMIC-NEXT: bufferization.dealloc (%[[SECOND_ALLOC]] : {{.*}}) if (%true{{[0-9_]*}})
// CHECK-DYNAMIC-NEXT: return
@@ -121,14 +117,39 @@ func.func private @memref_in_function_results(
// CHECK-DYNAMIC-LABEL: func private @memref_in_function_results
// CHECK-DYNAMIC: (%[[ARG0:.*]]: memref<5xf32>, %[[ARG1:.*]]: memref<10xf32>,
-// CHECK-DYNAMIC-SAME: %[[RESULT:.*]]: memref<5xf32>, %[[ARG3:.*]]: i1, %[[ARG4:.*]]: i1, %[[ARG5:.*]]: i1)
+// CHECK-DYNAMIC-SAME: %[[RESULT:.*]]: memref<5xf32>)
// CHECK-DYNAMIC: %[[X:.*]] = memref.alloc()
// CHECK-DYNAMIC: %[[Y:.*]] = memref.alloc()
// CHECK-DYNAMIC: test.copy
-// CHECK-DYNAMIC: %[[BASE0:[a-zA-Z0-9_]+]], {{.+}} = memref.extract_strided_metadata %[[ARG0]]
-// CHECK-DYNAMIC: %[[BASE1:[a-zA-Z0-9_]+]], {{.+}} = memref.extract_strided_metadata %[[RESULT]]
// CHECK-DYNAMIC: bufferization.dealloc (%[[Y]] : {{.*}}) if (%true{{[0-9_]*}})
// CHECK-DYNAMIC-NOT: retain
-// CHECK-DYNAMIC: [[OWN:%.+]] = bufferization.dealloc (%[[BASE0]], %[[BASE1]] : {{.*}}) if (%[[ARG3]], %[[ARG5]]) retain (%[[ARG1]] :
-// CHECK-DYNAMIC: [[OR:%.+]] = arith.ori [[OWN]], %[[ARG4]]
-// CHECK-DYNAMIC: return %[[ARG1]], %[[X]], [[OR]], %true
+// CHECK-DYNAMIC: return %[[ARG1]], %[[X]], %false, %true
+
+// -----
+
+// CHECK-DYNAMIC-LABEL: func private @private_callee(
+// CHECK-DYNAMIC-SAME: %[[arg0:.*]]: memref<f32>) -> (memref<f32>, i1)
+// CHECK-DYNAMIC: %[[true:.*]] = arith.constant true
+// CHECK-DYNAMIC: %[[alloc:.*]] = memref.alloc() : memref<f32>
+// CHECK-DYNAMIC-NOT: bufferization.dealloc
+// CHECK-DYNAMIC: return %[[alloc]], %[[true]]
+func.func private @private_callee(%arg0: memref<f32>) -> memref<f32> {
+ %alloc = memref.alloc() : memref<f32>
+ return %alloc : memref<f32>
+}
+
+// CHECK-DYNAMIC: func @caller() -> f32
+// CHECK-DYNAMIC: %[[true:.*]] = arith.constant true
+// CHECK-DYNAMIC: %[[alloc:.*]] = memref.alloc() : memref<f32>
+// CHECK-DYNAMIC: %[[call:.*]]:2 = call @private_callee(%[[alloc]])
+// CHECK-DYNAMIC: memref.load
+// CHECK-DYNAMIC: %[[base:.*]], %[[offset:.*]] = memref.extract_strided_metadata %[[call]]#0
+// CHECK-DYNAMIC: bufferization.dealloc (%[[alloc]], %[[base]] : {{.*}}) if (%[[true]], %[[call]]#1)
+// CHECK-DYNAMIC-NOT: retain
+func.func @caller() -> (f32) {
+ %alloc = memref.alloc() : memref<f32>
+ %ret = call @private_callee(%alloc) : (memref<f32>) -> memref<f32>
+
+ %val = memref.load %ret[] : memref<f32>
+ return %val : f32
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/80655
More information about the Mlir-commits
mailing list