Skip to content

Commit

Permalink
[mlir][sparse] Factoring out SparseTensorType class
Browse files Browse the repository at this point in the history 
This change adds a new `SparseTensorType` class for making the "dim" vs "lvl" distinction more overt, and for abstracting over the differences between sparse-tensors and dense-tensors.  In addition, this change also adds new type aliases `Dimension`, `Level`, and `FieldIndex` to make code more self-documenting.

Although the diff is very large, the majority of the changes are mechanical in nature (e.g., changing types to use the new aliases, updating variable names to match, etc).  Along the way I also made many variables `const` when they could be; the majority of which required only adding the keyword.  A few places had conditional definitions of these variables, requiring actual code changes; however, that was only done when the overall change was extremely local and easy to extract.  All these changes are included in the current patch only because it would be too onerous to split them off into a separate patch.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D143800
  • Loading branch information
@wrengr
wrengr committed Feb 15, 2023
1 parent 8d995b2 commit f708a54
Show file tree
Hide file tree
Showing 19 changed files with 1,503 additions and 1,154 deletions.
160 changes: 76 additions & 84 deletions mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h
View file
Expand Up @@ -18,6 +18,49 @@
#include "mlir/Interfaces/InferTypeOpInterface.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"

//===----------------------------------------------------------------------===//
//
// Type aliases to help code be more self-documenting. Unfortunately
// these are not type-checked, so they only provide documentation rather
// than doing anything to prevent mixups.
//
// We must include these here (rather than in "SparseTensorType.h")
// because they are used by methods declared in the tablegen files.
//
//===----------------------------------------------------------------------===//

namespace mlir {
namespace sparse_tensor {

/// The type of dimension identifiers, and dimension-ranks. We use the
/// same type for both identifiers and ranks because the latter are used
/// mainly for ordering-comparisons against the former (just like how the
/// one-past-the-end iterators are used).
using Dimension = uint64_t;

/// The type of level identifiers, and level-ranks. We use the same
/// type for both identifiers and ranks because the latter are used
/// mainly for ordering-comparisons against the former (just like how
/// the one-past-the-end iterators are used).
using Level = uint64_t;

/// The type for individual components of a compile-time shape. We avoid
/// calling this "size" because we use the term "sizes" to indicate the
/// actual run-time sizes, whereas this type also allows the value
/// `ShapedType::kDynamic`.
using DynSize = int64_t;

/// The type for individual components of a compile-time shape which
/// are known not to be `ShapedType::kDynamic`.
using StaticSize = int64_t;

} // namespace sparse_tensor
} // namespace mlir

//===----------------------------------------------------------------------===//
// TableGen-defined classes
//===----------------------------------------------------------------------===//

// We must include Enums.h.inc before AttrDefs.h.inc due to dependency between
// StorageSpecifierKindAttr and StorageSpeciferKind Enum.

Expand All @@ -35,6 +78,10 @@

#include "mlir/Dialect/SparseTensor/IR/SparseTensorOpsDialect.h.inc"

//===----------------------------------------------------------------------===//
// Additional convenience methods.
//===----------------------------------------------------------------------===//

