13 #ifndef MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSORTYPE_H_
14 #define MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSORTYPE_H_
19 namespace sparse_tensor {
68 assert(rtp &&
"got null RankedTensorType");
111 return withEncoding(enc.withBitWidths(posWidth, crdWidth));
150 template <
typename T,
typename = std::enable_if_t<
151 std::is_convertible_v<RankedTensorType, T>>>
164 return rtp == other.rtp;
168 return !(*
this == other);
196 CrdTransDirectionKind dir)
const {
197 return enc.translateCrds(builder, loc, crds, dir);
238 : (dimToLvl == other.dimToLvl);
253 CrdTransDirectionKind::dim2lvl);
270 enc.withoutDimToLvl());
277 assert(d <
getDimRank() &&
"Dimension is out of bounds");
306 assert(l < lvlRank &&
"Level out of bounds");
307 return enc.getLvlType(l);
325 unsigned getCrdWidth()
const {
return enc ? enc.getCrdWidth() : 0; }
328 unsigned getPosWidth()
const {
return enc ? enc.getPosWidth() : 0; }
357 RankedTensorType
getCOOType(
bool ordered)
const;
364 const RankedTensorType rtp;
365 const SparseTensorEncodingAttr enc;
377 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.
bool isWithPos(Level l) const
Type getElementType() const
bool isLooseCompressedLvl(Level l) const
SparseTensorType(ShapedType stp, SparseTensorEncodingAttr enc)
unsigned getCrdWidth() const
Returns the coordinate-overhead bitwidth, defaulting to zero.
bool operator!=(const SparseTensorType &other) const
unsigned getBatchLvlRank() const
Returns the batched level-rank.
SmallVector< Size > getBatchLvlShape() const
Returns the batched level-shape.
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.
bool isNOutOfMLvl(Level l) const
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
Level getAoSCOOStart() const
Returns the starting level of this sparse tensor type for a trailing COO region that spans at least t...
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
SmallVector< COOSegment > getCOOSegments() const
Returns a list of COO segments in the sparse tensor types.
bool isUniqueLT(LevelType lt)
bool isWithCrdLT(LevelType lt)
bool isWithPosLT(LevelType lt)
bool isOrderedLT(LevelType lt)
bool isSingletonLT(LevelType lt)
uint64_t Dimension
The type of dimension identifiers and dimension-ranks.
bool isCompressedLT(LevelType lt)
uint64_t Level
The type of level identifiers and level-ranks.
std::optional< SparseTensorType > tryGetSparseTensorType(Value val)
bool isLooseCompressedLT(LevelType lt)
int64_t Size
The type for individual components of a compile-time shape, including the value ShapedType::kDynamic ...
SparseTensorEncodingAttr getSparseTensorEncoding(Type type)
Convenience method to get a sparse encoding attribute from a type.
bool isDenseLT(LevelType lt)
SparseTensorType getSparseTensorType(Value val)
Convenience methods to obtain a SparseTensorType from a Value.
bool isNOutOfMLT(LevelType lt)
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...
A simple structure that encodes a range of levels in the sparse tensors that forms a COO segment.
bool inSegment(Level l) const
bool isSegmentStart(Level l) const
std::pair< Level, Level > lvlRange
This enum defines all the sparse representations supportable by the SparseTensor dialect.