[mlir][bufferization] Never pass ownership to functions #80655
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-bufferization Author: Matthias Springer (matthias-springer) ChangesEven when 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:
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
+}
|
maerhart
approved these changes
Feb 5, 2024
Even when `private-function-dynamic-ownership` is set, ownership should never be passed to the callee. This can lead to double deallocs (llvm#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 llvm#77096.
matthias-springer
force-pushed
the
bufferization_priv_func_double_dealloc
branch
from
44223f2 to
5bed6d0
Compare
agozillon
pushed a commit
to agozillon/llvm-project
that referenced
this pull request
Feb 5, 2024
Even when `private-function-dynamic-ownership` is set, ownership should never be passed to the callee. This can lead to double deallocs (llvm#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 llvm#77096.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Even when
private-function-dynamic-ownershipis 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.