Skip to content

Commit 44223f2

Browse files
matthias-springermatthias-springer
committed
[mlir][bufferization] Never pass ownership to functions
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.
1 parent 3e230bb commit 44223f2

File tree

5 files changed

+74
-80
lines changed

5 files changed

+74
-80
lines changed

‎mlir/docs/Bufferization.md

Lines changed: 22 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -502,31 +502,26 @@ accordingly:
502502

503503
### Example
504504

505-
The following example contains a few interesting cases:
506-
* Basic block arguments are modified to also pass along the ownership
507-
indicator, but not for entry bocks of non-private functions (assuming the
508-
`private-function-dynamic-ownership` pass option is disabled) where the
509-
function boundary ABI is applied instead. "Private" in this context refers
510-
to functions that cannot be called externally.
511-
* The result of `arith.select` initially has 'Unknown' assigned as ownership,
512-
but once the `bufferization.dealloc` operation is inserted it is put in the
513-
'retained' list (since it has uses in a later basic block) and thus the
514-
'Unknown' ownership can be replaced with a 'Unique' ownership using the
515-
corresponding result of the dealloc operation.
516-
* The `cf.cond_br` operation has more than one successor and thus has to
517-
insert two `bufferization.dealloc` operations (one for each successor).
518-
While they have the same list of MemRefs to deallocate (because they perform
519-
the deallocations for the same block), it must be taken into account that
520-
some MemRefs remain *live* for one branch but not the other (thus set
521-
intersection is performed on the *live-out* of the current block and the
522-
*live-in* of the target block). Also, `cf.cond_br` supports separate
523-
forwarding operands for each successor. To make sure that no MemRef is
524-
deallocated twice (because there are two `bufferization.dealloc` operations
525-
with the same MemRefs to deallocate), the condition operands are adjusted to
526-
take the branch condition into account. While a generic lowering for such
527-
terminator operations could be implemented, a specialized implementation can
528-
take all the semantics of this particular operation into account and thus
529-
generate a more efficient lowering.
505+
The following example contains a few interesting cases: * Basic block arguments
506+
are modified to also pass along the ownership indicator, but not for entry bocks
507+
where the function boundary ABI is applied instead. * The result of
508+
`arith.select` initially has 'Unknown' assigned as ownership, but once the
509+
`bufferization.dealloc` operation is inserted it is put in the 'retained' list
510+
(since it has uses in a later basic block) and thus the 'Unknown' ownership can
511+
be replaced with a 'Unique' ownership using the corresponding result of the
512+
dealloc operation. * The `cf.cond_br` operation has more than one successor and
513+
thus has to insert two `bufferization.dealloc` operations (one for each
514+
successor). While they have the same list of MemRefs to deallocate (because they
515+
perform the deallocations for the same block), it must be taken into account
516+
that some MemRefs remain *live* for one branch but not the other (thus set
517+
intersection is performed on the *live-out* of the current block and the
518+
*live-in* of the target block). Also, `cf.cond_br` supports separate forwarding
519+
operands for each successor. To make sure that no MemRef is deallocated twice
520+
(because there are two `bufferization.dealloc` operations with the same MemRefs
521+
to deallocate), the condition operands are adjusted to take the branch condition
522+
into account. While a generic lowering for such terminator operations could be
523+
implemented, a specialized implementation can take all the semantics of this
524+
particular operation into account and thus generate a more efficient lowering.
530525

