[mlir][Linalg] Create a tool to generate named Linalg ops from a Tensor Comprehension...
authorNicolas Vasilache <ntv@google.com>
Fri, 10 Apr 2020 17:54:08 +0000 (13:54 -0400)
committerNicolas Vasilache <ntv@google.com>
Fri, 10 Apr 2020 17:59:25 +0000 (13:59 -0400)
commit882ba484743763b8560b08f483ae21d26ab336f9
tree6655577bef50793c5acef98a99f65ab2f6100e8a
parenta04ab2ec08014d36337d686d43ecef7668622f1b
[mlir][Linalg] Create a tool to generate named Linalg ops from a Tensor Comprehensions-like specification.

Summary:

This revision adds a tool that generates the ODS and C++ implementation for "named" Linalg ops according to the [RFC discussion](https://llvm.discourse.group/t/rfc-declarative-named-ops-in-the-linalg-dialect/745).

While the mechanisms and language aspects are by no means set in stone, this revision allows connecting the pieces end-to-end from a mathematical-like specification.

Some implementation details and short-term decisions taken for the purpose of bootstrapping and that are not set in stone include:

    1. using a "[Tensor Comprehension](https://arxiv.org/abs/1802.04730)-inspired" syntax
    2. implicit and eager discovery of dims and symbols when parsing
    3. using EDSC ops to specify the computation (e.g. std_addf, std_mul_f, ...)

A followup revision will connect this tool to tablegen mechanisms and allow the emission of named Linalg ops that automatically lower to various loop forms and run end to end.

For the following "Tensor Comprehension-inspired" string:

```
    def batch_matmul(A: f32(Batch, M, K), B: f32(K, N)) -> (C: f32(Batch, M, N)) {
      C(b, m, n) = std_addf<k>(std_mulf(A(b, m, k), B(k, n)));
    }
```

With -gen-ods-decl=1, this emits (modulo formatting):

```
      def batch_matmulOp : LinalgNamedStructured_Op<"batch_matmul", [
        NInputs<2>,
        NOutputs<1>,
        NamedStructuredOpTraits]> {
          let arguments = (ins Variadic<LinalgOperand>:$views);
          let results = (outs Variadic<AnyRankedTensor>:$output_tensors);
          let extraClassDeclaration = [{
            llvm::Optional<SmallVector<StringRef, 8>> referenceIterators();
            llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps();
            void regionBuilder(ArrayRef<BlockArgument> args);
          }];
          let hasFolder = 1;
      }
```

With -gen-ods-impl, this emits (modulo formatting):

```
      llvm::Optional<SmallVector<StringRef, 8>> batch_matmul::referenceIterators() {
          return SmallVector<StringRef, 8>{ getParallelIteratorTypeName(),
                                            getParallelIteratorTypeName(),
                                            getParallelIteratorTypeName(),
                                            getReductionIteratorTypeName() };
      }
      llvm::Optional<SmallVector<AffineMap, 8>> batch_matmul::referenceIndexingMaps()
      {
        MLIRContext *context = getContext();
        AffineExpr d0, d1, d2, d3;
        bindDims(context, d0, d1, d2, d3);
        return SmallVector<AffineMap, 8>{
            AffineMap::get(4, 0, {d0, d1, d3}),
            AffineMap::get(4, 0, {d3, d2}),
            AffineMap::get(4, 0, {d0, d1, d2}) };
      }
      void batch_matmul::regionBuilder(ArrayRef<BlockArgument> args) {
        using namespace edsc;
        using namespace intrinsics;
        ValueHandle _0(args[0]), _1(args[1]), _2(args[2]);

        ValueHandle _4 = std_mulf(_0, _1);
        ValueHandle _5 = std_addf(_2, _4);
        (linalg_yield(ValueRange{ _5 }));
      }
```

Differential Revision: https://reviews.llvm.org/D77067
mlir/docs/Dialects/Linalg.md
mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
mlir/include/mlir/IR/AffineExpr.h
mlir/lib/IR/AffineExpr.cpp
mlir/test/CMakeLists.txt
mlir/test/lit.cfg.py
mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc [new file with mode: 0644]
mlir/tools/CMakeLists.txt
mlir/tools/mlir-linalg-ods-gen/CMakeLists.txt [new file with mode: 0644]
mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp [new file with mode: 0644]