[mlir][Linalg][RFC] Add basic plumbing for "named ops" in Linalg

[nicolasvasilache]

Summary

This revision sketches how Linalg ops can be defined declaratively using a language like Tensor Comprehensions.

For now this does not invest more in the infrastructure than the bare minimum to demonstrate end-to-end op generation with Tablegen. At the moment the language is left unspecified and we only do basic string matching to a predefined C++ implementation.

In the longer term, a large fraction of ops can be written declaratively as custom ops with a single mathematical source of truth that resembles:

"C(i,j) +=! A(i,k) * B(k,j)"

[nicolasvasilache]

Remove duplicate assemblyFormat

[nicolasvasilache]

This is now absolete as of: https://reviews.llvm.org/D79135

mlir/include/mlir/Dialect/Linalg/IR/CMakeLists.txt

 add_mlir_dialect(LinalgOps linalg LinalgDoc)
 set(LLVM_TARGET_DEFINITIONS LinalgStructuredOps.td)
 mlir_tablegen(LinalgStructuredOps.h.inc -gen-op-decls)
 mlir_tablegen(LinalgStructuredOps.cpp.inc -gen-op-defs)
+mlir_tablegen(LinalgNamedStructuredOpsInterface.cpp.inc -gen-linalg-named-structured-ops-defs)
 add_public_tablegen_target(MLIRLinalgStructuredOpsIncGen)
 
 set(LLVM_TARGET_DEFINITIONS LinalgStructuredOpsInterface.td)
 mlir_tablegen(LinalgStructuredOpsInterfaces.h.inc -gen-op-interface-decls)
 mlir_tablegen(LinalgStructuredOpsInterfaces.cpp.inc -gen-op-interface-defs)
 add_public_tablegen_target(MLIRLinalgStructuredOpsInterfaceIncGen)

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

let description = [{
     future.
   }];
 
   let verifier = [{ return ::verify(*this); }];
 
   let hasFolder = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// Named Linalg ops, implemented as a declarative configurations of generic ops.
+//===----------------------------------------------------------------------===//
+
+def NamedStructuredOpTraits : NativeOpTrait<"linalg::NamedStructuredOpTraits">;
+
+class LinalgNamedStructured_Op<string mnemonic, list<OpTrait> props>
+  : Op<Linalg_Dialect, mnemonic,
+       !listconcat(props, [StructuredOpTraits, LinalgStructuredInterface])> {
+  string spec = ?;
+  let assemblyFormat = "`(` operands `)` attr-dict `:` "
+    "functional-type(operands, results)";
+}
+
+def PointwiseAddOp : LinalgNamedStructured_Op<"pointwise_add", [
+    NInputs<2>,
+    NOutputs<1>,
+    NamedStructuredOpTraits]> {
+  let arguments = (ins Variadic<LinalgOperand>:$views);
+  let results = (outs Variadic<AnyRankedTensor>:$output_tensors);
+  let spec = "A(...) = B(...) + C(...)";
+}
+
+def PointwiseMulOp : LinalgNamedStructured_Op<"pointwise_mul", [
+    NInputs<2>,
+    NOutputs<1>,
+    NamedStructuredOpTraits]> {
+  let arguments = (ins Variadic<LinalgOperand>:$views);
+  let results = (outs Variadic<AnyRankedTensor>:$output_tensors);
+  let spec = "A(...) = B(...) * C(...)";
+}
+
 #endif // LINALG_STRUCTURED_OPS

mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h

public:
   static LogicalResult verifyTrait(Operation *op) {
     auto nOperands = cast<ConcreteType>(op).getNumInputsAndOutputBuffers();
     if (failed(OpTrait::impl::verifyAtLeastNOperands(op, nOperands)))
       return failure();
     return success();
   }
 };
 
+/// This class provides the API for named Linalg StructuredOps.
+template <typename ConcreteType>
+class NamedStructuredOpTraits
+    : public OpTrait::TraitBase<ConcreteType, NamedStructuredOpTraits> {
+public:
+  llvm::Optional<SmallVector<StringRef, 8>> referenceIterators();
+  llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps();
+  std::function<void(OpBuilder &, Location, ArrayRef<Value>)>
+  emitScalarImplementation();
+};
+
 } // namespace linalg
 } // namespace OpTrait
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_LINALGTRAITS_H_

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

