MLIR  19.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 
15 #include "mlir/IR/BuiltinTypes.h"
16 #include "mlir/IR/Dialect.h"
17 #include "mlir/IR/OpDefinition.h"
19 #include "mlir/IR/TensorEncoding.h"
22 
23 //===----------------------------------------------------------------------===//
24 //
25 // Type aliases to help code be more self-documenting. Unfortunately
26 // these are not type-checked, so they only provide documentation rather
27 // than doing anything to prevent mixups.
28 //
29 //===----------------------------------------------------------------------===//
30 
31 namespace mlir {
32 namespace sparse_tensor {
33 
34 /// The type of dimension identifiers and dimension-ranks.
35 using Dimension = uint64_t;
36 
37 /// The type of level identifiers and level-ranks.
38 using Level = uint64_t;
39 
40 /// The type for individual components of a compile-time shape,
41 /// including the value `ShapedType::kDynamic` (for shapes).
42 using Size = int64_t;
43 
44 /// A simple structure that encodes a range of levels in the sparse tensors
45 /// that forms a COO segment.
46 struct COOSegment {
47  std::pair<Level, Level> lvlRange; // [low, high)
48  bool isSoA;
49 
50  bool isAoS() const { return !isSoA; }
51  bool isSegmentStart(Level l) const { return l == lvlRange.first; }
52  bool inSegment(Level l) const {
53  return l >= lvlRange.first && l < lvlRange.second;
54  }
55 };
56 
57 } // namespace sparse_tensor
58 } // namespace mlir
59 
60 //===----------------------------------------------------------------------===//
61 // TableGen-defined classes
62 //===----------------------------------------------------------------------===//
63 
64 #define GET_ATTRDEF_CLASSES
65 #include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrEnums.h.inc"
66 
67 #define GET_ATTRDEF_CLASSES
68 #include "mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.h.inc"
69 
70 #define GET_TYPEDEF_CLASSES
71 #include "mlir/Dialect/SparseTensor/IR/SparseTensorTypes.h.inc"
72 
73 #define GET_OP_CLASSES
74 #include "mlir/Dialect/SparseTensor/IR/SparseTensorOps.h.inc"
75 
76 #include "mlir/Dialect/SparseTensor/IR/SparseTensorOpsDialect.h.inc"
77 
78 //===----------------------------------------------------------------------===//
79 // Additional convenience methods.
80 //===----------------------------------------------------------------------===//
81 
82 namespace mlir {
83 namespace sparse_tensor {
84 
85 /// Convenience method to abbreviate casting `getType()`.
86 template <typename T>
87 inline RankedTensorType getRankedTensorType(T &&t) {
88  assert(static_cast<bool>(std::forward<T>(t)) &&
89  "getRankedTensorType got null argument");
90  return dyn_cast<RankedTensorType>(std::forward<T>(t).getType());
91 }
92 
93 /// Convenience method to abbreviate casting `getType()`.
94 template <typename T>
95 inline MemRefType getMemRefType(T &&t) {
96  assert(static_cast<bool>(std::forward<T>(t)) &&
97  "getMemRefType got null argument");
98  return cast<MemRefType>(std::forward<T>(t).getType());
99 }
100 
101 /// Convenience method to get a sparse encoding attribute from a type.
102 /// Returns null-attribute for any type without an encoding.
103 SparseTensorEncodingAttr getSparseTensorEncoding(Type type);
104 
105 /// Returns true iff the type range has any sparse tensor type.
106 inline bool hasAnySparseType(TypeRange types) {
107  return llvm::any_of(types, [](Type type) {
108  return getSparseTensorEncoding(type) != nullptr;
109  });
110 }
111 
112 /// Returns true iff MLIR operand has any sparse operand.
113 inline bool hasAnySparseOperand(Operation *op) {
114  return hasAnySparseType(op->getOperands().getTypes());
115 }
116 
117 /// Returns true iff MLIR operand has any sparse result.
118 inline bool hasAnySparseResult(Operation *op) {
119  return hasAnySparseType(op->getResults().getTypes());
120 }
121 
122 /// Returns true iff MLIR operand has any sparse operand or result.
124  return hasAnySparseOperand(op) || hasAnySparseResult(op);
125 }
126 
127 /// Returns true iff MLIR operation has any sparse tensor with non-identity
128 /// dim2lvl maps.
130 
131 //
132 // Inference.
133 //
134 
135 /// Given the dimToLvl map, infers the lvlToDim map, or returns
136 /// empty Affine map when inference fails.
137 AffineMap inferLvlToDim(AffineMap dimToLvl, MLIRContext *context);
138 
139 /// Returns the lvlToDim map for the given dimToLvl map specific
140 /// to the block sparse cases.
141 /// Asserts on failure (so only use when known to succeed).
143 
144 /// Given the dimToLvl map, returns the block sizes in a vector.
145 /// For instance, a 2x3 block will return [2, 3]. Unblocked dimension i
146 /// will return 0, and i floordiv 1, i mod 1 will return 1. Therefore,
147 /// the example below will return [0, 1].
148 /// map = ( i, j ) ->
149 /// ( i : dense,
150 /// j floordiv 1 : compressed,
151 /// j mod 1 : dense
152 /// )
153 /// Only valid block sparsity will be accepted.
155 
156 /// Given the dimToLvl map, returns if it's block sparsity.
157 bool isBlockSparsity(AffineMap dimToLvl);
158 
159 //
160 // Reordering.
161 //
162 
163 /// Convenience method to translate the given level to the corresponding
164 /// dimension.
165 /// Requires: `enc` has a permuted dim2lvl map and `0 <= l < lvlRank`.
166 Dimension toDim(SparseTensorEncodingAttr enc, Level l);
167 
168 /// Convenience method to translate the given dimension to the corresponding
169 /// level.
170 /// Requires: `enc` has a permuted dim2lvl map and `0 <= d < dimRank`.
171 Level toLvl(SparseTensorEncodingAttr enc, Dimension d);
172 
173 } // namespace sparse_tensor
174 } // namespace mlir
175 
176 #endif // MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSOR_H_
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Definition: AffineMap.h:47
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
This class provides an abstraction over the various different ranges of value types.
Definition: TypeRange.h:36
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
bool hasAnySparseOperandOrResult(Operation *op)
Returns true iff MLIR operand has any sparse operand or result.
Definition: SparseTensor.h:123
Dimension toDim(SparseTensorEncodingAttr enc, Level l)
Convenience method to translate the given level to the corresponding dimension.
uint64_t Dimension
The type of dimension identifiers and dimension-ranks.
Definition: SparseTensor.h:35
uint64_t Level
The type of level identifiers and level-ranks.
Definition: SparseTensor.h:38
int64_t Size
The type for individual components of a compile-time shape, including the value ShapedType::kDynamic ...
Definition: SparseTensor.h:42
RankedTensorType getRankedTensorType(T &&t)
Convenience method to abbreviate casting getType().
Definition: SparseTensor.h:87
AffineMap inferLvlToDim(AffineMap dimToLvl, MLIRContext *context)
Given the dimToLvl map, infers the lvlToDim map, or returns empty Affine map when inference fails.
SparseTensorEncodingAttr getSparseTensorEncoding(Type type)
Convenience method to get a sparse encoding attribute from a type.
MemRefType getMemRefType(T &&t)
Convenience method to abbreviate casting getType().
Definition: SparseTensor.h:95
bool hasAnySparseType(TypeRange types)
Returns true iff the type range has any sparse tensor type.
Definition: SparseTensor.h:106
Level toLvl(SparseTensorEncodingAttr enc, Dimension d)
Convenience method to translate the given dimension to the corresponding level.
bool isBlockSparsity(AffineMap dimToLvl)
Given the dimToLvl map, returns if it's block sparsity.
bool hasAnyNonIdentityOperandsOrResults(Operation *op)
Returns true iff MLIR operation has any sparse tensor with non-identity dim2lvl maps.
bool hasAnySparseResult(Operation *op)
Returns true iff MLIR operand has any sparse result.
Definition: SparseTensor.h:118
bool hasAnySparseOperand(Operation *op)
Returns true iff MLIR operand has any sparse operand.
Definition: SparseTensor.h:113
SmallVector< unsigned > getBlockSize(AffineMap dimToLvl)
Given the dimToLvl map, returns the block sizes in a vector.
AffineMap inverseBlockSparsity(AffineMap dimToLvl, MLIRContext *context)
Returns the lvlToDim map for the given dimToLvl map specific to the block sparse cases.
Include the generated interface declarations.
A simple structure that encodes a range of levels in the sparse tensors that forms a COO segment.
Definition: SparseTensor.h:46
bool inSegment(Level l) const
Definition: SparseTensor.h:52
bool isSegmentStart(Level l) const
Definition: SparseTensor.h:51
std::pair< Level, Level > lvlRange
Definition: SparseTensor.h:47