531526
```mlir
532527
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) {
543538
After running `--ownership-based-buffer-deallocation`, it looks as follows:
544539

545540
```mlir
546-
// Since this is not a private function, the signature will not be modified even
547-
// when private-function-dynamic-ownership is enabled. Instead the function
548-
// boundary ABI has to be applied which means that ownership of `%memref` will
549-
// never be acquired.
541+
// Function boundary ABI: ownership of `%memref` will never be acquired.
550542
func.func @example(%memref: memref<?xi8>, %select_cond: i1, %br_cond: i1) {
551543
%false = arith.constant false
552544
%true = arith.constant true
@@ -602,7 +594,7 @@ func.func @example(%memref: memref<?xi8>, %select_cond: i1, %br_cond: i1) {
602594
: memref<i8>, memref<i8>, memref<i8>)
603595
if (%false, %not_br_cond, %false)
604596
retain (%memref, %select : memref<?xi8>, memref<?xi8>)
605-
597+
606598
// Because %select is used in ^bb1 without passing it via block argument, we
607599
// need to update it's ownership value here by merging the ownership values
608600
// returned by the dealloc operations

‎mlir/include/mlir/Dialect/Bufferization/Pipelines/Passes.h

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,9 +24,9 @@ struct BufferDeallocationPipelineOptions
2424
PassOptions::Option<bool> privateFunctionDynamicOwnership{
2525
*this, "private-function-dynamic-ownership",
2626
llvm::cl::desc(
27-
"Allows to add additional arguments to private functions to "
28-
"dynamically pass ownership of memrefs to callees. This can enable "
29-
"earlier deallocations."),
27+
"Allows to add additional results to private functions to return "
28+
"ownership of returned memrefs to callers. This can avoid spurious "
29+
"buffer clones in the callee."),
3030
llvm::cl::init(false)};
3131
};
3232

‎mlir/lib/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation.cpp

