[all-commits] [llvm/llvm-project] b4fa28: [mlir] Add ElementwiseMappable trait and apply it ...

[Sean Silva via All-commits]

  Branch: refs/heads/master
  Home:   https://github.com/llvm/llvm-project
  Commit: b4fa28b408662a638f7d835b302db08633ccb91e
      https://github.com/llvm/llvm-project/commit/b4fa28b408662a638f7d835b302db08633ccb91e
  Author: Sean Silva <silvasean at google.com>
  Date:   2020-11-10 (Tue, 10 Nov 2020)

  Changed paths:
    M mlir/docs/Traits.md
    M mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
    M mlir/include/mlir/IR/OpBase.td
    M mlir/include/mlir/IR/OpDefinition.h
    M mlir/lib/IR/Operation.cpp
    M mlir/test/Dialect/Standard/invalid.mlir
    M mlir/test/IR/invalid-ops.mlir
    M mlir/test/IR/traits.mlir
    M mlir/test/lib/Dialect/Test/TestOps.td

  Log Message:
  -----------
  [mlir] Add ElementwiseMappable trait and apply it to std elementwise ops.

This patch adds an `ElementwiseMappable` trait as discussed in the RFC
here:
https://llvm.discourse.group/t/rfc-std-elementwise-ops-on-tensors/2113/23

This trait can power a number of transformations and analyses.
A subsequent patch adds a convert-elementwise-to-linalg pass exhibits
how this trait allows writing generic transformations.
See https://reviews.llvm.org/D90354 for that patch.

This trait slightly changes some verifier messages, but the diagnostics
are usually about as good. I fiddled with the ordering of the trait in
the .td file trait lists to minimize the changes here.

Differential Revision: https://reviews.llvm.org/D90731


  Commit: 53a0d45db6d0f33dfbb724c99ce2560ae25473c2
      https://github.com/llvm/llvm-project/commit/53a0d45db6d0f33dfbb724c99ce2560ae25473c2
  Author: Sean Silva <silvasean at google.com>
  Date:   2020-11-10 (Tue, 10 Nov 2020)

  Changed paths:
    M mlir/include/mlir/Dialect/Linalg/Passes.h
    M mlir/include/mlir/Dialect/Linalg/Passes.td
    A mlir/integration_test/Dialect/Linalg/CPU/test-elementwise.mlir
    M mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
    A mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp
    A mlir/test/Dialect/Linalg/convert-elementwise-to-linalg.mlir

  Log Message:
  -----------
  [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


Compare: https://github.com/llvm/llvm-project/compare/8262e94a6da7...53a0d45db6d0