Add a generic Linalg op

This CL introduces a linalg.generic op to represent generic tensor contraction operations on views.

A linalg.generic operation requires a numbers of attributes that are sufficient to emit the computation in scalar form as well as compute the appropriate subviews to enable tiling and fusion.

These attributes are very similar to the attributes for existing operations such as linalg.matmul etc and existing operations can be implemented with the generic form.

In the future, most existing operations can be implemented using the generic form.

This CL starts by splitting out most of the functionality of the linalg::NInputsAndOutputs trait into a ViewTrait that queries the per-instance properties of the op. This allows using the attribute informations.

This exposes an ordering of verifiers issue where ViewTrait::verify uses attributes but the verifiers for those attributes have not been run. The desired behavior would be for the verifiers of the attributes specified in the builder to execute first but it is not the case atm. As a consequence, to emit proper error messages and avoid crashing, some of the
linalg.generic methods are defensive as such:
```
    unsigned getNumInputs() {
      // This is redundant with the `n_views` attribute verifier but ordering of verifiers
      // may exhibit cases where we crash instead of emitting an error message.
      if (!getAttr("n_views") || n_views().getValue().size() != 2)
        return 0;
```

In pretty-printed form, the specific attributes required for linalg.generic are factored out in an independent dictionary named "_". When parsing its content is flattened and the "_name" is dropped. This allows using aliasing for reducing boilerplate at each linalg.generic invocation while benefiting from the Tablegen'd verifier form for each named attribute in the dictionary.

For instance, implementing linalg.matmul in terms of linalg.generic resembles:

```
func @mac(%a: f32, %b: f32, %c: f32) -> f32 {
  %d = mulf %a, %b: f32
  %e = addf %c, %d: f32
  return %e: f32
}
#matmul_accesses = [
  (m, n, k) -> (m, k),
  (m, n, k) -> (k, n),
  (m, n, k) -> (m, n)
]
#matmul_trait = {
  doc = "C(m, n) += A(m, k) * B(k, n)",
  fun = @mac,
  indexing_maps = #matmul_accesses,
  library_call = "linalg_matmul",
  n_views = [2, 1],
  n_loop_types = [2, 1, 0]
}
```

And can be used in multiple places as:
```
  linalg.generic #matmul_trait %A, %B, %C [other-attributes] :
    !linalg.view<?x?xf32>, !linalg.view<?x?xf32>, !linalg.view<?x?xf32>
```

In the future it would be great to have a mechanism to alias / register a new
linalg.op as a pair of linalg.generic, #trait.

Also, note that with one could theoretically only specify the `doc` string and parse all the attributes from it.

PiperOrigin-RevId: 261338740

Author: Nicolas Vasilache

include/mlir/AffineOps/AffineOps.td

@@ -1,4 +1,4 @@
-//===- Ops.td - Affine operation definitions ---------------*- tablegen -*-===//
+//===- AffineOps.td - Affine operation definitions ---------*- tablegen -*-===//
 //
 // Copyright 2019 The MLIR Authors.
 //
@@ -28,6 +28,8 @@
 include "mlir/IR/OpBase.td"
 #endif // OP_BASE
 