if (auto dilations = op.dilations()) {
     if (failed(verifyStrideOrDilation(op, dilations->getValue(),
                                       /*isStride=*/false)))
       return failure();
   }
   return success();
 }
 
 namespace mlir {
+namespace OpTrait {
+namespace linalg {
+
+// Linalg "named" traits live in the mlir::OpTrait::linalg namespace.
+#include "mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsInterface.cpp.inc"
+
+} // namespace linalg
+} // namespace OpTrait
+} // namespace mlir
+
+namespace mlir {
 namespace linalg {
 
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterfaces.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Linalg/IR/LinalgOps.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
 
 } // namespace linalg
 } // namespace mlir
 
 AffineMap mlir::linalg::extractOrIdentityMap(Optional<AffineMap> maybeMap,
                                              unsigned rank,
                                              MLIRContext *context) {
   if (maybeMap)
     return maybeMap.getValue();
-  if (rank == 0)
-    return AffineMap::get(context);
   return AffineMap::getMultiDimIdentityMap(rank, context);
 }
 
 SmallVector<AffineExpr, 4>
 mlir::linalg::makeAffineDimExprs(unsigned num, unsigned &startIdx,
                                  MLIRContext *context) {
   SmallVector<AffineExpr, 4> res;
   res.reserve(num);

mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp

for (unsigned idx = 0, e = map.getNumResults(); idx < e; ++idx) {
     res.push_back(
         linalg_range(std_constant_index(0), sizes[idx], std_constant_index(1)));
   }
   return res;
 }
 
 namespace {
 template <typename IndexedValueType, typename LinalgOpType>
-class LinalgScopedEmitter {};
+class LinalgScopedEmitter {
+public:
+  static void emitScalarImplementation(ArrayRef<Value> allIvs,
+                                       LinalgOpType linalgOp) {
+    assert(linalgOp.hasBufferSemantics() &&
+           "expected linalg op with buffer semantics");
+    llvm_unreachable("NYI");
+    linalgOp.emitScalarImplementation()(ScopedContext::getBuilder(),
+                                        ScopedContext::getLocation(), allIvs);
+  }
+};
 
 template <typename IndexedValueType>
 class LinalgScopedEmitter<IndexedValueType, CopyOp> {
 public:
   static void emitScalarImplementation(ArrayRef<Value> allIvs, CopyOp copyOp) {
     assert(copyOp.hasBufferSemantics() &&
            "expected linalg op with buffer semantics");
     auto nPar = copyOp.getNumParallelLoops();

mlir/lib/IR/AffineMap.cpp

 }
 
 AffineMap AffineMap::getMultiDimIdentityMap(unsigned numDims,
                                             MLIRContext *context) {
   SmallVector<AffineExpr, 4> dimExprs;
   dimExprs.reserve(numDims);
   for (unsigned i = 0; i < numDims; ++i)
     dimExprs.push_back(mlir::getAffineDimExpr(i, context));
-  return get(/*dimCount=*/numDims, /*symbolCount=*/0, dimExprs);
+  return dimExprs.empty() ?
+      get(/*dimCount=*/0, /*symbolCount=*/0, context) :
+      get(/*dimCount=*/numDims, /*symbolCount=*/0, dimExprs);
 }
 
 MLIRContext *AffineMap::getContext() const { return map->context; }
 
 bool AffineMap::isIdentity() const {
   if (getNumDims() != getNumResults())
     return false;
   ArrayRef<AffineExpr> results = getResults();

mlir/test/Dialect/Linalg/roundtrip.mlir

%2 = linalg.reshape %arg2 [affine_map<(i, j, k) -> (i, j)>,
     memref<?x?x?xf32, offset : ?, strides : [?, ?, 1]> into
     memref<?x?xf32, offset : ?, strides : [?, 1]>
   %r2 = linalg.reshape %2 [affine_map<(i, j, k) -> (i, j)>,
                            affine_map<(i, j, k) -> (k)>] :
     memref<?x?xf32, offset : ?, strides : [?, 1]> into
     memref<?x?x?xf32, offset : ?, strides : [?, ?, 1]>
   return
 }
-// CHECK-LABEL: func @reshape
-//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
-//  CHECK-SAME:     memref<?x?x?xf32> into memref<?x?xf32>
-//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
-//  CHECK-SAME:     memref<?x?xf32> into memref<?x?x?xf32>
-//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
-//  CHECK-SAME:     memref<?x?x?xf32, #[[strided3DOFF0]]> into memref<?x?xf32, #[[strided2DOFF0]]>
-//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
-//  CHECK-SAME:     memref<?x?xf32, #[[strided2DOFF0]]> into memref<?x?x?xf32, #[[strided3DOFF0]]>
-//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
-//  CHECK-SAME:     memref<?x?x?xf32, #[[strided3D]]> into memref<?x?xf32, #[[strided2D]]>
-//       CHECK:   linalg.reshape {{.*}} [#[[reshapeD01]], #[[reshapeD2]]]
-//  CHECK-SAME:     memref<?x?xf32, #[[strided2D]]> into memref<?x?x?xf32, #[[strided3D]]>
+
+func @pointwise1d(%a: memref<?xf32>, %b: memref<?xf32>, %c: memref<?xf32>,
+                  %ta: tensor<?xf32>, %tb: tensor<?xf32>, %tc: tensor<?xf32>) {
+  linalg.pointwise_add(%a, %b, %c): (memref<?xf32>, memref<?xf32>, memref<?xf32>) -> ()
+  %dadd = linalg.pointwise_add(%a, %b): (memref<?xf32>, memref<?xf32>) -> tensor<?xf32>
+  %eadd = linalg.pointwise_add(%ta, %b): (tensor<?xf32>, memref<?xf32>) -> tensor<?xf32>
+  linalg.pointwise_add(%ta, %tb, %c): (tensor<?xf32>, tensor<?xf32>, memref<?xf32>) -> ()
+
+  linalg.pointwise_mul(%a, %b, %c): (memref<?xf32>, memref<?xf32>, memref<?xf32>) -> ()
+  %dmul = linalg.pointwise_mul(%a, %b): (memref<?xf32>, memref<?xf32>) -> tensor<?xf32>
+  %emul = linalg.pointwise_mul(%ta, %b): (tensor<?xf32>, memref<?xf32>) -> tensor<?xf32>
+  linalg.pointwise_mul(%ta, %tb, %c): (tensor<?xf32>, tensor<?xf32>, memref<?xf32>) -> ()
+
+  return
+}
+// CHECK-LABEL: func @pointwise1d
+//       CHECK:   linalg.pointwise_add
+//  CHECK-SAME:     (memref<?xf32>, memref<?xf32>, memref<?xf32>) -> ()
+//       CHECK:   linalg.pointwise_add
+//  CHECK-SAME:     (memref<?xf32>, memref<?xf32>) -> tensor<?xf32>
+//       CHECK:   linalg.pointwise_add
+//  CHECK-SAME:     (tensor<?xf32>, memref<?xf32>) -> tensor<?xf32>
+//       CHECK:   linalg.pointwise_add
+//  CHECK-SAME:     (tensor<?xf32>, tensor<?xf32>, memref<?xf32>) -> ()
+//       CHECK:   linalg.pointwise_mul
+//  CHECK-SAME:     (memref<?xf32>, memref<?xf32>, memref<?xf32>) -> ()
+//       CHECK:   linalg.pointwise_mul
+//  CHECK-SAME:     (memref<?xf32>, memref<?xf32>) -> tensor<?xf32>
+//       CHECK:   linalg.pointwise_mul
+//  CHECK-SAME:     (tensor<?xf32>, memref<?xf32>) -> tensor<?xf32>
+//       CHECK:   linalg.pointwise_mul
+//  CHECK-SAME:     (tensor<?xf32>, tensor<?xf32>, memref<?xf32>) -> ()

mlir/tools/mlir-tblgen/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   MLIRTableGen
   Support
   )
 
 add_tablegen(mlir-tblgen MLIR
   DialectGen.cpp
   EnumsGen.cpp
+  LinalgNamedOpsGen.cpp
   LLVMIRConversionGen.cpp
   LLVMIRIntrinsicGen.cpp
   mlir-tblgen.cpp
   OpDefinitionsGen.cpp
   OpDocGen.cpp
   OpFormatGen.cpp
   OpInterfacesGen.cpp
   RewriterGen.cpp
   SPIRVUtilsGen.cpp
   StructsGen.cpp
   )
 set_target_properties(mlir-tblgen PROPERTIES FOLDER "Tablegenning")

mlir/tools/mlir-tblgen/LinalgNamedOpsGen.cpp

+//===- LinalgNamedOpsGen.cpp - MLIR Linalg op generator -------------------===//
+//
+// Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// LinalgNamedOpsGen uses the description of structured operations in the Linalg
+// dialect to generate op definitions, parsers, pretty-printers and matchers
+// (e.g. linalg.matmul).
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Support/STLExtras.h"
+#include "mlir/TableGen/GenInfo.h"
+#include "mlir/TableGen/OpClass.h"
+#include "mlir/TableGen/Operator.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/Signals.h"
+#include "llvm/TableGen/Error.h"
+#include "llvm/TableGen/Record.h"
+#include "llvm/TableGen/TableGenBackend.h"
+
+using namespace llvm;
+using namespace mlir;
+using namespace mlir::tblgen;
+
+using mlir::tblgen::Operator;
+
+// Emits a function with signature:
+//    `llvm::Optional<SmallVector<StringRef, 8>> referenceIterators()`
+// This is an instance method as opposed to a class method (i.e. a static
+// method) so that it can work in a rank-agnostic fashion.
+// TODO: lang + parser.
+void iteratorTypesBody(const Record &r, StringRef opName,
+                       llvm::raw_ostream &os) {
+  if (r.getValueAsString("spec") == "A(...) = B(...) + C(...)" ||
+      r.getValueAsString("spec") == "A(...) = B(...) * C(...)") {
+    // Rank-polymorphic pointwise iterator types: `getNumParallelLoops` is
+    // retrieved from the generic op's referenceIterators. The
+    // referenceIterators are built by looking at the operands of the op
+    // instance soit can be polymorphic at compiler compile time and static at
+    // MLIR compile time.
+    os << opName << " concreteOp = cast<" << opName << ">(getOperation());";
+    os << "unsigned nPar = concreteOp.getNumParallelLoops();";
+    os << "return SmallVector<StringRef, 8>(nPar, "
+          "getParallelIteratorTypeName());";
+  }
+}
+
+// Emits a instance method with signature:
+//    `llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps()`
+// This is an instance method as opposed to a class method (i.e. a static
+// method) so that it can work in a rank-agnostic fashion.
+// TODO: lang + parser.
+void indexingMapsBody(const Record &r, StringRef opName,
+                      llvm::raw_ostream &os) {
+  if (r.getValueAsString("spec") == "A(...) = B(...) + C(...)" ||
+      r.getValueAsString("spec") == "A(...) = B(...) * C(...)") {
+    // Rank-polymorphic pointwise indexing maps.
+    os << opName << " concreteOp = cast<" << opName << ">(getOperation());";
+    os << R"SPEC(
+      MLIRContext *context = getOperation()->getContext();
+      auto id = AffineMap::getMultiDimIdentityMap(
+        concreteOp.getNumParallelLoops(), context);
+      return SmallVector<AffineMap, 8>{id, id, id};
+    )SPEC";
+  }
+}
+
+// Emits a instance method with signature:
+//    `void emitScalarImplementationBody()`
+// This is an instance method as opposed to a class method (i.e. a static
+// method) so that it can work in a rank-agnostic fashion.
+// TODO: lang + parser.
+void emitScalarImplementationBody(const Record &r, StringRef opName,
+                                  llvm::raw_ostream &os) {
+  if (r.getValueAsString("spec") == "A(...) = B(...) + C(...)" ||
+      r.getValueAsString("spec") == "A(...) = B(...) * C(...)") {
+    // Rank-polymorphic pointwise indexing maps.
+    os << R"SPEC(
+      return [](OpBuilder &, Location, ArrayRef<Value>) {
+        llvm_unreachable("NYI in LinalgNamedOpsGen.cpp");
+      };
+    )SPEC";
+  }
+}
+
+template <typename Lambda>
+void emitMethodWithBody(StringRef traitName, StringRef opName,
+                        StringRef resultType, StringRef funName,
+                        StringRef operands, const Record &r,
+                        llvm::raw_ostream &os, Lambda fun) {
+  os << "template<>\n";
+  os << resultType << " " << traitName << "<" << opName << ">::" << funName
+     << "(" << operands << ") {\n";
+  fun(r, opName, os);
+  os << "\n}\n\n";
+}
+
+bool emitLinalgNamedOpsInterfaceFunctions(const llvm::RecordKeeper &records,
+                                          llvm::raw_ostream &os) {
+  llvm::emitSourceFileHeader("Operations for Linalg Named Structured Ops", os);
+  os << "#include \"mlir/Dialect/Linalg/IR/LinalgTraits.h\"\n\n";
+
+  auto defs = records.getAllDerivedDefinitions("LinalgNamedStructured_Op");
+  for (const llvm::Record *r : defs) {
+    const auto *traitName = "NamedStructuredOpTraits";
+    auto opName = r->getName();
+    const auto *resultType = "llvm::Optional<SmallVector<StringRef, 8>>";
+    const auto *funName = "referenceIterators";
+    emitMethodWithBody(traitName, opName, resultType, funName, /*operands=*/"",
+                       *r, os, iteratorTypesBody);
+
+    resultType = "llvm::Optional<SmallVector<AffineMap, 8>>";
+    funName = "referenceIndexingMaps";
+    emitMethodWithBody(traitName, opName, resultType, funName, /*operands=*/"",
+                       *r, os, indexingMapsBody);
+
+    resultType = "std::function<void(OpBuilder &, Location, ArrayRef<Value>)>";
+    funName = "emitScalarImplementation";
+    emitMethodWithBody(traitName, opName, resultType, funName, /*operands=*/"",
+                       *r, os, emitScalarImplementationBody);
+  }
+
+  return false;
+}
+
+static mlir::GenRegistration
+    genLinalgNamedOps("gen-linalg-named-structured-ops-defs",
+                      "Generate Linalg named StructuredOps",
+                      emitLinalgNamedOpsInterfaceFunctions);