[mlir][sparse] update doc for new surface syntax encoding (#67072)

aartbik · web-flow · commit edefc18c46b0 · 2023-09-21T17:08:11.000-07:00
diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
@@ -14,22 +14,15 @@ include "mlir/IR/EnumAttr.td"
 include "mlir/Dialect/SparseTensor/IR/SparseTensorBase.td"
 include "mlir/IR/TensorEncoding.td"
 
-// All of the Tensor attributes will extend this class.
+// All of the sparse tensor attributes will extend this class.
 class SparseTensor_Attr<string name,
                         list<Trait> traits = []>
     : AttrDef<SparseTensor_Dialect, name, traits>;
 
 //===----------------------------------------------------------------------===//
-// Type aliases.
-//
-// These attributes are just like `IndexAttr` (include/mlir/IR/OpBase.td),
-// except that:
-// (1) the `summary` is more specific (i.e., the fourth parameter to
-//     `TypedAttrBase`), which helps tablegen provide better error messages.
-// (2) tablegen-generated getters will have the given `returnType`, in
-//     lieu of the `APInt` that `IndexAttr` uses.  This avoids the boilerplate
-//     of needing to say `get{FOO}().getZExtValue()`, as well as using
-//     C++ types which better document intent.
+// These attributes are just like `IndexAttr` except that they clarify whether
+// the index refers to a dimension (an axis of the semantic tensor) or a level
+// (an axis of the actual storage format).
 //===----------------------------------------------------------------------===//
 
 def DimensionAttr :
@@ -107,79 +100,71 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
   let mnemonic = "encoding";
 
   let description = [{
-    An attribute to encode TACO-style information on sparsity properties
-    of tensors. The encoding is eventually used by a **sparse compiler**
-    pass to generate sparse code fully automatically for all tensor
-    expressions that involve tensors with a sparse encoding. Compiler
-    passes that run before this sparse compiler pass need to be
-    aware of the semantics of tensor types with such an encoding.
-
-    Each sparse tensor comes equipped with two different sets of axes for
-    describing the tensor's multi-dimensional structure.  We use the term
-    "dimension" to refer to the axes of the semantic tensor itself; whereas,
-    we use the term "level" to refer to the axes of the storage scheme,
-    which is the operational representation of that tensor.  Therefore,
-    the fields of the encoding attribute (further explained below) satisfy
-    the following correspondences:
-
-    - Dimensions:
-        - the shape of the tensor type
-        - the `dimSlices` field
-        - the arguments of the `dimToLvl` field
-    - Levels:
-        - the results of the `dimToLvl` field
-        - the `lvlTypes` field
-
-    The attribute consists of the following fields.
-
-    - Level-type for each level of a tensor type:
-        - **dense** : all entries along this level are stored.
-        - **compressed** : only nonzeros along this level are stored.
-        - **singleton** : a variant of the compressed level-format,
-          for when coordinates are guaranteed to have no siblings at this level.
-      By default, each level-type has the property of being unique (no
-      duplicates at that level) and ordered (coordinates appear sorted
-      at that level).  The following two suffixes can be used to specify
-      that the level should instead be non-unique (duplicates may appear)
-      and/or non-ordered (coordinates may appear unsorted).
-        - **-nu** : not unique
-        - **-no** : not ordered
-      Currently, these suffixes (if present) must appear in this order.
-      In the future, we may introduce additional level-types and
-      properties, and split up how the level-format and properties are
-      specified rather than using this suffix mechanism.
-
-    - An optional affine map from dimension-coordinates to level-coordinates;
-      defaulting to the identity map.  For example, given a 2-d tensor:
-      `(i, j) -> (i, j)` specifies row-wise storage, `(i, j) -> (j, i)`
-      specifies column-wise storage, and
-      `(i, j) -> (i floordiv 2, j floordiv 3, i mod 2, j mod 3)`
-      specifies 2x3 block-sparsity.  For block-sparsity, blocks are typically
-      stored with compression while dense storage is used within each block
-      (although hybrid schemes are possible as well).
-
-      (The following will be corrected in an upcoming change that completely
-      overhauls the syntax of this attribute.)
-
-      The dimToLvl mapping also provides a notion of "counting a
-      dimension", where every stored element with the same coordinate
-      is mapped to a new slice.  For instance, ELL storage of a 2-d
-      tensor can be defined with the mapping `(i, j) -> (#i, i, j)`
-      using the notation of [Chou20].  Lacking the `#` symbol in MLIR's
-      affine mapping, we use a free symbol `c` to define such counting,
-      together with a constant that denotes the number of resulting
-      slices.  For example, the mapping `(i, j)[c] -> (c * 3 * i, i, j)`
-      with the level-types `["dense", "dense", "compressed"]` denotes ELL
-      storage with three jagged diagonals that count the dimension `i`.
-
-    - The required bitwidth for "position" storage (integral offsets
+    An attribute to encode information on sparsity properties of tensors, inspired
+    by the TACO formalization of sparse tensors. This encoding is eventually used
+    by a **sparsifier** pass to generate sparse code fully automatically from a
+    sparsity-agnostic representation of the computation, i.e., an implicit sparse
+    representation is converted to an explicit sparse representation where co-iterating
+    loops operate on sparse storage formats rather than tensors with a sparsity
+    encoding. Compiler passes that run before this sparse compiler pass need to
+    be aware of the semantics of tensor types with such a sparsity encoding.
+
+    In this encoding, we use `dimension` to refer to the axes of the semantic tensor,
+    and `level` to refer to the axes of the actual storage format, i.e., the
+    operational representation of the sparse tensor in memory. The number of
+    dimensions is usually the same as the number of levels (such as CSR storage format).
+    However, the encoding can also map dimensions to higher-order levels (for example,
+    to encode a block-sparse BSR storage format) or to lower-order levels
+    (for example, to linearize dimensions as a single level in the storage).
+
+    The encoding contains a `map` that provides the following:
+
+    - An ordered sequence of dimension specifications, each of which defines:
+      - the dimension-size (implicit from the tensor’s dimension-shape)
+      - a **dimension-expression**
+    - An ordered sequence of level specifications, each of which includes a required
+      **level-type**, which defines how the level should be stored. Each level-type
+      consists of:
+      - a **level-format**
+      - a collection of **level-properties** that apply to the level-format
+      - a **level-expression**, which defines what is stored
+
+    Each level-expression is an affine expression over dimension-variables. Thus, the
+    level-expressions collectively define an affine map from dimension-coordinates to
+    level-coordinates. The dimension-expressions collectively define the inverse map,
+    which only needs to be provided for elaborate cases where it cannot be inferred
+    automatically. Within the sparse storage format, we refer to indices that are
+    stored explicitly as `coordinates` and indices into the storage format as `positions`.
+
+    The supported level-formats are the following:
+
+    - **dense** : all entries along this level are stored
+    - **compressed** : only nonzeros along this level are stored
+    - **singleton** : a variant of the compressed format, where coordinates have no siblings
+
+    Different level-formats may have different collections of level-properties.
+    By default, each level-type has the property of being unique (no duplicate
+    coordinates at that level), ordered (coordinates appear sorted at that
+    level), and, for compression, storing the positions in a compact way where
+    an interval is defined by a lower bound "pos(i)" and an upper bound "pos(i+1)-1".
+    The following properties can be added to a level-format to change this
+    default behavior:
+
+    - **nonunique** : duplicate coordinates may appear at the level
+    - **nonordered** : coordinates may appear in arbribratry order
+    - **high** : the upper bound is stored explicitly in a separate array
+    - **block2_4** : the compression uses a 2:4 encoding per 1x4 block
+
+    In addition to the  `map`, the following two fields are optional:
+
+    - The required bitwidth for `position` storage (integral offsets
       into the sparse storage scheme).  A narrow width reduces the memory
       footprint of overhead storage, as long as the width suffices to
       define the total required range (viz. the maximum number of stored
       entries over all indirection levels).  The choices are `8`, `16`,
       `32`, `64`, or, the default, `0` to indicate the native bitwidth.
 
-    - The required bitwidth for "coordinate" storage (the coordinates
+    - The required bitwidth for `coordinate` storage (the coordinates
       of stored entries).  A narrow width reduces the memory footprint
       of overhead storage, as long as the width suffices to define
       the total required range (viz. the maximum value of each tensor
@@ -194,41 +179,46 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
     ```mlir
     // Sparse vector.
     #SparseVector = #sparse_tensor.encoding<{
-      map = (d0) -> (d0 : compressed)
+      map = (i) -> (i : compressed)
     }>
     ... tensor<?xf32, #SparseVector> ...
 
-    // Sorted Coordinate Scheme.
+    // Sorted coordinate scheme.
     #SortedCOO = #sparse_tensor.encoding<{
-      map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton)
+      map = (i, j) -> (i : compressed(nonunique), j : singleton)
     }>
     ... tensor<?x?xf64, #SortedCOO> ...
 
+    // Batched sorted coordinate scheme, with high encoding.
+    #BCOO = #sparse_tensor.encoding<{
+      map = (i, j, k) -> (i : dense, j : compressed(nonunique, high), k : singleton)
+    }>
+    ... tensor<10x10xf32, #BCOO> ...
+
+    // Compressed sparse row.
+    #CSR = #sparse_tensor.encoding<{
+      map = (i, j) -> (i : dense, j : compressed)
+    }>
+    ... tensor<100x100xbf16, #CSR> ...
+
     // Doubly compressed sparse column storage with specific bitwidths.
     #DCSC = #sparse_tensor.encoding<{
-      map = (d0, d1) -> (d1 : compressed, d0 : compressed),
+      map = (i, j) -> (j : compressed, i : compressed),
       posWidth = 32,
       crdWidth = 8
     }>
     ... tensor<8x8xf64, #DCSC> ...
 
     // Block sparse row storage (2x3 blocks).
-    #BCSR = #sparse_tensor.encoding<{
+    #BSR = #sparse_tensor.encoding<{
       map = ( i, j ) ->
       ( i floordiv 2 : compressed,
         j floordiv 3 : compressed,
         i mod 2      : dense,
         j mod 3      : dense
       )
     }>
-    ... tensor<20x30xf32, #BCSR> ...
-
-    // ELL storage (4 jagged diagonals, i.e., at most 4 nonzeros per row).
-    #ELL = #sparse_tensor.encoding<{
-      lvlTypes = [ "dense", "dense", "compressed" ],
-      dimToLvl = affine_map<(i, j)[c] -> (c * 4 * i, i, j)>
-    }>
-    ... tensor<?x?xf64, #ELL> ...
+    ... tensor<20x30xf32, #BSR> ...
 
     // CSR slice (offset = 0, size = 4, stride = 1 on the first dimension;
     // offset = 0, size = 8, and a dynamic stride on the second dimension).
@@ -444,7 +434,6 @@ def AnyRankedSparseTensor : RankedSparseTensorOf<[AnyType]>;
 class ScalarLikeOf<list<Type> allowedTypes>
   : AnyTypeOf<[0DTensorOf<allowedTypes>, AnyTypeOf<allowedTypes>]>;
 
-
 //===----------------------------------------------------------------------===//
 // Sparse Tensor Sorting Algorithm Attribute.
 //===----------------------------------------------------------------------===//