+include "mlir/AffineOps/AffineOpsBase.td"
+
 def Affine_Dialect : Dialect {
   let name = "affine";
   let cppNamespace = "";

include/mlir/AffineOps/AffineOpsBase.td

@@ -0,0 +1,44 @@
+//===- AffineOpsBase.td - Affine operation definitions -----*- tablegen -*-===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// Defines base support for MLIR affine operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifdef AFFINE_OPS_BASE
+#else
+#define AFFINE_OPS_BASE
+
+#ifdef OP_BASE
+#else
+include "mlir/IR/OpBase.td"
+#endif // OP_BASE
+
+// Attributes containing affine maps.
+def AffineMapAttr : Attr<
+    CPred<"$_self.isa<AffineMapAttr>()">, "AffineMap attribute"> {
+  let storageType = [{ AffineMapAttr }];
+  let returnType = [{ AffineMap }];
+  let constBuilderCall = "$_builder.getAffineMapAttr($0)";
+}
+
+def AffineMapArrayAttr : TypedArrayAttrBase<AffineMapAttr,
+                                      "AffineMap array attribute"> {
+  let constBuilderCall = "$_builder.getAffineMapArrayAttr($0)";
+}
+
+#endif // AFFINE_OPS_BASE

include/mlir/EDSC/Intrinsics.h

@@ -190,6 +190,7 @@ using alloc = ValueBuilder<AllocOp>;
 using affine_apply = ValueBuilder<AffineApplyOp>;
 using affine_load = ValueBuilder<AffineLoadOp>;
 using affine_store = OperationBuilder<AffineStoreOp>;
+using call = OperationBuilder<mlir::CallOp>;
 using constant_float = ValueBuilder<ConstantFloatOp>;
 using constant_index = ValueBuilder<ConstantIndexOp>;
 using constant_int = ValueBuilder<ConstantIntOp>;

include/mlir/IR/AffineMap.h

@@ -153,11 +153,12 @@ AffineMap simplifyAffineMap(AffineMap map);
 
 /// Returns a map of codomain to domain dimensions such that the first codomain
 /// dimension for a particular domain dimension is selected.
-/// Returns an empty map if the input map is empty.
+/// Returns an empty map if the input map is empty or if `map` is not invertible
+/// (i.e. `map` does not contain a subset that is a permutation of full domain
+/// rank).
 ///
 /// Prerequisites:
-///   1. `map` must contain a subset that is a permutation of full domain rank.
-///   2. `map` has no symbols.
+///   1. `map` has no symbols.
 ///
 /// Example 1:
 ///

include/mlir/IR/Builders.h

@@ -130,9 +130,11 @@ class Builder {
   FloatAttr getF16FloatAttr(float value);
   FloatAttr getF32FloatAttr(float value);
   FloatAttr getF64FloatAttr(double value);
+
   IntegerAttr getI32IntegerAttr(int32_t value);
   IntegerAttr getI64IntegerAttr(int64_t value);
 
+  ArrayAttr getAffineMapArrayAttr(ArrayRef<AffineMap> values);
   ArrayAttr getI32ArrayAttr(ArrayRef<int32_t> values);
   ArrayAttr getI64ArrayAttr(ArrayRef<int64_t> values);
   ArrayAttr getF32ArrayAttr(ArrayRef<float> values);

include/mlir/Linalg/IR/LinalgLibraryOps.td

@@ -20,12 +20,13 @@
 //
 //===----------------------------------------------------------------------===//
 
-include "mlir/Linalg/IR/LinalgBase.td"
-
 #ifdef LINALG_LIBRARY_OPS
 #else
 #define LINALG_LIBRARY_OPS
 
+include "mlir/AffineOps/AffineOpsBase.td"
+include "mlir/Linalg/IR/LinalgBase.td"
+
 class LinalgParametricNativeOpTrait<string prop, string parameters> :
   NativeOpTrait<"linalg::" # prop # parameters>
 {}
@@ -65,29 +66,28 @@ class ViewRanks<list<int> ranks> :
 LinalgParametricIntNativeOpTrait<"ViewRanks", ranks>
 {}
 
+def ViewTraits : NativeOpTrait<"linalg::ViewTraits">;
+
 // Base Tablegen class for Linalg ops.
 // Linalg ops that correspond to library calls operate on linalg::View as their
 // first operands. These may be optionally followed by non-view operands
 // depending on the specific Linalg op.
-class LinalgLibrary_Op<string mnemonic, list<OpTrait> props>
-  : Op<Linalg_Dialect, mnemonic, props> {
+class LinalgLibraryBase_Op<string mnemonic, list<OpTrait> props>
+  : Op<Linalg_Dialect, mnemonic, !listconcat(props, [ViewTraits])> {
   let parser = [{ return parseLinalgLibraryOp(parser, result); }];
   let printer = [{ printLinalgLibraryOp(p, *this); }];
+}
 
-  let extraClassDeclaration = [{
-    static StringRef getLibraryCallName() {
+class LinalgLibrary_Op<string mnemonic, list<OpTrait> props>
+  : LinalgLibraryBase_Op<mnemonic, props> {
+
+  code classDeclaration = [{
+    StringRef getLibraryCallName() {
       return "linalg_}] # mnemonic # [{";
     }
   }];
 }
 
-def AffineMapAttr : Attr<
-    CPred<"$_self.isa<AffineMapAttr>()">, "AffineMap attribute"> {
-  let storageType = [{ AffineMapAttr }];
-  let returnType = [{ AffineMap }];
-  let constBuilderCall = "$_builder.getAffineMapAttr($0)";
-}
-
 ////////////////////////////////////////////////////////////////////////////////
 // Concrete Linalg ops.
 ////////////////////////////////////////////////////////////////////////////////
@@ -138,7 +138,7 @@ def CopyOp : LinalgLibrary_Op<"copy", [NInputsAndOutputs<1, 1>]> {
     return build(
       builder, result, input, output, AffineMapAttr(), AffineMapAttr());
   }]>];
-  let extraClassDeclaration = [{
+  let extraClassDeclaration = classDeclaration # [{
     unsigned getNumParallelLoops() {
       auto *view = *(getOperands().begin());
       return view->getType().cast<ViewType>().getRank();
@@ -151,7 +151,7 @@ def CopyOp : LinalgLibrary_Op<"copy", [NInputsAndOutputs<1, 1>]> {
 
 def FillOp : LinalgLibrary_Op<"fill", [NInputsAndOutputs<0, 1>]> {
   let arguments = (ins View, AnyTypeOf<[AnyFloat, AnyInteger, AnyVector]>);
-  let extraClassDeclaration = [{
+  let extraClassDeclaration = classDeclaration # [{
     unsigned getNumParallelLoops() {
       auto *view = *(getOperands().begin());
       return view->getType().cast<ViewType>().getRank();
@@ -170,20 +170,23 @@ def DotOp : LinalgLibrary_Op<"dot",
                              NLoopTypes<0, 1, 0>,
                              ViewRanks<[1, 1, 0]>]> {
   let arguments = (ins View, View, View);
+  let extraClassDeclaration = classDeclaration;
 }
 
 def MatvecOp : LinalgLibrary_Op<"matvec",
                                   [NInputsAndOutputs<2, 1>,
                                    NLoopTypes<1, 1, 0>,
                                    ViewRanks<[2, 1, 1]>]> {
   let arguments = (ins View, View, View);
+  let extraClassDeclaration = classDeclaration;
 }
 
 def MatmulOp : LinalgLibrary_Op<"matmul",
                                   [NInputsAndOutputs<2, 1>,
                                    NLoopTypes<2, 1, 0>,
                                    ViewRanks<[2, 2, 2]>]> {
   let arguments = (ins View, View, View);
+  let extraClassDeclaration = classDeclaration;
 }
 
 def ConvOp : LinalgLibrary_Op<"conv", [NInputsAndOutputs<2, 1>]> {
@@ -208,7 +211,7 @@ def ConvOp : LinalgLibrary_Op<"conv", [NInputsAndOutputs<2, 1>]> {
   let arguments = (ins View:$filter, View:$input, View:$output,
                    OptionalAttr<I64ArrayAttr>:$strides,
                    OptionalAttr<I64ArrayAttr>:$dilations);
-  let extraClassDeclaration = [{
+  let extraClassDeclaration = classDeclaration # [{
     // TODO(ntv) extend to support more than 1 dimensions and potentially
     // grouping too.
     unsigned getNumBatchDimensions() { return 1; }
@@ -248,4 +251,163 @@ def ConvOp : LinalgLibrary_Op<"conv", [NInputsAndOutputs<2, 1>]> {
   let verifier = [{ return ::verify(*this); }];
 }
 
+def GenericOp : LinalgLibraryBase_Op<"generic", []> {
+  let description = [{
+    Generic Linalg op form where the key properties of the computation are
+    specified as attributes. In pretty form, a linalg.generic op is written as:
+
+      ```
+        linalg.generic #trait_attribute %A, %B, %C {other-attributes} :
+          !linalg.view<?x?xf32>, !linalg.view<?x?xf32>, !linalg.view<?x?xf32>
+      ```
+
+    Where #trait_attributes is an alias of a dictionary attribute containing:
+      - doc [optional]: a documentation string
+      - fun: a SymbolRefAttr that must resolve to an existing function symbol.
+        To support inplace updates in a generic fashion, the signature of the
+        function must be:
+        ```
+          fun([input views element types], [output views element types])
+            -> ([output views element types])
+        ```
+      - indexing_maps: a list of AffineMapAttr, one AffineMapAttr per each input
+        and output view. Such AffineMapAttr specifies the mapping between the
+        loops and the indexing within each view.
+      - library_call [optional]: a StringAttr containing the name of an
+        external library function that the linalg.generic operation maps to.
+        The external library is assumed to be dynamically linked and no strong
+        compile-time guarantees are provided. In the absence of such a library
+        call, linalg.generic will always lower to loops.
+      - n_loops: a triple of I64Attr representing the number of enclosing
+        [parallel, reduction, window] loops respectively.
+      - n_views: a pair of I64Attr representing the number of input (readonly)
+        and output (readwrite) views.
+
+    Example:
+    Defining a #matmul_trait attribute in MLIR can be done as follows:
+      ```
+        func @fma(%a: f32, %b: f32, %c: f32) -> f32 {
+          %d = mulf %a, %b: f32
+          %e = addf %c, %d: f32
+          return %e: f32
+        }
+        #matmul_accesses = [
+          (m, n, k) -> (m, k),
+          (m, n, k) -> (k, n),
+          (m, n, k) -> (m, n)
+        ]
+        #matmul_trait = {
+          doc = "C(m, n) += A(m, k) * B(k, n)",
+          fun = @fma,
+          indexing_maps = #matmul_accesses,
+          library_call = "linalg_matmul",
+          n_views = [2, 1],
+          n_loop_types = [2, 1, 0]
+        }
+      ```
+
+    And can be reused in multiple places as:
+      ```
+        linalg.generic #matmul_trait %A, %B, %C [other-attributes] :
+          !linalg.view<?x?xf32>, !linalg.view<?x?xf32>, !linalg.view<?x?xf32>
+      ```
+
+    This may lower to either:
+      ```
+        call @linalg_matmul(%A, %B, %C) :
+          (!linalg.view<?x?xf32>, !linalg.view<?x?xf32>, !linalg.view<?x?xf32>)
+          -> ()
+      ```
+
+    or IR resembling:
+    ```
+    loop.for %m = %c0 to %M step %c1 {
+      loop.for %n = %c0 to %N step %c1 {
+        loop.for %k = %c0 to %K step %c1 {
+          %a = linalg.load %A[%m, %k] : !linalg.view<?x?xf32>
+          %b = linalg.load %B[%k, %n] : !linalg.view<?x?xf32>
+          %c = linalg.load %C[%m, %n] : !linalg.view<?x?xf32>
+          %d = call @mac(%a, %b, %c) : (f32, f32, f32) -> (f32)
+          linalg.store %d, %C[%m, %n] : !linalg.view<?x?x?xf32>
+        }
+      }
+    }
+    ```
+  }];
+  let arguments = (ins Variadic<View>:$views,
+                   SymbolRefAttr:$fun,
+                   AffineMapArrayAttr:$indexing_maps,
+                   I64ArrayAttr:$n_loop_types,
+                   I64ArrayAttr:$n_views,
+                   OptionalAttr<StrAttr>:$doc,
+                   OptionalAttr<StrAttr>:$library_call);
+  let extraClassDeclaration = [{
+    SmallVector<StringRef, 8> linalgTraitAttrNames() {
+      return SmallVector<StringRef, 8>{
+        "doc", "fun", "indexing_maps", "library_call", "n_loop_types", "n_views"
+      };
+    }
+    unsigned getNumInputs() {
+      if (!getAttr("n_views") || n_views().getValue().size() != 2)
+        return 0;
+      auto val = n_views().getValue()[0].cast<IntegerAttr>().getValue();
+      assert(val.getSExtValue() >= 0);
+      return val.getZExtValue();
+    }
+    unsigned getNumOutputs() {
+      if (!getAttr("n_views") || n_views().getValue().size() != 2)
+        return 0;
+      auto val = n_views().getValue()[1].cast<IntegerAttr>().getValue();
+      assert(val.getSExtValue() >= 0);
+      return val.getZExtValue();
+    }
+    unsigned getNumParallelLoops() {
+      if (!getAttr("n_loop_types") || n_loop_types().getValue().size() != 3)
+        return 0;
+      auto val = n_loop_types().getValue()[0].cast<IntegerAttr>().getValue();
+      assert(val.getSExtValue() >= 0);
+      return val.getZExtValue();
+    }
+    unsigned getNumReductionLoops() {
+      if (!getAttr("n_loop_types") || n_loop_types().getValue().size() != 3)
+        return 0;
+      auto val = n_loop_types().getValue()[1].cast<IntegerAttr>().getValue();
+      assert(val.getSExtValue() >= 0);
+      return val.getZExtValue();
+    }
+    unsigned getNumWindowLoops() {
+      if (!getAttr("n_loop_types") || n_loop_types().getValue().size() != 3)
+        return 0;
+      auto val = n_loop_types().getValue()[2].cast<IntegerAttr>().getValue();
+      assert(val.getSExtValue() >= 0);
+      return val.getZExtValue();
+    }
+    unsigned getNumLoops() {
+      return getNumParallelLoops() + getNumReductionLoops() +
+        getNumWindowLoops();
+    }
+    StringRef getFunName() {
+      return fun();
+    }
+    StringRef getLibraryCallName() {
+      return library_call().hasValue() ? library_call().getValue() : "";
+    }
+    AffineMap getIndexingMap(unsigned i) {
+      assert(i < getNumInputsAndOutputs());
+      return indexing_maps().getValue()[i].cast<AffineMapAttr>().getValue();
+    }
+    AffineMap getInputIndexingMap(unsigned i) {
+      assert(i < getNumInputs());
+      return indexing_maps().getValue()[i].cast<AffineMapAttr>().getValue();
+    }
+    AffineMap getOutputIndexingMap(unsigned i) {
+      assert(i < getNumOutputs());
+      return indexing_maps().getValue()[i + getNumInputs()]
+          .cast<AffineMapAttr>().getValue();
+    }
+  }];
+  let printer = [{ return ::print(p, *this); }];
+  let verifier = [{ return ::verify(*this); }];
+  let parser = [{ return ::parse$cppClass(parser, result); }];
+}
 #endif // LINALG_LIBRARY_OPS

include/mlir/Linalg/IR/LinalgOps.h

@@ -18,6 +18,7 @@
 #ifndef MLIR_LINALG_LINALGOPS_H_
 #define MLIR_LINALG_LINALGOPS_H_
 
+#include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Linalg/IR/LinalgTraits.h"