13 #ifndef MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSORTYPE_H_
14 #define MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSORTYPE_H_
19 namespace sparse_tensor {
55 assert(rtp &&
"got null RankedTensorType");
98 return withEncoding(enc.withBitWidths(posWidth, crdWidth));
137 template <
typename T,
typename = std::enable_if_t<
138 std::is_convertible_v<RankedTensorType, T>>>
151 return rtp == other.rtp;
155 return !(*
this == other);
183 CrdTransDirectionKind dir)
const {
184 return enc.translateCrds(builder, loc, crds, dir);
225 : (dimToLvl == other.dimToLvl);
240 CrdTransDirectionKind::dim2lvl);
246 enc.withoutDimToLvl());
253 assert(d <
getDimRank() &&
"Dimension is out of bounds");
282 assert(l < lvlRank &&
"Level out of bounds");
283 return enc.getLvlType(l);
301 unsigned getCrdWidth()
const {
return enc ? enc.getCrdWidth() : 0; }
304 unsigned getPosWidth()
const {
return enc ? enc.getPosWidth() : 0; }
331 RankedTensorType
getCOOType(
bool ordered)
const;
335 const RankedTensorType rtp;
336 const SparseTensorEncodingAttr enc;
348 if (
auto rtp = dyn_cast<RankedTensorType>(val.
getType()))
bool isUnique(It begin, It end)
static ArrayRef< int64_t > getShape(Type type)
Returns the shape of the given type.
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
static AffineMap getMultiDimIdentityMap(unsigned numDims, MLIRContext *context)
Returns an AffineMap with 'numDims' identity result dim exprs.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
MLIRContext is the top-level object for a collection of MLIR operations.
This class helps build Operations.
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
This class provides an abstraction over the different types of ranges over Values.
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Type getType() const
Return the type of this value.
A wrapper around RankedTensorType, which has three goals:
bool isSingletonLvl(Level l) const
MLIRContext * getContext() const
Size getDynamicDimSize(Dimension d) const
Safely looks up the requested dimension-DynSize.
Level getCOOStart() const
Returns the starting level of this sparse tensor type for a trailing COO region that spans at least t...
bool isWithPos(Level l) const
Type getElementType() const
bool isLooseCompressedLvl(Level l) const
SparseTensorType(ShapedType stp, SparseTensorEncodingAttr enc)
bool is2OutOf4Lvl(Level l) const
unsigned getCrdWidth() const
Returns the coordinate-overhead bitwidth, defaulting to zero.
bool operator!=(const SparseTensorType &other) const
ArrayRef< LevelType > getLvlTypes() const
bool hasEncoding() const
Returns true for tensors which have an encoding, and false for those which do not.
ArrayRef< Size > getDimShape() const
Returns the dimension-shape.
bool isAllOrdered() const
Returns true for tensors where every level is ordered.
SmallVector< Size > getLvlShape() const
Returns the Level-shape.
bool operator==(const SparseTensorType &other) const
SparseTensorType withEncoding(SparseTensorEncodingAttr newEnc) const
bool isCOOType(Level startLvl=0, bool isUnique=true) const
Returns true iff this sparse tensor type has a trailing COO region starting at the given level.
Dimension getDimRank() const
Returns the dimension-rank.
AffineMap getLvlToDim() const
Returns the lvlToDiml mapping (or the null-map for the identity).
SparseTensorType withoutDimToLvl() const
SparseTensorType withDimToLvl(AffineMap dimToLvl) const
SparseTensorType(const SparseTensorType &)=default
SparseTensorType withDimToLvl(const SparseTensorType &dimToLvlSTT) const
bool isAllDense() const
Returns true for tensors where every level is dense.
int64_t getNumDynamicDims() const
Returns the number of dimensions which have dynamic sizes.
SparseTensorType withDimSlices(ArrayRef< SparseTensorDimSliceAttr > dimSlices) const
Type getCrdType() const
Returns the coordinate-overhead MLIR type, defaulting to IndexType.
bool isIdentity() const
Returns true if the dimToLvl mapping is the identity.
bool hasDynamicDimShape() const
Returns true if any dimension has dynamic size.
SparseTensorType(RankedTensorType rtp)
bool hasSameDimToLvl(const SparseTensorType &other) const
Returns true iff the two types have the same mapping.
RankedTensorType getRankedTensorType() const
Explicitly convert to RankedTensorType.
SparseTensorType withoutBitWidths() const
SparseTensorType(SparseTensorEncodingAttr enc)
bool isCompressedLvl(Level l) const
bool hasStaticDimShape() const
Returns true if no dimension has dynamic size.
bool isWithCrd(Level l) const
SparseTensorType withoutDimSlices() const
RankedTensorType getDemappedType() const
Returns the type with an identity mapping.
AffineMap getExpandedDimToLvl() const
Returns the dimToLvl mapping, where the identity map is expanded out into a full AffineMap.
SparseTensorType withBitWidths(unsigned posWidth, unsigned crdWidth) const
Level getLvlRank() const
Returns the level-rank.
unsigned getPosWidth() const
Returns the position-overhead bitwidth, defaulting to zero.
RankedTensorType getCOOType(bool ordered) const
Returns [un]ordered COO type for this sparse tensor type.
bool isPermutation() const
Returns true if the dimToLvl mapping is a permutation.
SparseTensorType withDimToLvl(SparseTensorEncodingAttr dimToLvlEnc) const
SparseTensorEncodingAttr getEncoding() const
bool isDynamicDim(Dimension d) const
Returns true if the given dimension has dynamic size.
bool isDenseLvl(Level l) const
LevelType getLvlType(Level l) const
AffineMap getDimToLvl() const
Returns the dimToLvl mapping (or the null-map for the identity).
ValueRange translateCrds(OpBuilder &builder, Location loc, ValueRange crds, CrdTransDirectionKind dir) const
Translates between level / dimension coordinate space.
Type getPosType() const
Returns the position-overhead MLIR type, defaulting to IndexType.
SparseTensorType & operator=(const SparseTensorType &)=delete
bool isOrderedLvl(Level l) const
bool isUniqueLvl(Level l) const
constexpr bool isWithPosLT(LevelType lt)
Check if the LevelType needs positions array.
constexpr bool isLooseCompressedLT(LevelType lt)
Check if the LevelType is loose compressed (regardless of properties).
uint64_t Dimension
The type of dimension identifiers and dimension-ranks.
constexpr bool isUniqueLT(LevelType lt)
Check if the LevelType is unique (regardless of storage format).
uint64_t Level
The type of level identifiers and level-ranks.
std::optional< SparseTensorType > tryGetSparseTensorType(Value val)
int64_t Size
The type for individual components of a compile-time shape, including the value ShapedType::kDynamic ...
constexpr bool isWithCrdLT(LevelType lt)
Check if the LevelType needs coordinates array.
constexpr bool is2OutOf4LT(LevelType lt)
Check if the LevelType is 2OutOf4 (regardless of properties).
constexpr bool isDenseLT(LevelType lt)
Check if the LevelType is dense (regardless of properties).
constexpr bool isSingletonLT(LevelType lt)
Check if the LevelType is singleton (regardless of properties).
constexpr bool isOrderedLT(LevelType lt)
Check if the LevelType is ordered (regardless of storage format).
LevelType
This enum defines all the sparse representations supportable by the SparseTensor dialect.
SparseTensorEncodingAttr getSparseTensorEncoding(Type type)
Convenience method to get a sparse encoding attribute from a type.
constexpr bool isCompressedLT(LevelType lt)
Check if the LevelType is compressed (regardless of properties).
SparseTensorType getSparseTensorType(Value val)
Convenience methods to obtain a SparseTensorType from a Value.
Include the generated interface declarations.
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...