Add a generic Linalg op
authorNicolas Vasilache <ntv@google.com>
Fri, 2 Aug 2019 16:53:08 +0000 (09:53 -0700)
committerA. Unique TensorFlower <gardener@tensorflow.org>
Fri, 2 Aug 2019 16:53:41 +0000 (09:53 -0700)
commit600c47e77b008994e39d2f526c3155c5da59f11c
treee57c841f75f958ca841f74db2a392bcb52614f50
parent192039e8bef5e3977995b9ac5f0c3199a1e5b8b4
Add a generic Linalg op

This CL introduces a linalg.generic op to represent generic tensor contraction operations on views.

A linalg.generic operation requires a numbers of attributes that are sufficient to emit the computation in scalar form as well as compute the appropriate subviews to enable tiling and fusion.

These attributes are very similar to the attributes for existing operations such as linalg.matmul etc and existing operations can be implemented with the generic form.

In the future, most existing operations can be implemented using the generic form.

This CL starts by splitting out most of the functionality of the linalg::NInputsAndOutputs trait into a ViewTrait that queries the per-instance properties of the op. This allows using the attribute informations.

This exposes an ordering of verifiers issue where ViewTrait::verify uses attributes but the verifiers for those attributes have not been run. The desired behavior would be for the verifiers of the attributes specified in the builder to execute first but it is not the case atm. As a consequence, to emit proper error messages and avoid crashing, some of the
linalg.generic methods are defensive as such:
```
    unsigned getNumInputs() {
      // This is redundant with the `n_views` attribute verifier but ordering of verifiers
      // may exhibit cases where we crash instead of emitting an error message.
      if (!getAttr("n_views") || n_views().getValue().size() != 2)
        return 0;
```

In pretty-printed form, the specific attributes required for linalg.generic are factored out in an independent dictionary named "_". When parsing its content is flattened and the "_name" is dropped. This allows using aliasing for reducing boilerplate at each linalg.generic invocation while benefiting from the Tablegen'd verifier form for each named attribute in the dictionary.

For instance, implementing linalg.matmul in terms of linalg.generic resembles:

```
func @mac(%a: f32, %b: f32, %c: f32) -> f32 {
  %d = mulf %a, %b: f32
  %e = addf %c, %d: f32
  return %e: f32
}
#matmul_accesses = [
  (m, n, k) -> (m, k),
  (m, n, k) -> (k, n),
  (m, n, k) -> (m, n)
]
#matmul_trait = {
  doc = "C(m, n) += A(m, k) * B(k, n)",
  fun = @mac,
  indexing_maps = #matmul_accesses,
  library_call = "linalg_matmul",
  n_views = [2, 1],
  n_loop_types = [2, 1, 0]
}
```

And can be used in multiple places as:
```
  linalg.generic #matmul_trait %A, %B, %C [other-attributes] :
    !linalg.view<?x?xf32>, !linalg.view<?x?xf32>, !linalg.view<?x?xf32>
```

In the future it would be great to have a mechanism to alias / register a new
linalg.op as a pair of linalg.generic, #trait.

Also, note that with one could theoretically only specify the `doc` string and parse all the attributes from it.

PiperOrigin-RevId: 261338740
16 files changed:
mlir/include/mlir/AffineOps/AffineOps.td
mlir/include/mlir/AffineOps/AffineOpsBase.td [new file with mode: 0644]
mlir/include/mlir/EDSC/Intrinsics.h
mlir/include/mlir/IR/AffineMap.h
mlir/include/mlir/IR/Builders.h
mlir/include/mlir/Linalg/IR/LinalgLibraryOps.td
mlir/include/mlir/Linalg/IR/LinalgOps.h
mlir/include/mlir/Linalg/IR/LinalgTraits.h
mlir/lib/IR/AffineMap.cpp
mlir/lib/IR/Builders.cpp
mlir/lib/Linalg/IR/LinalgOps.cpp
mlir/lib/Linalg/Transforms/LowerToLLVMDialect.cpp
mlir/lib/Linalg/Transforms/Tiling.cpp
mlir/test/Linalg/invalid-generic-op.mlir [new file with mode: 0644]
mlir/test/Linalg/loops.mlir
mlir/test/Linalg/roundtrip.mlir