| # 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 |
| ``` |