namespace mlir {
namespace sparse_tensor {

Expand All @@ -54,103 +101,48 @@ inline MemRefType getMemRefType(T t) {
/// Returns null-attribute for any type without an encoding.
SparseTensorEncodingAttr getSparseTensorEncoding(Type type);

/// Returns true iff the given type is a type for a COO tensor with the last
/// dimension level type being unique.
/// Returns true iff the given type is a COO type where the last level
/// is unique.
bool isUniqueCOOType(TensorType tp);

/// Returns the starting dimension for a trailing COO region that spans across
/// at least two dimensions. If no such COO region is found, returns the rank
/// of the tensor.
unsigned getCOOStart(SparseTensorEncodingAttr enc);
/// Returns the starting level for a trailing COO region that spans
/// at least two levels. If no such COO region is found, then returns
/// the level-rank.
Level getCOOStart(SparseTensorEncodingAttr enc);

/// Helpers to setup a COO type.
RankedTensorType getCOOFromTypeWithOrdering(RankedTensorType src,
AffineMap ordering, bool ordered);

RankedTensorType getCOOFromType(RankedTensorType src, bool ordered);

//
// Dimension level types.
//

// MSVC does not allow this function to be constexpr, because
// `SparseTensorEncodingAttr::operator bool` isn't declared constexpr.
// And therefore all functions calling it cannot be constexpr either.
// TODO: since Clang does allow these to be constexpr, perhaps we should
// define a macro to abstract over `inline` vs `constexpr` annotations.
inline DimLevelType getDimLevelType(SparseTensorEncodingAttr enc, uint64_t d) {
if (enc) {
auto types = enc.getDimLevelType();
assert(d < types.size() && "Dimension out of bounds");
return types[d];
}
return DimLevelType::Dense; // unannotated tensor is dense
}

inline DimLevelType getDimLevelType(RankedTensorType type, uint64_t d) {
return getDimLevelType(getSparseTensorEncoding(type), d);
}

/// Convenience function to test for dense dimension (0 <= d < rank).
inline bool isDenseDim(RankedTensorType type, uint64_t d) {
return isDenseDLT(getDimLevelType(type, d));
}

/// Convenience function to test for compressed dimension (0 <= d < rank).
inline bool isCompressedDim(RankedTensorType type, uint64_t d) {
return isCompressedDLT(getDimLevelType(type, d));
}

/// Convenience function to test for singleton dimension (0 <= d < rank).
inline bool isSingletonDim(RankedTensorType type, uint64_t d) {
return isSingletonDLT(getDimLevelType(type, d));
}

/// Convenience function to test for dense dimension (0 <= d < rank).
inline bool isDenseDim(SparseTensorEncodingAttr enc, uint64_t d) {
return isDenseDLT(getDimLevelType(enc, d));
}

/// Convenience function to test for compressed dimension (0 <= d < rank).
inline bool isCompressedDim(SparseTensorEncodingAttr enc, uint64_t d) {
return isCompressedDLT(getDimLevelType(enc, d));
}

/// Convenience function to test for singleton dimension (0 <= d < rank).
inline bool isSingletonDim(SparseTensorEncodingAttr enc, uint64_t d) {
return isSingletonDLT(getDimLevelType(enc, d));
}

//
// Dimension level properties.
//

/// Convenience function to test for ordered property in the
/// given dimension (0 <= d < rank).
inline bool isOrderedDim(RankedTensorType type, uint64_t d) {
return isOrderedDLT(getDimLevelType(type, d));
}

/// Convenience function to test for unique property in the
/// given dimension (0 <= d < rank).
inline bool isUniqueDim(RankedTensorType type, uint64_t d) {
return isUniqueDLT(getDimLevelType(type, d));
}

//
// Reordering.
//

uint64_t toOrigDim(SparseTensorEncodingAttr enc, uint64_t d);
uint64_t toStoredDim(SparseTensorEncodingAttr enc, uint64_t d);

/// Convenience method to translate the given stored dimension
/// to the original dimension (0 <= d < rank).
uint64_t toOrigDim(RankedTensorType type, uint64_t d);

/// Convenience method to translate the given original dimension
/// to the stored dimension (0 <= d < rank).
uint64_t toStoredDim(RankedTensorType type, uint64_t d);
// This CPP guard is to disable deprecation warnings for the LLVM
// build-bot, while making it easy to re-enable it for local development.
#if 0
#define DEPRECATED \
LLVM_DEPRECATED("The toOrigDim/toStoredDim functions are deprecated " \
"because they only work for permutations; therefore any " \
"code using them cannot support non-permutations.", \
"")
#else
#define DEPRECATED
#endif

/// [deprecated] Convenience method to translate the given level to the
/// corresponding dimension. Requires: `0 <= l < lvlRank`.
DEPRECATED Dimension toOrigDim(SparseTensorEncodingAttr enc, Level l);
DEPRECATED Dimension toOrigDim(RankedTensorType type, Level l);

/// [deprecated] Convenience method to translate the given dimension to
/// the corresponding level. Requires: `0 <= d < dimRank`.
DEPRECATED Level toStoredDim(SparseTensorEncodingAttr enc, Dimension d);
DEPRECATED Level toStoredDim(RankedTensorType type, Dimension d);

#undef DEPRECATED

} // namespace sparse_tensor
} // namespace mlir
Expand Down
39 changes: 31 additions & 8 deletions mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
View file
Expand Up @@ -260,22 +260,45 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
/// reset to the default/identity.
SparseTensorEncodingAttr withoutOrdering() const;

/// Return true if every level is dense in the encoding.
/// Returns true if every level is dense. Also returns true for
/// the null encoding (since dense-tensors are always all-dense).
bool isAllDense() const;

/// Return true if the encoding has an identity dimension ordering.
/// Returns true if every level is ordered. Also returns true for
/// the null encoding (since dense-tensors are always all-ordered).
bool isAllOrdered() const;

/// Returns true if the encoding has an identity dimension ordering.
/// Also returns true for the null encoding (since dense-tensors
/// always have the identity ordering).
bool hasIdDimOrdering() const;

/// Returns the number of storage levels. Asserts that the encoding
/// is non-null (since there is no fixed result that's valid for
/// every dense-tensor).
::mlir::sparse_tensor::Level getLvlRank() const;

/// Safely looks up the level-type for the requested level. (Returns
/// `DimLevelType::Dense` for the null encoding, since dense-tensors
/// are always all-dense.)
::mlir::sparse_tensor::DimLevelType getLvlType(::mlir::sparse_tensor::Level l) const;

bool isDenseLvl(::mlir::sparse_tensor::Level l) const { return isDenseDLT(getLvlType(l)); }
bool isCompressedLvl(::mlir::sparse_tensor::Level l) const { return isCompressedDLT(getLvlType(l)); }
bool isSingletonLvl(::mlir::sparse_tensor::Level l) const { return isSingletonDLT(getLvlType(l)); }
bool isOrderedLvl(::mlir::sparse_tensor::Level l) const { return isOrderedDLT(getLvlType(l)); }
bool isUniqueLvl(::mlir::sparse_tensor::Level l) const { return isUniqueDLT(getLvlType(l)); }

bool isSlice() const {
return !getDimSlices().empty();
}

std::optional<uint64_t> getStaticDimSliceOffset(unsigned dim) const;
std::optional<uint64_t> getStaticDimSliceSize(unsigned dim) const;
std::optional<uint64_t> getStaticDimSliceStride(unsigned dim) const;
std::optional<uint64_t> getStaticLvlSliceOffset(unsigned lvl) const;
std::optional<uint64_t> getStaticLvlSliceSize(unsigned lvl) const;
std::optional<uint64_t> getStaticLvlSliceStride(unsigned lvl) const;
std::optional<uint64_t> getStaticDimSliceOffset(::mlir::sparse_tensor::Dimension dim) const;
std::optional<uint64_t> getStaticDimSliceSize(::mlir::sparse_tensor::Dimension dim) const;
std::optional<uint64_t> getStaticDimSliceStride(::mlir::sparse_tensor::Dimension dim) const;
std::optional<uint64_t> getStaticLvlSliceOffset(::mlir::sparse_tensor::Level lvl) const;
std::optional<uint64_t> getStaticLvlSliceSize(::mlir::sparse_tensor::Level lvl) const;
std::optional<uint64_t> getStaticLvlSliceStride(::mlir::sparse_tensor::Level lvl) const;
}];

let genVerifyDecl = 1;
Expand Down

0 comments on commit f708a54

Please sign in to comment.