[Mlir-commits] [mlir] [mlir][NFC] Improve bufferization documentation (PR #89495)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Apr 20 01:53:01 PDT 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Matthias Springer (matthias-springer)
<details>
<summary>Changes</summary>
* Add example for `test-analysis-only` and `print-conflicts`.
* Mention other bufferization-related passes.
* Update outdated documentation.
---
Patch is 303.48 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/89495.diff
3 Files Affected:
- (modified) mlir/docs/Bufferization.md (+214-90)
- (added) mlir/docs/includes/img/bufferization_passes.svg (+1)
- (added) mlir/docs/includes/img/bufferization_tensor_insert_dst.svg (+1)
``````````diff
diff --git a/mlir/docs/Bufferization.md b/mlir/docs/Bufferization.md
index 80853582221258..6a49bea9c68c5d 100644
--- a/mlir/docs/Bufferization.md
+++ b/mlir/docs/Bufferization.md
@@ -5,35 +5,45 @@
## Overview
Bufferization in MLIR is the process of converting ops with `tensor` semantics
-to ops with `memref` semantics. MLIR provides an infrastructure that bufferizes
-an entire program in a single pass (*One-Shot Bufferize*). This infrastructure
-bufferizes all ops that implement the
-[`BufferizableOpInterface`](https://github.com/llvm/llvm-project/blob/17a68065c378da74805e4e1b9a5b78cc9f83e580/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td)
-can be bufferized.
-
-MLIR has an older bufferization infrastructure built around
-[dialect conversion](DialectConversion.md). Most dialect conversion
-bufferization patterns have been migrated to One-Shot Bufferize, but some
-functionality such as function boundary bufferization still depends on dialect
-conversion and its type converter. New projects should use One-Shot Bufferize,
-as the dialect conversion-based bufferization will eventually be deprecated.
-Moreover, One-Shot Bufferize results in better bufferization with fewer memory
-allocations and buffer copies. This documentation is mostly about One-Shot
-Bufferize, but also describes how to gradually migrate a project from dialect
-conversion-based bufferization to One-Shot Bufferize.
+to ops with `memref` semantics. There are multiple MLIR passes that are related
+to bufferization. These passes typically run as one of the last steps in a
+pass pipeline, right before lowering to `memref` ops to LLVM. That is because
+many transformations are easier or only supported in tensor land; e.g.,
+[tile/fuse/… on tensors first](https://llvm.discourse.group/t/rfc-linalg-on-tensors-update-and-comprehensive-bufferization-rfc/3373),
+then bufferize the remaining IR.
+
+
+
+The most important bufferization pass is *One-Shot Bufferize*: This pass
+rewrites `tensor` IR to `memref` IR. There are additional helper passes that
+preprocess IR (e.g., so that IR can be bufferized more efficiently), perform
+buffer-level optimizations such as allocation hoisting, and
+[insert buffer deallocation ops](OwnershipBasedBufferDeallocation.md) so that
+the resulting `memref` IR has no memory leaks.
+
+## Deprecated Passes
+
+The old dialect conversion-based bufferization passes have been deprecated and
+should not be used anymore. Most of those passes have already been removed from
+MLIR. One-Shot Bufferize produces in better bufferization results with fewer
+memory allocations and buffer copies.
+
+The buffer deallocation pass has been deprecated in favor of the ownership-based
+buffer deallocation pipeline. The deprecated pass has some limitations that may
+cause memory leaks in the resulting IR.
## What is One-Shot Bufferize?
-One-Shot Bufferize is a new tensor bufferization pass designed for IR in
+One-Shot Bufferize is a tensor bufferization pass designed for IR in
[destination-passing style](https://www.microsoft.com/en-us/research/wp-content/uploads/2016/11/dps-fhpc17.pdf),
and with aggressive in-place bufferization.
One-Shot Bufferize is:
-* **Monolithic**: A single MLIR pass does the entire work, whereas the
- previous bufferization in MLIR was split across multiple passes residing in
- different dialects. In One-Shot Bufferize, `BufferizableOpInterface`
- implementations are spread across different dialects.
+* **Monolithic**: A single MLIR pass does the entire work.
+
+* **Extensible** via an op interface: All ops that implement
+ `BufferizableOpInterface` can be bufferized.
* A **whole-function at a time analysis**. In-place bufferization decisions
are made by analyzing SSA use-def chains on tensors. Op interface
@@ -41,10 +51,7 @@ One-Shot Bufferize is:
ops, but also helper methods for One-Shot Bufferize's analysis to query
information about an op's bufferization/memory semantics.
-* **Extensible** via an op interface: All ops that implement
- `BufferizableOpInterface` can be bufferized.
-
-* **2-Pass**: Bufferization is internally broken down into 2 steps: First,
+* **2-Phase**: Bufferization is internally broken down into 2 steps: First,
analyze the entire IR and make bufferization decisions. Then, bufferize
(rewrite) the IR. The analysis has access to exact SSA use-def information.
It incrementally builds alias and equivalence sets and does not rely on a
@@ -60,27 +67,17 @@ One-Shot Bufferize is:
of `AnalysisState` that implements a small number virtual functions can
serve as a custom analysis. It is even possible to run One-Shot Bufferize
without any analysis (`AlwaysCopyAnalysisState`), in which case One-Shot
- Bufferize behaves exactly like the old dialect conversion-based
- bufferization (i.e., copy every buffer before writing to it).
+ Bufferize copies every buffer before writing to it.
-To reduce complexity, One-Shot Bufferize should be
-[run after other transformations](https://llvm.discourse.group/t/rfc-linalg-on-tensors-update-and-comprehensive-bufferization-rfc/3373),
-typically as one of the last steps right before lowering memref ops. Many
-transformations are easier in tensor land; e.g., tile/fuse/… on tensors first,
-then bufferize the remaining IR.
-
-From an architecture perspective, One-Shot Bufferize consists of
-[BufferizableOpInterface](https://github.com/llvm/llvm-project/blob/17a68065c378da74805e4e1b9a5b78cc9f83e580/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td)
-(and its implementations) and an
-[analysis](https://github.com/llvm/llvm-project/blob/ae2764e835a26bad9774803eca0a6530df2a3e2d/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h#L164)
-of tensor SSA values that decides if a buffer can be used directly or must be
-copied. The [bufferize] method of the op interface inspects analysis results and
-rewrites tensor ops into memref ops.
+Note that One-Shot Bufferize does not deallocate buffers. That is done by the
+[Ownership-based Buffer Deallocation passes](OwnershipBasedBufferDeallocation.md).
## Goals of Bufferization
-The high-level goal of every bufferization technique is to: 1. Use as little
-memory as possible. 2. Copy as little memory as possible.
+The high-level goal of every bufferization technique is to:
+
+1. Use as little memory as possible.
+2. Copy as little memory as possible.
This implies reusing already allocated buffers when possible, turning
bufferization into an algorithmically complex problem with similarities to
@@ -102,40 +99,46 @@ choosing an already existing buffer, we must be careful not to accidentally
overwrite data that is still needed later in the program.
To simplify this problem, One-Shot Bufferize was designed to take advantage of
-*destination-passing style*. This form exists in itself independently of
-bufferization and is tied to SSA semantics: many ops are “updating” part of
-their input SSA variable. For example the LLVM instruction
+*destination-passing style* (DPS). In MLIR, DPS op should implement the
+[`DestinationStyleOpInterface`](https://github.com/llvm/llvm-project/blob/792d437b56adfb3416daf8105942d4899fb82763/mlir/include/mlir/Interfaces/DestinationStyleOpInterface.td).
+DPS exists in itself independently of bufferization and is tied to SSA
+semantics: many ops are "updating" a part of their input SSA variables. For
+example the LLVM instruction
[`insertelement`](https://llvm.org/docs/LangRef.html#insertelement-instruction)
is inserting an element inside a vector. Since SSA values are immutable, the
operation returns a copy of the input vector with the element inserted.
-Another example in MLIR is `linalg.generic`, which always has an extra `outs`
-operand which provides the initial values to update (for example when the
-operation is doing a reduction).
+Another example in MLIR is `linalg.generic` on tensors, which always has an
+extra `outs` operand for each result, which provides the initial values to
+update (for example when the operation is doing a reduction).
-This input is referred to as "destination" in the following (quotes are
+`outs` operands are referred to as "destinations" in the following (quotes are
important as this operand isn't modified in place but copied) and comes into
place in the context of bufferization as a possible "anchor" for the
bufferization algorithm. This allows the user to shape the input in a form that
guarantees close to optimal bufferization result when carefully choosing the
SSA value used as "destination".
-For every tensor result, a "destination-passing" style op has a corresponding
-tensor operand. If there aren't any other uses of this tensor, the bufferization
-can alias it with the op result and perform the operation "in-place" by reusing
-the buffer allocated for this "destination" input.
+For every tensor result, a DPS op has a corresponding tensor operand. If there
+aren't any other conflicting uses of this tensor, the bufferization can alias
+it with the op result and perform the operation "in-place" by reusing the buffer
+allocated for this "destination" input.
-As an example, consider the following op: `%0 = tensor.insert %cst into
-%t[%idx] : tensor<?xf32>`
+As an example, consider the following op: `%r = tensor.insert %f into
+%t[%idx] : tensor<5xf32>`
+
+
`%t` is the "destination" in this example. When choosing a buffer for the result
-`%0`, denoted as `buffer(%0)`, One-Shot Bufferize considers only two options:
+`%r`, denoted as `buffer(%r)`, One-Shot Bufferize considers only two options:
-1. `buffer(%0) = buffer(%t)` : alias the "destination" tensor with the
- result and perform the operation in-place.
-2. `buffer(%0)` is a newly allocated buffer.
+1. `buffer(%r) = buffer(%t)`: store the result in the existing `buffer(%t)`.
+ Note that this is not always possible. E.g., if the old contents of
+ `buffer(%t)` are still needed. One-Shot Bufferize's main task is to detect
+ such cases and fall back to the second option when necessary.
+2. `buffer(%r)` is a newly allocated buffer.
There may be other buffers in the same function that could potentially be used
-for `buffer(%0)`, but those are not considered by One-Shot Bufferize to keep the
+for `buffer(%r)`, but those are not considered by One-Shot Bufferize to keep the
bufferization simple. One-Shot Bufferize could be extended to consider such
buffers in the future to achieve a better quality of bufferization.
@@ -151,7 +154,7 @@ memory allocation. E.g.:
```
The result of `tensor.generate` does not have a "destination" operand, so
-bufferization allocates a new buffer. This could be avoided by choosing an
+bufferization allocates a new buffer. This could be avoided by instead using an
op such as `linalg.generic`, which can express the same computation with a
"destination" operand, as specified behind outputs (`outs`):
@@ -198,12 +201,61 @@ e.g.:
```mlir
%0 = "my_dialect.some_op"(%t) : (tensor<?xf32>) -> (tensor<?xf32>)
%1 = "my_dialect.another_op"(%0) : (tensor<?xf32>) -> (tensor<?xf32>)
+
+// "yet_another_op" likely needs to read the data of %0, so "another_op" cannot
+// in-place write to buffer(%0).
%2 = "my_dialect.yet_another_op"(%0) : (tensor<?xf32>) -> (tensor<?xf32>)
```
-One-Shot Bufferize has debug flags (`test-analysis-only print-conflicts`) that
-print the results of the analysis and explain to the user why buffer copies were
-inserted.
+## Tensor / MemRef Boundary
+
+The bufferization dialect provides a few helper ops to connect tensor IR (that
+should be bufferized) with existing buffers (that may be allocated/provided by
+a different runtime/library/etc.).
+
+`bufferization.to_memref %t` returns the future buffer of a tensor SSA value.
+`bufferization.to_tensor %m` returns a tensor SSA value for a given MemRef
+buffer. `bufferization.materialize_in_destination` indicates that a tensor value
+should materialize in a certain buffer.
+
+Consider the following example, where a TOSA matmul result should materialize in
+an existing buffer `%C`:
+
+```mlir
+// Batched TOSA matrix multiplication. %A and %B are the
+// inputs, %C is the output.
+func.func @test_matmul(%A: memref<1x17x19xf32>,
+ %B: memref<1x19x29xf32>,
+ %C: memref<1x17x29xf32>) {
+
+ %A_tensor = bufferization.to_tensor %A restrict : memref<1x17x19xf32>
+ %B_tensor = bufferization.to_tensor %B restrict : memref<1x19x29xf32>
+
+ %0 = tosa.matmul %A_tensor, %B_tensor
+ : (tensor<1x17x19xf32>, tensor<1x19x29xf32>) ->
+ tensor<1x17x29xf32>
+
+ bufferization.materialize_in_destination
+ %0 in restrict writable %C
+ : (tensor<1x17x29xf32>, memref<1x17x29xf32>) -> ()
+
+ return
+}
+```
+
+Note that all bufferization ops in this example have the `restrict` unit
+attribute set. This attribute is similar to the C restrict keyword and indicates
+that there is no other `to_tensor` or `materialize_in_destination` op with
+the same or an aliasing MemRef operand. Only such
+`to_tensor`/`materialize_in_destination` ops are supported. The `restrict`
+attribute gives strong aliasing guarantees to the bufferization analysis and
+allows us to look only at the tensor IR in a program. (Ops that do not operate
+on tensors are ignored by the One-Shot Bufferize.)
+
+Also note that `tosa.matmul` cannot be bufferized as is: there is no
+`BufferizableOpInterface` implementation for that op. However, the op can be
+lowered to a combination of `tensor.empty` and `linalg.matmul`, which can be
+bufferized.
## Using One-Shot Bufferize
@@ -221,17 +273,14 @@ By default, One-Shot Bufferize fails when it encounters an op with tensor
semantics (i.e., tensor result or tensor operand) that is not bufferizable
(i.e., does not implement `BufferizableOpInterface`). This can be avoided with
`allow-unknown-ops`. In that case, One-Shot Bufferize inserts
-`to_memref`/`to_tensor` ops around the bufferization boundary. These ops are
-named versions of `unrealized_conversion_cast`. Note that One-Shot Bufferize's
-analysis can currently not analyze these ops, so input IR with such ops may fail
-bufferization. Therefore, running One-Shot Bufferize multiple times in a
-sequence is also not supported at the moment.
+`to_memref`/`to_tensor` ops around the bufferization boundary.
One-Shot Bufferize can be configured to bufferize only ops from a set of
dialects with `dialect-filter`. This can be useful for gradually migrating from
dialect conversion-based bufferization to One-Shot Bufferize. One-Shot Bufferize
must run first in such a case, because dialect conversion-based bufferization
-generates `to_tensor`/`to_memref` ops which One-Shot Bufferize cannot analyze.
+generates `to_tensor` ops without the `restrict` unit attribute, which One-Shot
+Bufferize cannot analyze.
One-Shot Bufferize can also be called programmatically with
[`bufferization::runOneShotBufferize`](https://github.com/llvm/llvm-project/blob/ae2764e835a26bad9774803eca0a6530df2a3e2d/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h#L167).
@@ -240,6 +289,14 @@ Alternatively,
skips the analysis and inserts a copy on every buffer write, just like the
dialect conversion-based bufferization.
+By default, function boundaries are not bufferized. This is because there are
+currently limitations around function graph bufferization: recursive
+calls are not supported. As long as there are no recursive calls, function
+boundary bufferization can be enabled with `bufferize-function-boundaries`. Each
+tensor function argument and tensor function result is then turned into a
+memref. The layout map of the memref type can be controlled with
+`function-boundary-type-conversion`.
+
## Memory Layouts
One-Shot Bufferize bufferizes ops from top to bottom. This works well when all
@@ -319,6 +376,11 @@ To get a better intuition of the interface methods, we invite users to take a
look at existing implementations in MLIR, e.g., the implementation of
`tensor.insert` or `tensor.extract`.
+Interface implementations of DPS ops (that implement
+`DestinationStyleOpInterface`) can derive from
+`DstBufferizableOpInterfaceExternalModel`, which provides all necessary
+method implementations except for `bufferize`.
+
## Debugging Buffer Copies
To get a better understanding of why One-Shot Bufferize introduced a buffer
@@ -338,14 +400,90 @@ There are two reasons why a buffer copy may be inserted.
In the first case, `print-conflicts` illustrates the conflict in the form of a
("read", "conflicting write", "last write") tuple.
-## Understanding the SSA Use-Def Chain Analysis
+A RaW conflict consists of three parts, in the following order according to
+op dominance:
+
+1. **Definition:** A tensor `%t` is defined.
+2. **Conflicting Write:** An operation writes to `buffer(%t)`.
+3. **Read:** An operation reads `%t`.
+
+When such a RaW conflict is detected during the analysis phase, One-Shot
+Bufferize will insert a buffer copy for the conflicting write.
+
+**Example**
+
+```mlir
+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries test-analysis-only print-conflicts"
+func.func @test(%arg0: f32, %arg1: f32, %arg2: index, %arg3: index) -> (f32, tensor<3xf32>) {
+ // Create a new tensor with [%arg0, %arg0, %arg0].
+ %0 = tensor.from_elements %arg0, %arg0, %arg0 : tensor<3xf32>
+
+ // Insert something into the new tensor.
+ %1 = tensor.insert %arg1 into %0[%arg2] : tensor<3xf32>
+
+ // Read from the old tensor.
+ %r = tensor.extract %0[%arg3] : tensor<3xf32>
+
+ // Return the extracted value and the result of the insertion.
+ func.return %r, %1 : f32, tensor<3xf32>
+}
+```
+
+The output IR is as follows:
+
+```mlir
+func.func @test(%arg0: f32, %arg1: f32, %arg2: index, %arg3: index) -> (f32, tensor<3xf32>) {
+ %from_elements = tensor.from_elements %arg0, %arg0, %arg0 {"C_0[DEF: result 0]"} : tensor<3xf32>
+ %inserted = tensor.insert %arg1 into %from_elements[%arg2] {"C_0[CONFL-WRITE: 1]", __inplace_operands_attr__ = ["none", "false", "none"]} : tensor<3xf32>
+ %extracted = tensor.extract %from_elements[%arg3] {"C_0[READ: 0]", __inplace_operands_attr__ = ["true", "none"]} : tensor<3xf32>
+ return {__inplace_operands_attr__ = ["none", "true"]} %extracted, %inserted : f32, tensor<3xf32>
+}
+```
+
+Note that the IR was not bufferized. It was merely annotated with the results
+of the bufferization analysis. Every operation with tensor semantics has a
+`__inplace_operands_attr__` attribute with one value per operand. If an operand
+is not a tensor, the respective value is `none`. Otherwise, if the operand was
+decided to be bufferized in-place, the value is `true`. A value of `false`
+indicates a buffer copy. In the above example, a buffer copy would be inserted
+for `tensor.insert`, so that it does not overwrite `buffer(%from_elements)`,
+which is still needed for `tensor.extract`.
+
+For each RaW (there is only one in the example), three `C_i` attributes were
+added:
+
+* `C_0[DEF: result 0]`: A tensor is defined: 0-th result of
+ `tensor.from_elements`.
+* `C_0[CONFL-WRITE: 1]`: An operation (if bufferized in-place) would write into
+ the future buffer of the defined tensor: 1-st operand of `tensor.insert`.
+* `C_0[READ: 0]`: An operation reads the tensor definition: 0-th operand of
+ `tensor.extract`.
+
+The fully bufferized IR (with the inserted buffer copy) is as follows:
+
+```mlir
+func.func @test(%arg0: f32, %arg1: f32, %arg2: index, %arg3: index) -> (f32, memref<3xf32>) {
+ %c2 = arith.constant 2 : index
+ %c1 = arith.constant 1 : index
+ %c0 = arith.constant 0 : index
+ %alloc = memref.alloc() {alignment = 64 : i64} : memref<3xf32>
+ memref.store %arg0, %alloc[%c0] : memref<3xf32>
+ memref.store %arg0, %alloc[%c1] : memref<3xf32>
+ memref.store %arg0, %alloc[%c2] : memref<3xf32>
+ %alloc_0 = memref.alloc() {alignment = 64 : i64} : memref<3xf32>
+ memref.copy %alloc, %alloc_0 : memref<3xf32> to memref<3xf32>
+ memref.store %arg1, %alloc_0[%arg2] : memref<3xf32>
+ %0 = memref.load %alloc[%arg3] : memref<3xf32>
+ return %0, %alloc_0 : f32, memref<3xf32>
+}
+```
To get a better understanding of the SSA Use-Def Chain Analysis and the RaW
-conflict...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/89495
More information about the Mlir-commits
mailing list