Improve Linalg documentation following the Structured Ops presentation.

PiperOrigin-RevId: 284291653

3 files changed, 130 insertions, 46 deletions

diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td
index ca1b58d7961..edc81250aae 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td
@@ -27,29 +27,96 @@ include "mlir/IR/OpBase.td"
 def Linalg_Dialect : Dialect {
   let name = "linalg";
   let description = [{
-    The Linalg dialect groups together a set of types and operations that are
-    useful to implement a "linear algebra"-like abstraction where ops can lower
-    to scalar load/store and operations or to more general library calls.
+    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.
 
-    The Linalg dialect adopts a convention that is similar to BLAS when
+    The `linalg` dialect manipulates the following types and operations:
+
+    ### Core data types and special ops.
+
+    The following abstractions are used by the `linalg` dialect:
+
+    #### Views
+    The current implementation uses the strided memref abstraction. In the
+    future other abstractions than strided memref will be used.
+
+    #### `!linalg.range`
+    This data type is currently just a triple (`min`,`max`, `step`) that does
+    not pass function boundaries.
+
+    #### `linalg.yield`
+    This op is used as a terminator within the appropriate `linalg` regions.
+
+    In the future, richer `view` and `range` representations are expected, in
+    particular to represent sparse traversals.
+
+    ### Metadata Ops
+    A set of ops that manipulate metadata but do not move memory. These ops take
+    `view` operands + extra attributes and return new `view`s. The returned
+    `view`s generally alias the operand `view`. At the moment the existing ops
+    are:
+
+        * `std.view`,
+        * `std.subview`,
+        * `linalg.range`,
+        * `linalg.slice`,
+        * `linalg.transpose`.
+
+    Future ops are added on a per-need basis but should include:
+
+        * `linalg.reshape`,
+        * `linalg.tile`,
+        * `linalg.intersection`,
+        * `linalg.convex_union`,
+        * `linalg.difference` (would need to work on a list of views).
+
+    ### Payload Ops
+    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.
+
+    Additionally, `linalg` provides some common named operations:
+
+        * `linalg.copy`,
+        * `linalg.fill`,
+        * `linalg.dot`,
+        * `linalg.matmul`,
+        * `linalg.conv`.
+
+    Future ops are added on a per-need basis but should include:
+
+        * `linalg.pad`.
+
+    In an ideal world, all the named ops would be automatically generated from
+    a description in terms of only the `2` generic ops. Unfortunately we do not
+    have such support yet (contributions are most welcome).
+
+    ### Convention for external library interop
+    The `linalg` dialect adopts a convention that is similar to `BLAS` when
     offloading operations to fast library implementations: pass a non-owning
     pointer to input and output data with additional metadata. This convention
-    is also found in libraries such as MKL, OpenBLAS, cuBLAS, cuDNN, etc.. and
-    more generally at interface points across language boundaries (e.g. C++ /
-    Python).
+    is also found in libraries such as `MKL`, `OpenBLAS`, `BLIS`, `cuBLAS`,
+    `cuDNN`, etc.. and more generally at interface points across language
+    boundaries (e.g. C++ / Python).
 
-    Generally, Linalg passes non-owning pointers to strided memref data
+    Generally, `linalg` passes non-owning pointers to strided memref data
     structures to precompiled library calls linked externally. The name `view`
-    is used interchangeably in Linalg to signify strided memref.
+    is used interchangeably in `linalg` to signify strided memref discussed at
+    length in the [strided memref RFC](
+    https://groups.google.com/a/tensorflow.org/g/mlir/c/MaL8m2nXuio/m/a_v07o9yBwAJ).
   }];
 }
 
-// Whether a type is a BufferType.
-def LinalgIsBufferTypePred : CPred<"$_self.isa<BufferType>()">;
-def Buffer : Type<LinalgIsBufferTypePred, "buffer">;
-
 // Whether a type is a RangeType.
 def LinalgIsRangeTypePred : CPred<"$_self.isa<RangeType>()">;
 def Range : Type<LinalgIsRangeTypePred, "range">;
 
+// TODO(ntv): inject the doc for LinalgLibraryOps.td here.
+
 #endif // LINALG_BASE
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgLibraryOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgLibraryOps.td
index afaf039ffd5..74830d50b35 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgLibraryOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgLibraryOps.td
@@ -318,6 +318,7 @@ def ConvOp : LinalgLibrary_Op<"conv", [NInputsAndOutputs<2, 1>]> {
                        q]
     ```
   }];
+
   // TODO(ntv) padding.
   // Following the TF source of truth above, strides and dilations are integer
   // attributes of the same rank as the number of window dimensions.
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td
index 414a5869f72..7e4c0b3def6 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td
@@ -39,18 +39,21 @@ class Linalg_Op<string mnemonic, list<OpTrait> traits = []> :
   let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
-def RangeOp :
+def Linalg_RangeOp :
     Linalg_Op<"range", [NoSideEffect]>,
     Arguments<(ins Index:$min, Index:$max, Index:$step)>,
     Results<(outs Range)> {
-  let summary = "Create a range type value, used to create views";
+  let summary = "Create a `range` type value, used to create `view`s";
   let description = [{
-    The `linalg.range` op creates a linalg.range from 3 values of type `index`
-    that represent the min, max and step values of the range.
+    The `linalg.range` op creates a `!linalg.range` from 3 values of type
+    `index` that represent the min, max and step values of the `range`. This
+    type does not pass function boundaries at the moment.
 
     Example:
 
+    ```mlir
       %3 = linalg.range %0:%1:%2 : !linalg.range
+    ````
   }];
   let builders = [OpBuilder<
     "Builder *builder, OperationState &result, Value *min, Value *max, "
@@ -64,40 +67,48 @@ def RangeOp :
   let verifier = ?;
 }
 
