[mlir][Linalg] Extend generic ops to allow tensors

    Summary:
    This diff adds support to allow `linalg.generic` and
    `linalg.indexed_generic` to take tensor input and output
    arguments.

    The subset of output tensor operand types must appear
    verbatim in the result types after an arrow. The parser,
    printer and verifier are extended to accomodate this
    behavior.

    The Linalg operations now support variadic ranked tensor
    return values. This extension exhibited issues with the
    current handling of NativeCall in RewriterGen.cpp. As a
    consequence, an explicit cast to `SmallVector<Value, 4>`
    is added in the proper place to support the new behavior
    (better suggestions are welcome).

    Relevant cleanups and name uniformization are applied.

    Relevant invalid and roundtrip test are added.

    Reviewers: mehdi_amini, rriddle, jpienaar, antiagainst, ftynse

    Subscribers: burmako, shauheen, llvm-commits

    Tags: #llvm

    Differential Revision: https://reviews.llvm.org/D72022

Author: nicolasvasilache

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

@@ -19,9 +19,10 @@ def Linalg_Dialect : Dialect {
   let name = "linalg";
   let description = [{
     The `linalg` dialect groups together a set of types, operations and
-    transformations that are useful to implement a structured abstraction where
-    ops can lower to scalar load/store and operations or to more general library
-    calls.
+    transformations that are useful to implement a structured abstraction on
+    buffers and tensors. These abstractions are useful for transformations and
+    can lower to scalar load/store and other operations or to more general
+    library calls.
 
     The `linalg` dialect manipulates the following types and operations:
 
@@ -67,12 +68,13 @@ def Linalg_Dialect : Dialect {
     A set of payload carrying operations that implement the [structured ops](
     https://docs.google.com/presentation/d/1P-j1GrH6Q5gLBjao0afQ-GfvcAeF-QU4GXXeSy0eJ9I/edit#slide=id.p
     )
-    abstraction on buffers. `linalg` has `2` generic operations `linalg.generic`
-    and `linalg.indexed_generic` for expressing custom operations. This is
-    subject to further evolution as transformations and analyses continue to be
-    developed.
+    abstraction on tensors and buffers. `linalg` has `2` generic operations
+    `linalg.generic` and `linalg.indexed_generic` for expressing custom
+    operations.
+    This is subject to further evolution as transformations and analyses
+    continue to be developed.
 
-    Additionally, `linalg` provides some common named operations:
+    Additionally, `linalg` provides some commonly named operations:
 
         * `linalg.copy`,
         * `linalg.fill`,

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

@@ -59,7 +59,8 @@ def Linalg_RangeOp :
 }
 
 def Linalg_SliceOp : Linalg_Op<"slice", [NoSideEffect]>,
-    Arguments<(ins AnyStridedMemRef:$view, Variadic<AnyTypeOf<[Range, Index]>>:$indexings)>,
+    Arguments<(ins AnyStridedMemRef:$view,
+                   Variadic<AnyTypeOf<[Range, Index]>>:$indexings)>,
     Results<(outs AnyStridedMemRef)> {
   let summary = "Produce a rank-reduced `subview` of a base `view`.";
   let description = [{
@@ -108,11 +109,11 @@ def Linalg_SliceOp : Linalg_Op<"slice", [NoSideEffect]>,
 
   let extraClassDeclaration = [{
     enum { FirstIndexingOperand = 1 };
-    unsigned getRank() { return getViewType().getRank(); }
-    Type getElementType() { return getViewType().getElementType(); }
-    MemRefType getViewType() { return getType().cast<MemRefType>(); }
+    unsigned getRank() { return getShapedType().getRank(); }
+    Type getElementType() { return getShapedType().getElementType(); }
+    ShapedType getShapedType() { return getType().cast<ShapedType>(); }
     unsigned getBaseViewRank() { return getBaseViewType().getRank(); }
-    MemRefType getBaseViewType() { return view()->getType().cast<MemRefType>(); }
+    ShapedType getBaseViewType() { return view()->getType().cast<ShapedType>();}
 
     // Get the underlying indexing at a given rank.
     Value indexing(unsigned rank) { return *(indexings().begin() + rank); }
@@ -131,7 +132,7 @@ def Linalg_SliceOp : Linalg_Op<"slice", [NoSideEffect]>,
 def Linalg_TransposeOp : Linalg_Op<"transpose", [NoSideEffect]>,
     Arguments<(ins AnyStridedMemRef:$view, AffineMapAttr:$permutation)>,
     Results<(outs AnyStridedMemRef)> {
-  let summary = "transpose operation produces a new strided memref (metadata-only)";
+  let summary = "`transpose` produces a new strided memref (metadata-only)";
   let description = [{
     The `linalg.transpose` op produces a strided memref whose sizes and strides
     are a permutation of the original `view`. This is a pure metadata
@@ -151,14 +152,14 @@ def Linalg_TransposeOp : Linalg_Op<"transpose", [NoSideEffect]>,
   let verifier = [{
     if (!permutation().isPermutation())
       return emitOpError("expected a permutation map");
-    if (permutation().getNumDims() != getViewType().getRank())
+    if (permutation().getNumDims() != getShapedType().getRank())
       return emitOpError("expected a permutation map of same rank as the view");
     return success();
   }];
 
   let extraClassDeclaration = [{
     static StringRef getPermutationAttrName() { return "permutation"; }
-    MemRefType getViewType() { return view()->getType().cast<MemRefType>(); }
+    ShapedType getShapedType() { return view()->getType().cast<ShapedType>(); }
   }];
 }
 

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

@@ -89,23 +89,32 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
       "Value ", "getOutput", (ins "unsigned":$i)
     >,
     InterfaceMethod<[{
-        Query the index of the given input value, or `None` if the value is not
-        an input.
+        Return the index of the given input value `v`, or `None` if the value is
+        not an input.
       }],
-      "llvm::Optional<unsigned>", "getIndexOfInput", (ins "Value ":$view)
+      "llvm::Optional<unsigned>", "getIndexOfInput", (ins "Value ":$v)
     >,
     InterfaceMethod<[{
         Query the index of the given view value, or `None` if the value is not
-        an view.
+        a view.
       }],
       "llvm::Optional<unsigned>", "getIndexOfOutput", (ins "Value ":$view)
     >,
     InterfaceMethod<[{
-        Query the type of the input view at the given index.
-      }], "MemRefType", "getInputViewType", (ins "unsigned":$i)>,
+        Query the type of the input shape at the given index.
+      }], "ShapedType", "getInputShapedType", (ins "unsigned":$i)>,
     InterfaceMethod<[{
         Query the type of the output view at the given index.
-      }], "MemRefType", "getOutputViewType", (ins "unsigned":$i)>,
+      }], "ShapedType", "getOutputShapedType", (ins "unsigned":$i)>,
+    InterfaceMethod<[{
+        Query whether the op has only MemRef input and outputs.
+      }], "bool", "hasBufferSemantics">,
+    InterfaceMethod<[{
+        Query the subset of input operands that are of ranked tensor type.
+      }], "SmallVector<RankedTensorType, 4>", "getInputTensorTypes">,
+    InterfaceMethod<[{
+        Query the subset of output operands that are of ranked tensor type.
+      }], "SmallVector<RankedTensorType, 4>", "getOutputTensorTypes">,
 
     StaticInterfaceMethod<[{
         Create an operation of the current type with the given ___location,
@@ -340,7 +349,7 @@ def ConvOp : LinalgStructured_Op<"conv", [NInputs<2>, NOutputs<1>]> {
     ArrayAttr iterator_types() {
       // Outer parallel loops are always the number of output dimensions; i.e.
       // [ b, xs, q] in the TF notation above.
-      unsigned nPar = getOutputViewType(0).getRank();
+      unsigned nPar = getOutputShapedType(0).getRank();
       unsigned nRed = getNumInputFeatureDimensions();
       // Window loops are a special kind of reduction that is never tiled or
       // parallelized across; i.e. [zs] in the TF notation above whose number
@@ -374,15 +383,25 @@ def ConvOp : LinalgStructured_Op<"conv", [NInputs<2>, NOutputs<1>]> {
   let verifier = [{ return ::verify(*this); }];
 }
 
+def LinalgOperand: Type<
+  Or<[AnyRankedTensor.predicate, AnyStridedMemRef.predicate]>>;
+
+class LinalgOperandOfRank<int rank>: Type<
+  And<[
+    LinalgOperand.predicate,
+    CPred<"$_self.cast<ShapedType>().getRank() == " # rank>]
+  >>;
+
 class GenericOpBase<string mnemonic> : LinalgStructuredBase_Op<mnemonic, []> {
-  let arguments = (ins Variadic<AnyStridedMemRef>:$views,
+  let arguments = (ins Variadic<LinalgOperand>:$views,
                    I64Attr:$args_in,
                    I64Attr:$args_out,
                    AffineMapArrayAttr:$indexing_maps,
                    ArrayAttr:$iterator_types,
                    OptionalAttr<StrAttr>:$doc,
                    OptionalAttr<FlatSymbolRefAttr>:$fun,
                    OptionalAttr<StrAttr>:$library_call);
+  let results = (outs Variadic<AnyRankedTensor>:$output_tensors);
   let regions = (region AnyRegion:$region);
   let extraClassDeclaration = [{
     SmallVector<StringRef, 8> linalgTraitAttrNames() {
@@ -511,6 +530,28 @@ def GenericOp : GenericOpBase<"generic"> {
       }
     }
     ```
+
+    To allow progressive lowering from the value world (a.k.a tensor values) to
+    the buffer world (a.k.a memref values), a `linalg.generic` op accepts
+    mixing input and output ranked tensor values with input and output memrefs.
+
+    ```mlir
+      %1 = linalg.generic #trait_attribute %A, %B, %C {other-attributes} :
+        tensor<?x?xf32>,
+        memref<?x?xf32, stride_specification>,
+        tensor<?x?xf32>
+        -> (tensor<?x?xf32>)
+    ```
+
+    In this case, the number of return values must match the number of output
+    tensor arguments. The semantics is that the `linalg.generic` op
+    produces (i.e. allocates and fills) its return values.
+    Tensor values must be legalized by a buffer allocation pass before most
+    transformations can be applied. In particular, transformations that create
+    control flow around linalg.generic operations are not expected to mix with
+    tensors because SSA values do not escape naturally. Still, transformations
+    and rewrites that take advantage of tensor SSA values are expected to be
+    useful and will be added in the near future.
   }];
   let verifier = [{ return ::verify(*this); }];
 }
@@ -555,9 +596,11 @@ def IndexedGenericOp : GenericOpBase<"indexed_generic"> {
     Example:
     Defining a #matmul_trait attribute in MLIR can be done as follows:
       ```mlir
-        func @fma(%i: index, %j: index, %k: index, %a: f32, %b: f32, %c: f32)
+        func @fma(%offset_m: index, %offset_n: index, %offset_k: index,
+                  %a: f32, %b: f32, %c: f32)
           -> f32
         {
+          "some_optional_condition"(%offset_m, %offset_n, %offset_k)
           %d = mulf %a, %b: f32
           %e = addf %c, %d: f32
           return %e: f32
@@ -587,7 +630,7 @@ def IndexedGenericOp : GenericOpBase<"indexed_generic"> {
 
     This may lower to either:
       ```mlir
-        call @linalg_matmul(%A, %B, %C) :
+        call @linalg_matmul(%offset_m, %offset_n, %offset_k, %A, %B, %C) :
           (memref<?x?xf32, stride_specification>,
            memref<?x?xf32, stride_specification>,
            memref<?x?xf32, stride_specification>)
@@ -609,6 +652,29 @@ def IndexedGenericOp : GenericOpBase<"indexed_generic"> {
       }
     }
     ```
+
+    To allow progressive lowering from the value world (a.k.a tensor values) to
+    the buffer world (a.k.a memref values), a `linalg.indexed_generic` op
+    accepts mixing input and output ranked tensor values with input and output
+    memrefs.
+
+    ```mlir
+      %1 = linalg.indexed_generic #trait_attribute %A, %B, %C {other-attributes}
+      : tensor<?x?xf32>,
+        memref<?x?xf32, stride_specification>,
+        tensor<?x?xf32>
+        -> (tensor<?x?xf32>)
+    ```
+
+    In this case, the number of return values must match the number of output
+    tensor arguments. The semantics is that the `linalg.indexed_generic` op
+    produces (i.e. allocates and fills) its return values.
+    Tensor values must be legalized by a buffer allocation pass before most
+    transformations can be applied. In particular, transformations that create
+    control flow around linalg.generic operations are not expected to mix with
+    tensors because SSA values do not escape naturally. Still, transformations
+    and rewrites that take advantage of tensor SSA values are expected to be
+    useful and will be added in the near future.
   }];
   let verifier = [{ return ::verify(*this); }];
 }