MLIR  17.0.0git
SparseTensor.h
Go to the documentation of this file.
1 //===- SparseTensor.h - Sparse tensor dialect -------------------*- C++ -*-===//
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 MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSOR_H_
10 #define MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSOR_H_
11 
13 #include "mlir/IR/BuiltinTypes.h"
14 #include "mlir/IR/Dialect.h"
15 #include "mlir/IR/OpDefinition.h"
17 #include "mlir/IR/TensorEncoding.h"
20 
21 // We must include Enums.h.inc before AttrDefs.h.inc due to dependency between
22 // StorageSpecifierKindAttr and StorageSpeciferKind Enum.
23 
24 #define GET_ATTRDEF_CLASSES
25 #include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrEnums.h.inc"
26 
27 #define GET_ATTRDEF_CLASSES
28 #include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.h.inc"
29 
30 #define GET_TYPEDEF_CLASSES
31 #include "mlir/Dialect/SparseTensor/IR/SparseTensorTypes.h.inc"
32 
33 #define GET_OP_CLASSES
34 #include "mlir/Dialect/SparseTensor/IR/SparseTensorOps.h.inc"
35 
36 #include "mlir/Dialect/SparseTensor/IR/SparseTensorOpsDialect.h.inc"
37 
38 namespace mlir {
39 namespace sparse_tensor {
40 
41 /// Convenience method to abbreviate casting `getType()`.
42 template <typename T>
43 inline RankedTensorType getRankedTensorType(T t) {
44  return t.getType().template cast<RankedTensorType>();
45 }
46 
47 /// Convenience method to abbreviate casting `getType()`.
48 template <typename T>
49 inline MemRefType getMemRefType(T t) {
50  return t.getType().template cast<MemRefType>();
51 }
52 
53 /// Convenience method to get a sparse encoding attribute from a type.
54 /// Returns null-attribute for any type without an encoding.
55 SparseTensorEncodingAttr getSparseTensorEncoding(Type type);
56 
57 /// Returns true iff the given type is a type for a COO tensor with the last
58 /// dimension level type being unique.
59 bool isUniqueCOOType(RankedTensorType tp);
60 
61 /// Returns the starting dimension for a trailing COO region that spans across
62 /// at least two dimensions. If no such COO region is found, returns the rank
63 /// of the tensor.
64 unsigned getCOOStart(SparseTensorEncodingAttr enc);
65 
66 //
67 // Dimension level types.
68 //
69 
70 // MSVC does not allow this function to be constexpr, because
71 // `SparseTensorEncodingAttr::operator bool` isn't declared constexpr.
72 // And therefore all functions calling it cannot be constexpr either.
73 // TODO: since Clang does allow these to be constexpr, perhaps we should
74 // define a macro to abstract over `inline` vs `constexpr` annotations.
75 inline DimLevelType getDimLevelType(SparseTensorEncodingAttr enc, uint64_t d) {
76  if (enc) {
77  auto types = enc.getDimLevelType();
78  assert(d < types.size() && "Dimension out of bounds");
79  return types[d];
80  }
81  return DimLevelType::Dense; // unannotated tensor is dense
82 }
83 
84 inline DimLevelType getDimLevelType(RankedTensorType type, uint64_t d) {
86 }
87 
88 /// Convenience function to test for dense dimension (0 <= d < rank).
89 inline bool isDenseDim(RankedTensorType type, uint64_t d) {
90  return isDenseDLT(getDimLevelType(type, d));
91 }
92 
93 /// Convenience function to test for compressed dimension (0 <= d < rank).
94 inline bool isCompressedDim(RankedTensorType type, uint64_t d) {
95  return isCompressedDLT(getDimLevelType(type, d));
96 }
97 
98 /// Convenience function to test for singleton dimension (0 <= d < rank).
99 inline bool isSingletonDim(RankedTensorType type, uint64_t d) {
100  return isSingletonDLT(getDimLevelType(type, d));
101 }
102 
103 /// Convenience function to test for dense dimension (0 <= d < rank).
104 inline bool isDenseDim(SparseTensorEncodingAttr enc, uint64_t d) {
105  return isDenseDLT(getDimLevelType(enc, d));
106 }
107 
108 /// Convenience function to test for compressed dimension (0 <= d < rank).
109 inline bool isCompressedDim(SparseTensorEncodingAttr enc, uint64_t d) {
110  return isCompressedDLT(getDimLevelType(enc, d));
111 }
112 
113 /// Convenience function to test for singleton dimension (0 <= d < rank).
114 inline bool isSingletonDim(SparseTensorEncodingAttr enc, uint64_t d) {
115  return isSingletonDLT(getDimLevelType(enc, d));
116 }
117 
118 //
119 // Dimension level properties.
120 //
121 
122 /// Convenience function to test for ordered property in the
123 /// given dimension (0 <= d < rank).
124 inline bool isOrderedDim(RankedTensorType type, uint64_t d) {
125  return isOrderedDLT(getDimLevelType(type, d));
126 }
127 
128 /// Convenience function to test for unique property in the
129 /// given dimension (0 <= d < rank).
130 inline bool isUniqueDim(RankedTensorType type, uint64_t d) {
131  return isUniqueDLT(getDimLevelType(type, d));
132 }
133 
134 //
135 // Reordering.
136 //
137 
138 uint64_t toOrigDim(SparseTensorEncodingAttr enc, uint64_t d);
139 uint64_t toStoredDim(SparseTensorEncodingAttr enc, uint64_t d);
140 
141 /// Convenience method to translate the given stored dimension
142 /// to the original dimension (0 <= d < rank).
143 uint64_t toOrigDim(RankedTensorType type, uint64_t d);
144 
145 /// Convenience method to translate the given original dimension
146 /// to the stored dimension (0 <= d < rank).
147 uint64_t toStoredDim(RankedTensorType type, uint64_t d);
148 
149 } // namespace sparse_tensor
150 } // namespace mlir
151 
152 #endif // MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSOR_H_
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
bool isOrderedDim(RankedTensorType type, uint64_t d)
Convenience function to test for ordered property in the given dimension (0 <= d < rank).
Definition: SparseTensor.h:124
uint64_t toStoredDim(SparseTensorEncodingAttr enc, uint64_t d)
MemRefType getMemRefType(T t)
Convenience method to abbreviate casting getType().
Definition: SparseTensor.h:49
constexpr bool isSingletonDLT(DimLevelType dlt)
Check if the DimLevelType is singleton (regardless of properties).
Definition: Enums.h:228
constexpr bool isUniqueDLT(DimLevelType dlt)
Check if the DimLevelType is unique (regardless of storage format).
Definition: Enums.h:239
uint64_t toOrigDim(SparseTensorEncodingAttr enc, uint64_t d)
bool isCompressedDim(RankedTensorType type, uint64_t d)
Convenience function to test for compressed dimension (0 <= d < rank).
Definition: SparseTensor.h:94
constexpr bool isDenseDLT(DimLevelType dlt)
Check if the DimLevelType is dense.
Definition: Enums.h:213
unsigned getCOOStart(SparseTensorEncodingAttr enc)
Returns the starting dimension for a trailing COO region that spans across at least two dimensions.
bool isUniqueDim(RankedTensorType type, uint64_t d)
Convenience function to test for unique property in the given dimension (0 <= d < rank).
Definition: SparseTensor.h:130
bool isDenseDim(RankedTensorType type, uint64_t d)
Convenience function to test for dense dimension (0 <= d < rank).
Definition: SparseTensor.h:89
constexpr std::optional< DimLevelType > getDimLevelType(LevelFormat lf, bool ordered, bool unique)
Convert a LevelFormat to its corresponding DimLevelType with the given properties.
Definition: Enums.h:257
RankedTensorType getRankedTensorType(T t)
Convenience method to abbreviate casting getType().
Definition: SparseTensor.h:43
SparseTensorEncodingAttr getSparseTensorEncoding(Type type)
Convenience method to get a sparse encoding attribute from a type.
DimLevelType
This enum defines all the sparse representations supportable by the SparseTensor dialect.
Definition: Enums.h:148
constexpr bool isCompressedDLT(DimLevelType dlt)
Check if the DimLevelType is compressed (regardless of properties).
Definition: Enums.h:222
bool isUniqueCOOType(RankedTensorType tp)
Returns true iff the given type is a type for a COO tensor with the last dimension level type being u...
bool isSingletonDim(RankedTensorType type, uint64_t d)
Convenience function to test for singleton dimension (0 <= d < rank).
Definition: SparseTensor.h:99
constexpr bool isOrderedDLT(DimLevelType dlt)
Check if the DimLevelType is ordered (regardless of storage format).
Definition: Enums.h:234
Include the generated interface declarations.