[PATCH] D77067: [mlir][Linalg] Create a tool to generate named Linalg ops from a Tensor Comprehensions-like specification.

[Nicolas Vasilache via Phabricator via llvm-commits]

nicolasvasilache created this revision.
Herald added subscribers: llvm-commits, Joonsoo, liufengdb, aartbik, lucyrfox, mgester, arpith-jacob, antiagainst, shauheen, burmako, jpienaar, rriddle, mehdi_amini, mgorny.
Herald added a reviewer: rriddle.
Herald added a project: LLVM.

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 }));
  }


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D77067

Files:
  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
  mlir/tools/CMakeLists.txt
  mlir/tools/mlir-linalg-ods-gen/CMakeLists.txt
  mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

-------------- next part --------------
A non-text attachment was scrubbed...
Name: D77067.253607.patch
Type: text/x-patch
Size: 58592 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200330/7690161b/attachment-0001.bin>