Lines changed: 13 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -292,11 +292,10 @@ class BufferDeallocation {
292292
FailureOr<Operation *> handleInterface(RegionBranchOpInterface op);
293293

294294
/// If the private-function-dynamic-ownership pass option is enabled and the
295-
/// called function is private, additional arguments and results are added for
296-
/// each MemRef argument/result to pass the dynamic ownership indicator along.
297-
/// Otherwise, updates the ownership map and list of memrefs to be deallocated
298-
/// according to the function boundary ABI, i.e., assume ownership of all
299-
/// returned MemRefs.
295+
/// called function is private, additional results are added for each MemRef
296+
/// result to pass the dynamic ownership indicator along. Otherwise, updates
297+
/// the ownership map and list of memrefs to be deallocated according to the
298+
/// function boundary ABI, i.e., assume ownership of all returned MemRefs.
300299
///
301300
/// Example (assume `private-function-dynamic-ownership` is enabled):
302301
/// ```
@@ -309,17 +308,15 @@ class BufferDeallocation {
309308
/// becomes
310309
/// ```
311310
/// func.func @f(%arg0: memref<2xi32>) -> memref<2xi32> {...}
312-
/// func.func private @g(%arg0: memref<2xi32>) -> memref<2xi32> {...}
311+
/// func.func private @g(%arg0: memref<2xi32>) -> (memref<2xi32>, i1) {...}
313312
///
314313
/// %ret_f = func.call @f(%memref) : (memref<2xi32>) -> memref<2xi32>
315314
/// // set ownership(%ret_f) := true
316315
/// // remember to deallocate %ret_f
317316
///
318-
/// // (new_memref, own) = getmemrefWithUniqueOwnership(%memref)
319-
/// %ret_g:2 = func.call @g(new_memref, own) :
320-
/// (memref<2xi32>, i1) -> (memref<2xi32>, i1)
317+
/// %ret_g:2 = func.call @g(%memref) : (memref<2xi32>) -> (memref<2xi32>, i1)
321318
/// // set ownership(%ret_g#0) := %ret_g#1
322-
/// // remember to deallocate %ret_g
319+
/// // remember to deallocate %ret_g if it comes with ownership
323320
/// ```
324321
FailureOr<Operation *> handleInterface(CallOpInterface op);
325322

@@ -444,8 +441,8 @@ class BufferDeallocation {
444441
static LogicalResult verifyOperationPreconditions(Operation *op);
445442

446443
/// When the 'private-function-dynamic-ownership' pass option is enabled,
447-
/// additional `i1` arguments and return values are added for each MemRef
448-
/// value in the function signature. This function takes care of updating the
444+
/// additional `i1` return values are added for each MemRef result in the
445+
/// function signature. This function takes care of updating the
449446
/// `function_type` attribute of the function according to the actually
450447
/// returned values from the terminators.
451448
static LogicalResult updateFunctionSignature(FunctionOpInterface op);
@@ -650,7 +647,7 @@ LogicalResult BufferDeallocation::deallocate(Block *block) {
650647

651648
// Adhere to function boundary ABI: no ownership of function argument
652649
// MemRefs is taken.
653-
if (isFunctionWithoutDynamicOwnership(block->getParentOp()) &&
650+
if (isa<FunctionOpInterface>(block->getParentOp()) &&
654651
block->isEntryBlock()) {
655652
Value newArg = buildBoolValue(builder, arg.getLoc(), false);
656653
state.updateOwnership(arg, newArg);
@@ -838,26 +835,10 @@ FailureOr<Operation *> BufferDeallocation::handleInterface(CallOpInterface op) {
838835
isPrivate = symbol.isPrivate() && !symbol.isDeclaration();
839836

840837
// If the private-function-dynamic-ownership option is enabled and we are
841-
// calling a private function, we need to add an additional `i1`
842-
// argument/result for each MemRef argument/result to dynamically pass the
843-
// current ownership indicator rather than adhering to the function boundary
844-
// ABI.
838+
// calling a private function, we need to add an additional `i1` result for
839+
// each MemRef result to dynamically pass the current ownership indicator
840+
// rather than adhering to the function boundary ABI.
845841
if (options.privateFuncDynamicOwnership && isPrivate) {
846-
SmallVector<Value> newOperands, ownershipIndicatorsToAdd;
847-
for (Value operand : op.getArgOperands()) {
848-
if (!isMemref(operand)) {
849-
newOperands.push_back(operand);
850-
continue;
851-
}
852-
auto [memref, condition] =
853-
materializeUniqueOwnership(builder, operand, op->getBlock());
854-
newOperands.push_back(memref);
855-
ownershipIndicatorsToAdd.push_back(condition);
856-
}
857-
newOperands.append(ownershipIndicatorsToAdd.begin(),
858-
ownershipIndicatorsToAdd.end());
859-
op.getArgOperandsMutable().assign(newOperands);
860-
861842
unsigned numMemrefs = llvm::count_if(op->getResults(), isMemref);
862843
SmallVector<Type> ownershipTypesToAppend(numMemrefs, builder.getI1Type());
863844
unsigned ownershipCounter = op->getNumResults();

‎mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-callop-interface.mlir

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ func.func @function_call() {
3333
// CHECK-DYNAMIC-LABEL: func @function_call()
3434
// CHECK-DYNAMIC: [[ALLOC0:%.+]] = memref.alloc(
3535
// CHECK-DYNAMIC-NEXT: [[ALLOC1:%.+]] = memref.alloc(
36-
// CHECK-DYNAMIC-NEXT: [[RET:%.+]]:2 = call @f([[ALLOC0]], %true{{[0-9_]*}}) : (memref<f64>, i1) -> (memref<f64>, i1)
36+
// CHECK-DYNAMIC-NEXT: [[RET:%.+]]:2 = call @f([[ALLOC0]]) : (memref<f64>) -> (memref<f64>, i1)
3737
// CHECK-DYNAMIC-NEXT: test.copy
3838
// CHECK-DYNAMIC-NEXT: [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[RET]]#0
3939
// 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) {
102102
// CHECK-DYNAMIC: [[ALLOC0:%.+]] = memref.alloc(
103103
// CHECK-DYNAMIC-NEXT: [[ALLOC1:%.+]] = memref.alloca(
104104
// CHECK-DYNAMIC-NEXT: [[SELECT:%.+]] = arith.select [[ARG0]], [[ALLOC0]], [[ALLOC1]]
105-
// CHECK-DYNAMIC-NEXT: [[RET:%.+]]:2 = call @f([[SELECT]], [[ARG0]])
105+
// CHECK-DYNAMIC-NEXT: [[RET:%.+]]:2 = call @f([[SELECT]])
106106
// CHECK-DYNAMIC-NEXT: test.copy
107107
// CHECK-DYNAMIC-NEXT: [[BASE:%[a-zA-Z0-9_]+]]{{.*}} = memref.extract_strided_metadata [[RET]]#0
108108
// CHECK-DYNAMIC-NEXT: bufferization.dealloc ([[ALLOC0]], [[BASE]] :

‎mlir/test/Dialect/Bufferization/Transforms/OwnershipBasedBufferDeallocation/dealloc-function-boundaries.mlir

Lines changed: 34 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -24,11 +24,9 @@ func.func private @emptyUsesValue(%arg0: memref<4xf32>) {
2424
// CHECK-NEXT: return
2525

2626
// CHECK-DYNAMIC-LABEL: func private @emptyUsesValue(
27-
// CHECK-DYNAMIC-SAME: [[ARG0:%.+]]: memref<4xf32>, [[ARG1:%.+]]: i1)
27+
// CHECK-DYNAMIC-SAME: [[ARG0:%.+]]: memref<4xf32>)
2828
// CHECK-DYNAMIC: [[ALLOC:%.*]] = memref.alloc()
29-
// CHECK-DYNAMIC: [[BASE:%[a-zA-Z0-9_]+]], {{.*}} = memref.extract_strided_metadata [[ARG0]]
30-
// CHECK-DYNAMIC-NEXT: bufferization.dealloc ([[BASE]] :{{.*}}) if ([[ARG1]])
31-
// CHECK-DYNAMIC-NOT: retain
29+
// CHECK-DYNAMIC-NEXT: "test.read_buffer"
3230
// CHECK-DYNAMIC-NEXT: bufferization.dealloc ([[ALLOC]] :{{.*}}) if (%true{{[0-9_]*}})
3331
// CHECK-DYNAMIC-NOT: retain
3432
// CHECK-DYNAMIC-NEXT: return
@@ -74,13 +72,11 @@ func.func private @redundantOperations(%arg0: memref<2xf32>) {
7472
// CHECK-NEXT: return
7573

7674
// CHECK-DYNAMIC-LABEL: func private @redundantOperations
77-
// CHECK-DYNAMIC: (%[[ARG0:.*]]: memref{{.*}}, %[[ARG1:.*]]: i1)
75+
// CHECK-DYNAMIC: (%[[ARG0:.*]]: memref{{.*}})
7876
// CHECK-DYNAMIC: %[[FIRST_ALLOC:.*]] = memref.alloc()
7977
// CHECK-DYNAMIC-NEXT: test.buffer_based
8078
// CHECK-DYNAMIC: %[[SECOND_ALLOC:.*]] = memref.alloc()
8179
// CHECK-DYNAMIC-NEXT: test.buffer_based
82-
// CHECK-DYNAMIC-NEXT: %[[BASE:[a-zA-Z0-9_]+]], {{.*}} = memref.extract_strided_metadata %[[ARG0]]
83-
// CHECK-DYNAMIC-NEXT: bufferization.dealloc (%[[BASE]] : {{.*}}) if (%[[ARG1]])
8480
// CHECK-DYNAMIC-NEXT: bufferization.dealloc (%[[FIRST_ALLOC]] : {{.*}}) if (%true{{[0-9_]*}})
8581
// CHECK-DYNAMIC-NEXT: bufferization.dealloc (%[[SECOND_ALLOC]] : {{.*}}) if (%true{{[0-9_]*}})
8682
// CHECK-DYNAMIC-NEXT: return
@@ -121,14 +117,39 @@ func.func private @memref_in_function_results(
121117

122118
// CHECK-DYNAMIC-LABEL: func private @memref_in_function_results
123119
// CHECK-DYNAMIC: (%[[ARG0:.*]]: memref<5xf32>, %[[ARG1:.*]]: memref<10xf32>,
124-
// CHECK-DYNAMIC-SAME: %[[RESULT:.*]]: memref<5xf32>, %[[ARG3:.*]]: i1, %[[ARG4:.*]]: i1, %[[ARG5:.*]]: i1)
120+
// CHECK-DYNAMIC-SAME: %[[RESULT:.*]]: memref<5xf32>)
125121
// CHECK-DYNAMIC: %[[X:.*]] = memref.alloc()
126122
// CHECK-DYNAMIC: %[[Y:.*]] = memref.alloc()
127123
// CHECK-DYNAMIC: test.copy
128-
// CHECK-DYNAMIC: %[[BASE0:[a-zA-Z0-9_]+]], {{.+}} = memref.extract_strided_metadata %[[ARG0]]
129-
// CHECK-DYNAMIC: %[[BASE1:[a-zA-Z0-9_]+]], {{.+}} = memref.extract_strided_metadata %[[RESULT]]
130124
// CHECK-DYNAMIC: bufferization.dealloc (%[[Y]] : {{.*}}) if (%true{{[0-9_]*}})
131125
// CHECK-DYNAMIC-NOT: retain
132-
// CHECK-DYNAMIC: [[OWN:%.+]] = bufferization.dealloc (%[[BASE0]], %[[BASE1]] : {{.*}}) if (%[[ARG3]], %[[ARG5]]) retain (%[[ARG1]] :
133-
// CHECK-DYNAMIC: [[OR:%.+]] = arith.ori [[OWN]], %[[ARG4]]
134-
// CHECK-DYNAMIC: return %[[ARG1]], %[[X]], [[OR]], %true
126+
// CHECK-DYNAMIC: return %[[ARG1]], %[[X]], %false, %true
127+
128+
// -----
129+
130+
// CHECK-DYNAMIC-LABEL: func private @private_callee(
131+
// CHECK-DYNAMIC-SAME: %[[arg0:.*]]: memref<f32>) -> (memref<f32>, i1)
132+
// CHECK-DYNAMIC: %[[true:.*]] = arith.constant true
133+
// CHECK-DYNAMIC: %[[alloc:.*]] = memref.alloc() : memref<f32>
134+
// CHECK-DYNAMIC-NOT: bufferization.dealloc
135+
// CHECK-DYNAMIC: return %[[alloc]], %[[true]]
136+
func.func private @private_callee(%arg0: memref<f32>) -> memref<f32> {
137+
%alloc = memref.alloc() : memref<f32>
138+
return %alloc : memref<f32>
139+
}
140+
141+
// CHECK-DYNAMIC: func @caller() -> f32
142+
// CHECK-DYNAMIC: %[[true:.*]] = arith.constant true
143+
// CHECK-DYNAMIC: %[[alloc:.*]] = memref.alloc() : memref<f32>
144+
// CHECK-DYNAMIC: %[[call:.*]]:2 = call @private_callee(%[[alloc]])
145+
// CHECK-DYNAMIC: memref.load
146+
// CHECK-DYNAMIC: %[[base:.*]], %[[offset:.*]] = memref.extract_strided_metadata %[[call]]#0
147+
// CHECK-DYNAMIC: bufferization.dealloc (%[[alloc]], %[[base]] : {{.*}}) if (%[[true]], %[[call]]#1)
148+
// CHECK-DYNAMIC-NOT: retain
149+
func.func @caller() -> (f32) {
150+
%alloc = memref.alloc() : memref<f32>
151+
%ret = call @private_callee(%alloc) : (memref<f32>) -> memref<f32>
152+
153+
%val = memref.load %ret[] : memref<f32>
154+
return %val : f32
155+
}

0 commit comments

Comments
 (0)