aboutsummaryrefslogtreecommitdiff
path: root/llvm/lib/CodeGen/SplitKit.cpp
diff options
context:
space:
mode:
authorNicolas Vasilache <ntv@google.com>2020-04-10 13:54:08 -0400
committerNicolas Vasilache <ntv@google.com>2020-04-10 13:59:25 -0400
commit882ba484743763b8560b08f483ae21d26ab336f9 (patch)
tree6655577bef50793c5acef98a99f65ab2f6100e8a /llvm/lib/CodeGen/SplitKit.cpp
parenta04ab2ec08014d36337d686d43ecef7668622f1b (diff)
downloadllvm-882ba484743763b8560b08f483ae21d26ab336f9.zip
llvm-882ba484743763b8560b08f483ae21d26ab336f9.tar.gz
llvm-882ba484743763b8560b08f483ae21d26ab336f9.tar.bz2
[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
Diffstat (limited to 'llvm/lib/CodeGen/SplitKit.cpp')
0 files changed, 0 insertions, 0 deletions