213917be82
The `-buffer-deallocation` pass is not compatible with One-Shot Bufferize and has been replaced with the Ownership-based Buffer Deallocation pass about 1.5 years ago. To clean up the code base, this commit removes the deprecated `buffer-deallocation` pass. All uses of this deprecated pass within MLIR have already been migrated. Note for LLVM integration: If you depend on this pass, migrate to the Ownership-based Buffer Deallocation pass or copy the pass to your codebase. For details, see https://discourse.llvm.org/t/psa-bufferization-new-buffer-deallocation-pipeline/73375.
618 lines
29 KiB
Markdown
618 lines
29 KiB
Markdown
# Ownership-based Buffer Deallocation
|
|
|
|
[TOC]
|
|
|
|
One-Shot Bufferize does not deallocate any buffers that it allocates. After
|
|
running One-Shot Bufferize, the resulting IR may have a number of `memref.alloc`
|
|
ops, but no `memref.dealloc` ops. Buffer dellocation is delegated to the
|
|
`-ownership-based-buffer-deallocation` pass.
|
|
|
|
On a high level, buffers are "owned" by a basic block. Ownership materializes
|
|
as an `i1` SSA value and can be thought of as "responsibility to deallocate". It
|
|
is conceptually similar to `std::unique_ptr` in C++.
|
|
|
|
There are few additional preprocessing and postprocessing passes that should be
|
|
run together with the ownership-based buffer deallocation pass. The recommended
|
|
compilation pipeline is as follows:
|
|
|
|
```
|
|
one-shot-bufferize
|
|
| it's recommended to perform all bufferization here at latest,
|
|
| <- any allocations inserted after this point have to be handled
|
|
V manually
|
|
expand-realloc
|
|
V
|
|
ownership-based-buffer-deallocation
|
|
V
|
|
canonicalize <- mostly for scf.if simplifications
|
|
V
|
|
buffer-deallocation-simplification
|
|
V <- from this point onwards no tensor values are allowed
|
|
lower-deallocations
|
|
V
|
|
CSE
|
|
V
|
|
canonicalize
|
|
```
|
|
|
|
The entire deallocation pipeline (excluding `-one-shot-bufferize`) is exposed
|
|
as `-buffer-deallocation-pipeline`.
|
|
|
|
The ownership-based buffer deallocation pass processes operations implementing
|
|
`FunctionOpInterface` one-by-one without analysing the call-graph.
|
|
This means that there have to be [some rules](#function-boundary-abi) on how
|
|
MemRefs are handled when being passed from one function to another. The rest of
|
|
the pass revolves heavily around the `bufferization.dealloc` operation which is
|
|
inserted at the end of each basic block with appropriate operands and should be
|
|
optimized using the Buffer Deallocation Simplification pass
|
|
(`--buffer-deallocation-simplification`) and the regular canonicalizer
|
|
(`--canonicalize`). Lowering the result of the
|
|
`-ownership-based-buffer-deallocation` pass directly using
|
|
`--convert-bufferization-to-memref` without beforehand optimization is not
|
|
recommended as it will lead to very inefficient code (the runtime-cost of
|
|
`bufferization.dealloc` is `O(|memrefs|^2+|memref|*|retained|)`).
|
|
|
|
## Function boundary ABI
|
|
|
|
The Buffer Deallocation pass operates on the level of operations implementing
|
|
the `FunctionOpInterface`. Such operations can take MemRefs as arguments, but
|
|
also return them. To ensure compatibility among all functions (including
|
|
external ones), some rules have to be enforced:
|
|
* When a MemRef is passed as a function argument, ownership is never acquired.
|
|
It is always the caller's responsibility to deallocate such MemRefs.
|
|
* Returning a MemRef from a function always passes ownership to the caller,
|
|
i.e., it is also the caller's responsibility to deallocate memrefs returned
|
|
from a called function.
|
|
* A function must not return a MemRef with the same allocated base buffer as
|
|
one of its arguments (in this case a copy has to be created). Note that in
|
|
this context two subviews of the same buffer that don't overlap are also
|
|
considered to alias.
|
|
|
|
For external functions (e.g., library functions written externally in C), the
|
|
externally provided implementation has to adhere to these rules and they are
|
|
just assumed by the buffer deallocation pass. Functions on which the
|
|
deallocation pass is applied and for which the implementation is accessible are
|
|
modified by the pass such that the ABI is respected (i.e., buffer copies are
|
|
inserted when necessary).
|
|
|
|
## Inserting `bufferization.dealloc` operations
|
|
|
|
`bufferization.dealloc` and ownership indicators are the main abstractions in
|
|
the ownership-based buffer deallocation pass. `bufferization.dealloc`
|
|
deallocates all given buffers if the respective ownership indicator is set and
|
|
there is no aliasing buffer in the retain list.
|
|
|
|
|
|
|
|
`bufferization.dealloc` operations are unconditionally inserted at the end of
|
|
each basic block (just before the terminator). The majority of the pass is about
|
|
finding the correct operands for this operation. There are three variadic
|
|
operand lists to be populated, the first contains all MemRef values that may
|
|
need to be deallocated, the second list contains their associated ownership
|
|
values (of `i1` type), and the third list contains MemRef values that are still
|
|
needed at a later point and should thus not be deallocated (e.g., yielded or
|
|
returned buffers).
|
|
|
|
`bufferization.dealloc` allows us to deal with any kind of aliasing behavior: it
|
|
lowers to runtime aliasing checks when not enough information can be collected
|
|
statically. When enough aliasing information is statically available, operands
|
|
or the entire op may fold away.
|
|
|
|
**Ownerships**
|
|
|
|
To do so, we use a concept of ownership indicators of memrefs which materialize
|
|
as an `i1` value for any SSA value of `memref` type, indicating whether the
|
|
basic block in which it was materialized has ownership of this MemRef. Ideally,
|
|
this is a constant `true` or `false`, but might also be a non-constant SSA
|
|
value. To keep track of those ownership values without immediately materializing
|
|
them (which might require insertion of `bufferization.clone` operations or
|
|
operations checking for aliasing at runtime at positions where we don't actually
|
|
need a materialized value), we use the `Ownership` class. This class represents
|
|
the ownership in three states forming a lattice on a partial order:
|
|
```
|
|
forall X in SSA values. uninitialized < unique(X) < unknown
|
|
forall X, Y in SSA values.
|
|
unique(X) == unique(Y) iff X and Y always evaluate to the same value
|
|
unique(X) != unique(Y) otherwise
|
|
```
|
|
Intuitively, the states have the following meaning:
|
|
* Uninitialized: the ownership is not initialized yet, this is the default
|
|
state; once an operation is finished processing the ownership of all
|
|
operation results with MemRef type should not be uninitialized anymore.
|
|
* Unique: there is a specific SSA value that can be queried to check ownership
|
|
without materializing any additional IR
|
|
* Unknown: no specific SSA value is available without materializing additional
|
|
IR, typically this is because two ownerships in 'Unique' state would have to
|
|
be merged manually (e.g., the result of an `arith.select` either has the
|
|
ownership of the then or else case depending on the condition value,
|
|
inserting another `arith.select` for the ownership values can perform the
|
|
merge and provide a 'Unique' ownership for the result), however, in the
|
|
general case this 'Unknown' state has to be assigned.
|
|
|
|
Implied by the above partial order, the pass combines two ownerships in the
|
|
following way:
|
|
|
|
| Ownership 1 | Ownership 2 | Combined Ownership |
|
|
|:--------------|:--------------|:-------------------|
|
|
| uninitialized | uninitialized | uninitialized |
|
|
| unique(X) | uninitialized | unique(X) |
|
|
| unique(X) | unique(X) | unique(X) |
|
|
| unique(X) | unique(Y) | unknown |
|
|
| unknown | unique | unknown |
|
|
| unknown | uninitialized | unknown |
|
|
| <td colspan=3> + symmetric cases |
|
|
|
|
**Collecting the list of MemRefs that potentially need to be deallocated**
|
|
|
|
For a given block, the list of MemRefs that potentially need to be deallocated
|
|
at the end of that block is computed by keeping track of all values for which
|
|
the block potentially takes over ownership. This includes MemRefs provided as
|
|
basic block arguments, interface handlers for operations like `memref.alloc` and
|
|
`func.call`, but also liveness information in regions with multiple basic
|
|
blocks. More concretely, it is computed by taking the MemRefs in the 'in' set
|
|
of the liveness analysis of the current basic block B, appended by the MemRef
|
|
block arguments and by the set of MemRefs allocated in B itself (determined by
|
|
the interface handlers), then subtracted (also determined by the interface
|
|
handlers) by the set of MemRefs deallocated in B.
|
|
|
|
Note that we don't have to take the intersection of the liveness 'in' set with
|
|
the 'out' set of the predecessor block because a value that is in the 'in' set
|
|
must be defined in an ancestor block that dominates all direct predecessors and
|
|
thus the 'in' set of this block is a subset of the 'out' sets of each
|
|
predecessor.
|
|
|
|
```
|
|
memrefs = filter((liveIn(block) U
|
|
allocated(block) U arguments(block)) \ deallocated(block), isMemRef)
|
|
```
|
|
|
|
The list of conditions for the second variadic operands list of
|
|
`bufferization.dealloc` is computed by querying the stored ownership value for
|
|
each of the MemRefs collected as described above. The ownership state is updated
|
|
by the interface handlers while processing the basic block.
|
|
|
|
**Collecting the list of MemRefs to retain**
|
|
|
|
Given a basic block B, the list of MemRefs that have to be retained can be
|
|
different for each successor block S. For the two basic blocks B and S and the
|
|
values passed via block arguments to the destination block S, we compute the
|
|
list of MemRefs that have to be retained in B by taking the MemRefs in the
|
|
successor operand list of the terminator and the MemRefs in the 'out' set of the
|
|
liveness analysis for B intersected with the 'in' set of the destination block
|
|
S.
|
|
|
|
This list of retained values makes sure that we cannot run into use-after-free
|
|
situations even if no aliasing information is present at compile-time.
|
|
|
|
```
|
|
toRetain = filter(successorOperands + (liveOut(fromBlock) insersect
|
|
liveIn(toBlock)), isMemRef)
|
|
```
|
|
|
|
## Supported interfaces
|
|
|
|
The pass uses liveness analysis and a few interfaces:
|
|
* `FunctionOpInterface`
|
|
* `CallOpInterface`
|
|
* `MemoryEffectOpInterface`
|
|
* `RegionBranchOpInterface`
|
|
* `RegionBranchTerminatorOpInterface`
|
|
|
|
Due to insufficient information provided by the interface, it also special-cases
|
|
on the `cf.cond_br` operation and makes some assumptions about operations
|
|
implementing the `RegionBranchOpInterface` at the moment, but improving the
|
|
interfaces would allow us to remove those dependencies in the future.
|
|
|
|
## Limitations
|
|
|
|
The Buffer Deallocation pass has some requirements and limitations on the input
|
|
IR. These are checked in the beginning of the pass and errors are emitted
|
|
accordingly:
|
|
* The set of interfaces the pass operates on must be implemented (correctly).
|
|
E.g., if there is an operation present with a nested region, but does not
|
|
implement the `RegionBranchOpInterface`, an error is emitted because the
|
|
pass cannot know the semantics of the nested region (and does not make any
|
|
default assumptions on it).
|
|
* No explicit control-flow loops are present. Currently, only loops using
|
|
structural-control-flow are supported. However, this limitation could be
|
|
lifted in the future.
|
|
* Deallocation operations should not be present already. The pass should
|
|
handle them correctly already (at least in most cases), but it's not
|
|
supported yet due to insufficient testing.
|
|
* Terminators must implement either `RegionBranchTerminatorOpInterface` or
|
|
`BranchOpInterface`, but not both. Terminators with more than one successor
|
|
are not supported (except `cf.cond_br`). This is not a fundamental
|
|
limitation, but there is no use-case justifying the more complex
|
|
implementation at the moment.
|
|
|
|
## 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 blocks, 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) {
|
|
%alloc = memref.alloc() : memref<?xi8>
|
|
%alloca = memref.alloca() : memref<?xi8>
|
|
%select = arith.select %select_cond, %alloc, %alloca : memref<?xi8>
|
|
cf.cond_br %br_cond, ^bb1(%alloc : memref<?xi8>), ^bb1(%memref : memref<?xi8>)
|
|
^bb1(%bbarg: memref<?xi8>):
|
|
test.copy(%bbarg, %select) : (memref<?xi8>, memref<?xi8>)
|
|
return
|
|
}
|
|
```
|
|
|
|
After running `--ownership-based-buffer-deallocation`, it looks as follows:
|
|
|
|
```mlir
|
|
// 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
|
|
|
|
// The ownership of a MemRef defined by the `memref.alloc` operation is always
|
|
// assigned to be 'true'.
|
|
%alloc = memref.alloc() : memref<?xi8>
|
|
|
|
// The ownership of a MemRef defined by the `memref.alloca` operation is
|
|
// always assigned to be 'false'.
|
|
%alloca = memref.alloca() : memref<?xi8>
|
|
|
|
// The ownership of %select will be the join of the ownership of %alloc and
|
|
// the ownership of %alloca, i.e., of %true and %false. Because the pass does
|
|
// not know about the semantics of the `arith.select` operation (unless a
|
|
// custom handler is implemented), the ownership join will be 'Unknown'. If
|
|
// the materialized ownership indicator of %select is needed, either a clone
|
|
// has to be created for which %true is assigned as ownership or the result
|
|
// of a `bufferization.dealloc` where %select is in the retain list has to be
|
|
// used.
|
|
%select = arith.select %select_cond, %alloc, %alloca : memref<?xi8>
|
|
|
|
// We use `memref.extract_strided_metadata` to get the base memref since it is
|
|
// not allowed to pass arbitrary memrefs to `memref.dealloc`. This property is
|
|
// already enforced for `bufferization.dealloc`
|
|
%base_buffer_memref, ... = memref.extract_strided_metadata %memref
|
|
: memref<?xi8> -> memref<i8>, index, index, index
|
|
%base_buffer_alloc, ... = memref.extract_strided_metadata %alloc
|
|
: memref<?xi8> -> memref<i8>, index, index, index
|
|
%base_buffer_alloca, ... = memref.extract_strided_metadata %alloca
|
|
: memref<?xi8> -> memref<i8>, index, index, index
|
|
|
|
// The deallocation conditions need to be adjusted to incorporate the branch
|
|
// condition. In this example, this requires only a single negation, but might
|
|
// also require multiple arith.andi operations.
|
|
%not_br_cond = arith.xori %true, %br_cond : i1
|
|
|
|
// There are two dealloc operations inserted in this basic block, one per
|
|
// successor. Both have the same list of MemRefs to deallocate and the
|
|
// conditions only differ by the branch condition conjunct.
|
|
// Note, however, that the retained list differs. Here, both contain the
|
|
// %select value because it is used in both successors (since it's the same
|
|
// block), but the value passed via block argument differs (%memref vs.
|
|
// %alloc).
|
|
%10:2 = bufferization.dealloc
|
|
(%base_buffer_memref, %base_buffer_alloc, %base_buffer_alloca
|
|
: memref<i8>, memref<i8>, memref<i8>)
|
|
if (%false, %br_cond, %false)
|
|
retain (%alloc, %select : memref<?xi8>, memref<?xi8>)
|
|
|
|
%11:2 = bufferization.dealloc
|
|
(%base_buffer_memref, %base_buffer_alloc, %base_buffer_alloca
|
|
: 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
|
|
%new_ownership = arith.select %br_cond, %10#1, %11#1 : i1
|
|
|
|
// The terminator is modified to pass along the ownership indicator values
|
|
// with each MemRef value.
|
|
cf.cond_br %br_cond, ^bb1(%alloc, %10#0 : memref<?xi8>, i1),
|
|
^bb1(%memref, %11#0 : memref<?xi8>, i1)
|
|
|
|
// All non-entry basic blocks are modified to have an additional i1 argument for
|
|
// each MemRef value in the argument list.
|
|
^bb1(%13: memref<?xi8>, %14: i1): // 2 preds: ^bb0, ^bb0
|
|
test.copy(%13, %select) : (memref<?xi8>, memref<?xi8>)
|
|
|
|
%base_buffer_13, ... = memref.extract_strided_metadata %13
|
|
: memref<?xi8> -> memref<i8>, index, index, index
|
|
%base_buffer_select, ... = memref.extract_strided_metadata %select
|
|
: memref<?xi8> -> memref<i8>, index, index, index
|
|
|
|
// Here, we don't have a retained list, because the block has no successors
|
|
// and the return has no operands.
|
|
bufferization.dealloc (%base_buffer_13, %base_buffer_select
|
|
: memref<i8>, memref<i8>)
|
|
if (%14, %new_ownership)
|
|
return
|
|
}
|
|
```
|
|
|
|
## Buffer Deallocation Simplification Pass
|
|
|
|
The [semantics of the `bufferization.dealloc` operation](#bufferizationdealloc-bufferizationdeallocop)
|
|
provide a lot of opportunities for optimizations which can be conveniently split
|
|
into patterns using the greedy pattern rewriter. Some of those patterns need
|
|
access to additional analyses such as an analysis that can determine whether two
|
|
MemRef values must, may, or never originate from the same buffer allocation.
|
|
These patterns are collected in the Buffer Deallocation Simplification pass,
|
|
while patterns that don't need additional analyses are registered as part of the
|
|
regular canonicalizer pass. This pass is best run after
|
|
`--ownership-based-buffer-deallocation` followed by `--canonicalize`.
|
|
|
|
The pass applies patterns for the following simplifications:
|
|
* Remove MemRefs from retain list when guaranteed to not alias with any value
|
|
in the 'memref' operand list. This avoids an additional aliasing check with
|
|
the removed value.
|
|
* Split off values in the 'memref' list to new `bufferization.dealloc`
|
|
operations only containing this value in the 'memref' list when it is
|
|
guaranteed to not alias with any other value in the 'memref' list. This
|
|
avoids at least one aliasing check at runtime and enables using a more
|
|
efficient lowering for this new `bufferization.dealloc` operation.
|
|
* Remove values from the 'memref' operand list when it is guaranteed to alias
|
|
with at least one value in the 'retained' list and may not alias any other
|
|
value in the 'retain' list.
|
|
|
|
## Lower Deallocations Pass
|
|
|
|
The `-lower-deallocations` pass transforms all `bufferization.dealloc`
|
|
operations to `memref.dealloc` operations and may also insert operations from
|
|
the `scf`, `func`, and `arith` dialects to make deallocations conditional and
|
|
check whether two MemRef values come from the same allocation at runtime (when
|
|
the `buffer-deallocation-simplification` pass wasn't able to determine it
|
|
statically).
|
|
|
|
The same lowering of the `bufferization.dealloc` operation is also part of the
|
|
`-convert-bufferization-to-memref` conversion pass which also lowers all the
|
|
other operations of the bufferization dialect.
|
|
|
|
We distinguish multiple cases in this lowering pass to provide an overall more
|
|
efficient lowering. In the general case, a library function is created to avoid
|
|
quadratic code size explosion (relative to the number of operands of the dealloc
|
|
operation). The specialized lowerings aim to avoid this library function because
|
|
it requires allocating auxiliary MemRefs of index values.
|
|
|
|
### Generic Lowering
|
|
|
|
A library function is generated to avoid code-size blow-up. On a high level, the
|
|
base-memref of all operands is extracted as an index value and stored into
|
|
specifically allocated MemRefs and passed to the library function which then
|
|
determines whether they come from the same original allocation. This information
|
|
is needed to avoid double-free situations and to correctly retain the MemRef
|
|
values in the `retained` list.
|
|
|
|
**Dealloc Operation Lowering**
|
|
|
|
This lowering supports all features the dealloc operation has to offer. It
|
|
computes the base pointer of each memref (as an index), stores it in a
|
|
new memref helper structure and passes it to the helper function generated
|
|
in `buildDeallocationLibraryFunction`. The results are stored in two lists
|
|
(represented as MemRefs) of booleans passed as arguments. The first list
|
|
stores whether the corresponding condition should be deallocated, the
|
|
second list stores the ownership of the retained values which can be used
|
|
to replace the result values of the `bufferization.dealloc` operation.
|
|
|
|
Example:
|
|
```mlir
|
|
%0:2 = bufferization.dealloc (%m0, %m1 : memref<2xf32>, memref<5xf32>)
|
|
if (%cond0, %cond1)
|
|
retain (%r0, %r1 : memref<1xf32>, memref<2xf32>)
|
|
```
|
|
lowers to (simplified):
|
|
```mlir
|
|
%c0 = arith.constant 0 : index
|
|
%c1 = arith.constant 1 : index
|
|
%dealloc_base_pointer_list = memref.alloc() : memref<2xindex>
|
|
%cond_list = memref.alloc() : memref<2xi1>
|
|
%retain_base_pointer_list = memref.alloc() : memref<2xindex>
|
|
%m0_base_pointer = memref.extract_aligned_pointer_as_index %m0
|
|
memref.store %m0_base_pointer, %dealloc_base_pointer_list[%c0]
|
|
%m1_base_pointer = memref.extract_aligned_pointer_as_index %m1
|
|
memref.store %m1_base_pointer, %dealloc_base_pointer_list[%c1]
|
|
memref.store %cond0, %cond_list[%c0]
|
|
memref.store %cond1, %cond_list[%c1]
|
|
%r0_base_pointer = memref.extract_aligned_pointer_as_index %r0
|
|
memref.store %r0_base_pointer, %retain_base_pointer_list[%c0]
|
|
%r1_base_pointer = memref.extract_aligned_pointer_as_index %r1
|
|
memref.store %r1_base_pointer, %retain_base_pointer_list[%c1]
|
|
%dyn_dealloc_base_pointer_list = memref.cast %dealloc_base_pointer_list :
|
|
memref<2xindex> to memref<?xindex>
|
|
%dyn_cond_list = memref.cast %cond_list : memref<2xi1> to memref<?xi1>
|
|
%dyn_retain_base_pointer_list = memref.cast %retain_base_pointer_list :
|
|
memref<2xindex> to memref<?xindex>
|
|
%dealloc_cond_out = memref.alloc() : memref<2xi1>
|
|
%ownership_out = memref.alloc() : memref<2xi1>
|
|
%dyn_dealloc_cond_out = memref.cast %dealloc_cond_out :
|
|
memref<2xi1> to memref<?xi1>
|
|
%dyn_ownership_out = memref.cast %ownership_out :
|
|
memref<2xi1> to memref<?xi1>
|
|
call @dealloc_helper(%dyn_dealloc_base_pointer_list,
|
|
%dyn_retain_base_pointer_list,
|
|
%dyn_cond_list,
|
|
%dyn_dealloc_cond_out,
|
|
%dyn_ownership_out) : (...)
|
|
%m0_dealloc_cond = memref.load %dyn_dealloc_cond_out[%c0] : memref<2xi1>
|
|
scf.if %m0_dealloc_cond {
|
|
memref.dealloc %m0 : memref<2xf32>
|
|
}
|
|
%m1_dealloc_cond = memref.load %dyn_dealloc_cond_out[%c1] : memref<2xi1>
|
|
scf.if %m1_dealloc_cond {
|
|
memref.dealloc %m1 : memref<5xf32>
|
|
}
|
|
%r0_ownership = memref.load %dyn_ownership_out[%c0] : memref<2xi1>
|
|
%r1_ownership = memref.load %dyn_ownership_out[%c1] : memref<2xi1>
|
|
memref.dealloc %dealloc_base_pointer_list : memref<2xindex>
|
|
memref.dealloc %retain_base_pointer_list : memref<2xindex>
|
|
memref.dealloc %cond_list : memref<2xi1>
|
|
memref.dealloc %dealloc_cond_out : memref<2xi1>
|
|
memref.dealloc %ownership_out : memref<2xi1>
|
|
// replace %0#0 with %r0_ownership
|
|
// replace %0#1 with %r1_ownership
|
|
```
|
|
|
|
**Library function**
|
|
|
|
A library function is built per compilation unit that can be called at
|
|
bufferization dealloc sites to determine whether two MemRefs come from the same
|
|
allocation and their new ownerships.
|
|
|
|
The generated function takes two MemRefs of indices and three MemRefs of
|
|
booleans as arguments:
|
|
* The first argument A should contain the result of the
|
|
extract_aligned_pointer_as_index operation applied to the MemRefs to be
|
|
deallocated
|
|
* The second argument B should contain the result of the
|
|
extract_aligned_pointer_as_index operation applied to the MemRefs to be
|
|
retained
|
|
* The third argument C should contain the conditions as passed directly
|
|
to the deallocation operation.
|
|
* The fourth argument D is used to pass results to the caller. Those
|
|
represent the condition under which the MemRef at the corresponding
|
|
position in A should be deallocated.
|
|
* The fifth argument E is used to pass results to the caller. It
|
|
provides the ownership value corresponding the the MemRef at the same
|
|
position in B
|
|
|
|
This helper function is supposed to be called once for each
|
|
`bufferization.dealloc` operation to determine the deallocation need and
|
|
new ownership indicator for the retained values, but does not perform the
|
|
deallocation itself.
|
|
|
|
Generated code:
|
|
```mlir
|
|
func.func @dealloc_helper(
|
|
%dyn_dealloc_base_pointer_list: memref<?xindex>,
|
|
%dyn_retain_base_pointer_list: memref<?xindex>,
|
|
%dyn_cond_list: memref<?xi1>,
|
|
%dyn_dealloc_cond_out: memref<?xi1>,
|
|
%dyn_ownership_out: memref<?xi1>) {
|
|
%c0 = arith.constant 0 : index
|
|
%c1 = arith.constant 1 : index
|
|
%true = arith.constant true
|
|
%false = arith.constant false
|
|
%num_dealloc_memrefs = memref.dim %dyn_dealloc_base_pointer_list, %c0
|
|
%num_retain_memrefs = memref.dim %dyn_retain_base_pointer_list, %c0
|
|
// Zero initialize result buffer.
|
|
scf.for %i = %c0 to %num_retain_memrefs step %c1 {
|
|
memref.store %false, %dyn_ownership_out[%i] : memref<?xi1>
|
|
}
|
|
scf.for %i = %c0 to %num_dealloc_memrefs step %c1 {
|
|
%dealloc_bp = memref.load %dyn_dealloc_base_pointer_list[%i]
|
|
%cond = memref.load %dyn_cond_list[%i]
|
|
// Check for aliasing with retained memrefs.
|
|
%does_not_alias_retained = scf.for %j = %c0 to %num_retain_memrefs
|
|
step %c1 iter_args(%does_not_alias_aggregated = %true) -> (i1) {
|
|
%retain_bp = memref.load %dyn_retain_base_pointer_list[%j]
|
|
%does_alias = arith.cmpi eq, %retain_bp, %dealloc_bp : index
|
|
scf.if %does_alias {
|
|
%curr_ownership = memref.load %dyn_ownership_out[%j]
|
|
%updated_ownership = arith.ori %curr_ownership, %cond : i1
|
|
memref.store %updated_ownership, %dyn_ownership_out[%j]
|
|
}
|
|
%does_not_alias = arith.cmpi ne, %retain_bp, %dealloc_bp : index
|
|
%updated_aggregate = arith.andi %does_not_alias_aggregated,
|
|
%does_not_alias : i1
|
|
scf.yield %updated_aggregate : i1
|
|
}
|
|
// Check for aliasing with dealloc memrefs in the list before the
|
|
// current one, i.e.,
|
|
// `fix i, forall j < i: check_aliasing(%dyn_dealloc_base_pointer[j],
|
|
// %dyn_dealloc_base_pointer[i])`
|
|
%does_not_alias_any = scf.for %j = %c0 to %i step %c1
|
|
iter_args(%does_not_alias_agg = %does_not_alias_retained) -> (i1) {
|
|
%prev_dealloc_bp = memref.load %dyn_dealloc_base_pointer_list[%j]
|
|
%does_not_alias = arith.cmpi ne, %prev_dealloc_bp, %dealloc_bp
|
|
%updated_alias_agg = arith.andi %does_not_alias_agg, %does_not_alias
|
|
scf.yield %updated_alias_agg : i1
|
|
}
|
|
%dealloc_cond = arith.andi %does_not_alias_any, %cond : i1
|
|
memref.store %dealloc_cond, %dyn_dealloc_cond_out[%i] : memref<?xi1>
|
|
}
|
|
return
|
|
}
|
|
```
|
|
|
|
### Specialized Lowerings
|
|
|
|
Currently, there are two special lowerings for common cases to avoid the library
|
|
function and thus unnecessary memory load and store operations and function
|
|
calls:
|
|
|
|
**One memref, no retained**
|
|
|
|
Lower a simple case without any retained values and a single MemRef. Ideally,
|
|
static analysis can provide enough information such that the
|
|
`buffer-deallocation-simplification` pass is able to split the dealloc
|
|
operations up into this simple case as much as possible before running this
|
|
pass.
|
|
|
|
Example:
|
|
```mlir
|
|
bufferization.dealloc (%arg0 : memref<2xf32>) if (%arg1)
|
|
```
|
|
is lowered to
|
|
```mlir
|
|
scf.if %arg1 {
|
|
memref.dealloc %arg0 : memref<2xf32>
|
|
}
|
|
```
|
|
|
|
In most cases, the branch condition is either constant 'true' or 'false' and can
|
|
thus be optimized away entirely by the canonicalizer pass.
|
|
|
|
**One memref, arbitrarily many retained**
|
|
|
|
A special case lowering for the deallocation operation with exactly one MemRef,
|
|
but an arbitrary number of retained values. The size of the code produced by
|
|
this lowering is linear to the number of retained values.
|
|
|
|
Example:
|
|
```mlir
|
|
%0:2 = bufferization.dealloc (%m : memref<2xf32>) if (%cond)
|
|
retain (%r0, %r1 : memref<1xf32>, memref<2xf32>)
|
|
return %0#0, %0#1 : i1, i1
|
|
```
|
|
is lowered to
|
|
```mlir
|
|
%m_base_pointer = memref.extract_aligned_pointer_as_index %m
|
|
%r0_base_pointer = memref.extract_aligned_pointer_as_index %r0
|
|
%r0_does_not_alias = arith.cmpi ne, %m_base_pointer, %r0_base_pointer
|
|
%r1_base_pointer = memref.extract_aligned_pointer_as_index %r1
|
|
%r1_does_not_alias = arith.cmpi ne, %m_base_pointer, %r1_base_pointer
|
|
%not_retained = arith.andi %r0_does_not_alias, %r1_does_not_alias : i1
|
|
%should_dealloc = arith.andi %not_retained, %cond : i1
|
|
scf.if %should_dealloc {
|
|
memref.dealloc %m : memref<2xf32>
|
|
}
|
|
%true = arith.constant true
|
|
%r0_does_alias = arith.xori %r0_does_not_alias, %true : i1
|
|
%r0_ownership = arith.andi %r0_does_alias, %cond : i1
|
|
%r1_does_alias = arith.xori %r1_does_not_alias, %true : i1
|
|
%r1_ownership = arith.andi %r1_does_alias, %cond : i1
|
|
return %r0_ownership, %r1_ownership : i1, i1
|
|
```
|