[all-commits] [llvm/llvm-project] 882ba4: [mlir][Linalg] Create a tool to generate named Lin...

[Nicolas Vasilache via All-commits]

  Branch: refs/heads/master
  Home:   https://github.com/llvm/llvm-project
  Commit: 882ba484743763b8560b08f483ae21d26ab336f9
      https://github.com/llvm/llvm-project/commit/882ba484743763b8560b08f483ae21d26ab336f9
  Author: Nicolas Vasilache <ntv at google.com>
  Date:   2020-04-10 (Fri, 10 Apr 2020)

  Changed paths:
    M mlir/docs/Dialects/Linalg.md
    M mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
    M mlir/include/mlir/IR/AffineExpr.h
    M mlir/lib/IR/AffineExpr.cpp
    M mlir/test/CMakeLists.txt
    M mlir/test/lit.cfg.py
    A mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc
    M mlir/tools/CMakeLists.txt
    A mlir/tools/mlir-linalg-ods-gen/CMakeLists.txt
    A mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

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