1//===-- SparseTensorAttrDefs.td - attributes definitions ---*- tablegen -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8
9#ifndef SPARSETENSOR_ATTRDEFS
10#define SPARSETENSOR_ATTRDEFS
11
12include "mlir/IR/AttrTypeBase.td"
13include "mlir/Dialect/SparseTensor/IR/SparseTensorBase.td"
14include "mlir/IR/TensorEncoding.td"
15
16// All of the Tensor attributes will extend this class.
17class SparseTensor_Attr<string name,
18                        list<Trait> traits = []>
19	: AttrDef<SparseTensor_Dialect, name, traits>;
20
21// Sparse tensor encoding attribute.
22def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
23         [ DeclareAttrInterfaceMethods<VerifiableTensorEncoding> ] > {
24  let mnemonic = "encoding";
25
26  let description = [{
27    An attribute to encode TACO-style information on sparsity properties
28    of tensors. The encoding is eventually used by a **sparse compiler**
29    pass to generate sparse code fully automatically for all tensor
30    expressions that involve tensors with a sparse encoding. Compiler
31    passes that run before this sparse compiler pass need to be
32    aware of the semantics of tensor types with such an encoding.
33
34    Example:
35
36    ```mlir
37    #DCSC = #sparse_tensor.encoding<{
38      dimLevelType = [ "compressed", "compressed" ],
39      dimOrdering = affine_map<(i,j) -> (j,i)>,
40      pointerBitWidth = 32,
41      indexBitWidth = 8
42    }>
43
44
45    ... tensor<8x8xf64, #DCSC> ...
46    ```
47  }];
48
49  // Data in sparse tensor encoding.
50  let parameters = (
51    ins
52    // A dimension level type for each dimension of a tensor type.
53    // The choices are `dense` (dimension should be stored in its entirety),
54    // `compressed` (only non-zero regions or elements should be stored),
55    // or `singleton` (no sibling elements for parent).
56    ArrayRefParameter<
57      "SparseTensorEncodingAttr::DimLevelType",
58      "Per-dimension level type"
59      >: $dimLevelType,
60    // A dimension order on the indices of this tensor type.
61    // Unlike dense storage, most sparse storage schemes do not provide
62    // fast random access. This affine map specifies the order of
63    // dimensions that should be support by the sparse storage scheme
64    // (e.g. (i,j) -> (i,j) requests 2-d row-wise and (i,j) -> (j,i)
65    // requests 2-d column-wise storage).
66    // TODO: block structure with higher-dim inputs
67    "AffineMap":$dimOrdering,
68    // The required bit width for pointer storage. A narrow width reduces
69    // the memory footprint of overhead storage, as long as the width
70    // suffices to define the total required range (viz. the maximum
71    // number of stored entries over all indirection dimensions). The choices
72    // are `8`, `16`, `32`, `64`, or `0` for a native width.
73    "unsigned":$pointerBitWidth,
74    // The required bit width for index storage. A narrow width reduces
75    // the memory footprint of overhead storage, as long as the width
76    // suffices to define the total required range (viz. the maximum
77    // value of each tensor index over all dimensions). The choices are `8`,
78    // `16`, `32`, `64`, or `0` for a native width.
79    "unsigned":$indexBitWidth
80  );
81
82  let genVerifyDecl = 1;
83  let hasCustomAssemblyFormat = 1;
84
85  let extraClassDeclaration = [{
86    // Dimension level types that define sparse tensors:
87    //   Dense      - dimension is dense, every entry is stored
88    //   Compressed - dimension is sparse, only nonzeros are stored
89    //   Singleton  - dimension contains single coordinate, no siblings
90    enum class DimLevelType {
91      Dense, Compressed, Singleton
92    };
93  }];
94}
95
96def IsSparseTensorPred
97  : CPred<"!!::mlir::sparse_tensor::getSparseTensorEncoding($_self)">;
98
99// The following four follow the same idiom as `TensorOf`, `AnyTensor`,
100// `RankedTensorOf`, `AnyRankedTensor`.
101
102class SparseTensorOf<list<Type> allowedTypes>
103  : TensorOf<allowedTypes, [IsSparseTensorPred], "sparse tensor">;
104
105def AnySparseTensor : SparseTensorOf<[AnyType]>;
106
107class RankedSparseTensorOf<list<Type> allowedTypes>
108  : RankedTensorOf<allowedTypes, [IsSparseTensorPred], "ranked sparse tensor">;
109
110def AnyRankedSparseTensor : RankedSparseTensorOf<[AnyType]>;
111
112#endif // SPARSETENSOR_ATTRDEFS
113