Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 2 commits into from
Sep 22, 2023
Merged

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

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
345dcdb
[mlir][sparse] update doc for new surface syntax encoding
aartbik Sep 21, 2023
85f9518
typos
aartbik Sep 21, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
173 changes: 81 additions & 92 deletions mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 :
Expand Down Expand Up @@ -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
Comment on lines +155 to +156
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why these two are treated as properties instead of new level format?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Perhaps we have a choice here, but currently they are implemented as a property, i.e. we only have dense, compressed, and singleton as level formats. So the doc follows the current syntax (and we can update when needed)


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
Expand All @@ -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).
Expand Down Expand Up @@ -444,7 +434,6 @@ def AnyRankedSparseTensor : RankedSparseTensorOf<[AnyType]>;
class ScalarLikeOf<list<Type> allowedTypes>
: AnyTypeOf<[0DTensorOf<allowedTypes>, AnyTypeOf<allowedTypes>]>;


//===----------------------------------------------------------------------===//
// Sparse Tensor Sorting Algorithm Attribute.
//===----------------------------------------------------------------------===//
Expand Down