-def SliceOp : Linalg_Op<"slice", [NoSideEffect]>,
+def Linalg_SliceOp : Linalg_Op<"slice", [NoSideEffect]>,
     Arguments<(ins AnyStridedMemRef:$view, Variadic<AnyTypeOf<[Range, Index]>>:$indexings)>,
     Results<(outs AnyStridedMemRef)> {
-  let summary = "Produce a linalg.view which is a subview of a base view.";
+  let summary = "Produce a rank-reduced `subview` of a base `view`.";
   let description = [{
-    The "linalg.slice" op produces a linalg.view which is a subview of a given
-    base view. This allows defining a subregion within the underlying buffer to
-    operate on only a subset of the buffer.
+    The `linalg.slice` op allows defining a subregion of a smaller rank than the
+    operand `view` within the underlying buffer.
 
-    A "linalg.slice" op takes a view and a variadic number of indexings and
-    produces a linalg.view of the same elemental type. An indexing is either:
-      1. a linalg.range, in which case it does not reduce the rank of the parent
-         view.
-      2. an index, in which case it reduces the rank of the parent view by one.
+    A `linalg.slice` op takes a view and a variadic number of indexings and
+    produces a `view` of the same elemental type. An indexing is either:
+      1. a `linalg.range`, in which case it does not reduce the rank of the
+         parent `view` along the corresponding dimension.
+      2. an `index`, in which case it reduces the rank of the parent view by
+         one.
 
-    If an indexing extends past the size of the view, the slice operation
-    automatically truncates it to be within the bounds.
+    If an indexing extends past the size of the `view`, this is undefined
+    behavior. Ideally the `linalg.slice` operation would automatically truncate
+    it to be within bounds but there are tradeoffs involved now that `std.view`
+    is a standard op.
 
     Examples:
 
-      1. rank-preserving slice:
+      1. rank-preserving `slice`:
 
-        %4 = linalg.slice %0[%1, %2] : memref<?x?xf32, stride_specification>,
-          !linalg.range, !linalg.range, memref<?x?xf32, stride_specification>
+      ```mlir
+        %4 = linalg.slice %0[%1, %2] : memref<?x?xf32, stride_spec>,
+          !linalg.range, !linalg.range, memref<?x?xf32, stride_spec>
+       ```
 
-      2. rank-reducing slice (from 2-D to 1-D):
+      2. rank-reducing `slice` (from 2-D to 1-D):
 
-        %4 = linalg.slice %0[%1, %2] : memref<?x?xf32, stride_specification>,
-          index, !linalg.range, memref<?x?xf32, stride_specification>
+      ```mlir
+        %4 = linalg.slice %0[%1, %2] : memref<?x?xf32, stride_spec>,
+          index, !linalg.range, memref<?x?xf32, stride_spec>
+      ```
 
-      3. rank-reducing slice (from 2-D to 0-D):
+      3. rank-reducing `slice` (from 2-D to 0-D):
 
-        %4 = linalg.slice %0[%1, %2] : memref<?x?xf32, stride_specification>,
-          index, index, memref<?x?xf32, stride_specification>
+      ```mlir
+        %4 = linalg.slice %0[%1, %2] : memref<?x?xf32, stride_spec>,
+          index, index, memref<?x?xf32, stride_spec>
+      ```
   }];
 
   let builders = [OpBuilder<
@@ -126,18 +137,20 @@ def SliceOp : Linalg_Op<"slice", [NoSideEffect]>,
   }];
 }
 
-def TransposeOp : Linalg_Op<"transpose", [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 description = [{
-    The "linalg.transpose" op produces a strided memref whose sizes and strides
-    are a permutation of the original. This is a pure metadata transformation.
+    The `linalg.transpose` op produces a strided memref whose sizes and strides
+    are a permutation of the original `view`. This is a pure metadata
+    transformation.
 
     Example:
 
-       %1 = linalg.transpose %0 (i, j) -> (j, i) :
-        memref<?x?xf32, stride_specification>
+    ```mlir
+       %1 = linalg.transpose %0 (i, j) -> (j, i) : memref<?x?xf32, stride_spec>
+    ```
   }];
 
   let builders = [OpBuilder<
@@ -158,16 +171,19 @@ def TransposeOp : Linalg_Op<"transpose", [NoSideEffect]>,
   }];
 }
 
-def YieldOp : Linalg_Op<"yield", [NativeOpTrait<"IsTerminator">]>,
+def Linalg_YieldOp : Linalg_Op<"yield", [NativeOpTrait<"IsTerminator">]>,
     Arguments<(ins Variadic<AnyType>:$values)> {
   let summary = "Linalg yield operation";
   let description = [{
-    "linalg.yield" is a special terminator operation for blocks inside regions
-    in linalg ops. It returns values to the immediately enclosing linalg op.
+    `linalg.yield` is a special terminator operation for blocks inside regions
+    in `linalg` generic ops. It returns values to the immediately enclosing
+    `linalg` generic op.
 
     Example:
 
+    ```mlir
        linalg.yield %f0, %f1 : f32, f32
+    ```
   }];
 }