diff options
author | Sean Silva <silvasean@google.com> | 2020-10-28 13:25:48 -0700 |
---|---|---|
committer | Sean Silva <silvasean@google.com> | 2020-11-10 13:44:44 -0800 |
commit | 53a0d45db6d0f33dfbb724c99ce2560ae25473c2 (patch) | |
tree | 7c04d4e857587522537604fdc51ea1fbf540dafc /llvm/lib/CodeGen/SafeStack.cpp | |
parent | b4fa28b408662a638f7d835b302db08633ccb91e (diff) | |
download | llvm-53a0d45db6d0f33dfbb724c99ce2560ae25473c2.zip llvm-53a0d45db6d0f33dfbb724c99ce2560ae25473c2.tar.gz llvm-53a0d45db6d0f33dfbb724c99ce2560ae25473c2.tar.bz2 |
[mlir] Add pass to convert elementwise ops to linalg.
This patch converts elementwise ops on tensors to linalg.generic ops
with the same elementwise op in the payload (except rewritten to
operate on scalars, obviously). This is a great form for later fusion to
clean up.
E.g.
```
// Compute: %arg0 + %arg1 - %arg2
func @f(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> tensor<?xf32> {
%0 = addf %arg0, %arg1 : tensor<?xf32>
%1 = subf %0, %arg2 : tensor<?xf32>
return %1 : tensor<?xf32>
}
```
Running this through
`mlir-opt -convert-std-to-linalg -linalg-fusion-for-tensor-ops` we get:
```
func @f(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> tensor<?xf32> {
%0 = linalg.generic {indexing_maps = [#map0, #map0, #map0, #map0], iterator_types = ["parallel"]} ins(%arg0, %arg1, %arg2 : tensor<?xf32>, tensor<?xf32>, tensor<?xf32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): // no predecessors
%1 = addf %arg3, %arg4 : f32
%2 = subf %1, %arg5 : f32
linalg.yield %2 : f32
} -> tensor<?xf32>
return %0 : tensor<?xf32>
}
```
So the elementwise ops on tensors have nicely collapsed into a single
linalg.generic, which is the form we want for further transformations.
Differential Revision: https://reviews.llvm.org/D90354
Diffstat (limited to 'llvm/lib/CodeGen/SafeStack.cpp')
0 files changed, 0 insertions, 0 deletions