Skip to content

[mlir][sparse] minor edits to support lib files #68137

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Oct 3, 2023
Merged

[mlir][sparse] minor edits to support lib files #68137

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
44b857d
[mlir][sparse] minor edits to support lib files
aartbik Oct 3, 2023
bfc7d00
minor edit
aartbik Oct 3, 2023
fab1ee6
format
aartbik Oct 3, 2023
f348685
reviewer suggestions
aartbik Oct 3, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 16 additions & 48 deletions mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,20 +23,19 @@ namespace mlir {
namespace sparse_tensor {

/// An element of a sparse tensor in coordinate-scheme representation
/// (i.e., a pair of coordinates and value). For example, a rank-1
/// (i.e., a pair of coordinates and value). For example, a rank-1
/// vector element would look like
/// ({i}, a[i])
/// and a rank-5 tensor element would look like
/// ({i,j,k,l,m}, a[i,j,k,l,m])
///
/// The coordinates are represented as a (non-owning) pointer into
/// a shared pool of coordinates, rather than being stored directly in
/// this object. This significantly improves performance because it:
/// (1) reduces the per-element memory footprint, and (2) centralizes
/// the memory management for coordinates. The only downside is that
/// the coordinates themselves cannot be retrieved without knowing the
/// rank of the tensor to which this element belongs (and that rank is
/// not stored in this object).
/// The coordinates are represented as a (non-owning) pointer into a
/// shared pool of coordinates, rather than being stored directly in this
/// object. This significantly improves performance because it reduces the
/// per-element memory footprint and centralizes the memory management for
/// coordinates. The only downside is that the coordinates themselves cannot
/// be retrieved without knowing the rank of the tensor to which this element
/// belongs (and that rank is not stored in this object).
template <typename V>
struct Element final {
Element(const uint64_t *coords, V val) : coords(coords), value(val){};
Expand All @@ -48,10 +47,6 @@ struct Element final {
template <typename V>
struct ElementLT final {
ElementLT(uint64_t rank) : rank(rank) {}

/// Compares two elements a la `operator<`.
///
/// Precondition: the elements must both be valid for `rank`.
bool operator()(const Element<V> &e1, const Element<V> &e2) const {
for (uint64_t d = 0; d < rank; ++d) {
if (e1.coords[d] == e2.coords[d])
Expand All @@ -60,13 +55,10 @@ struct ElementLT final {
}
return false;
}

const uint64_t rank;
};

/// The type of callback functions which receive an element. We avoid
/// packaging the coordinates and value together as an `Element` object
/// because this helps keep code somewhat cleaner.
/// The type of callback functions which receive an element.
template <typename V>
using ElementConsumer =
const std::function<void(const std::vector<uint64_t> &, V)> &;
Expand All @@ -89,27 +81,14 @@ class SparseTensorCOO final {
using size_type = typename vector_type::size_type;

/// Constructs a new coordinate-scheme sparse tensor with the given
/// sizes and initial storage capacity.
///
/// Asserts:
/// * `dimSizes` has nonzero size.
/// * the elements of `dimSizes` are nonzero.
/// sizes and an optional initial storage capacity.
explicit SparseTensorCOO(const std::vector<uint64_t> &dimSizes,
uint64_t capacity = 0)
: SparseTensorCOO(dimSizes.size(), dimSizes.data(), capacity) {}

// TODO: make a class for capturing known-valid sizes (a la PermutationRef),
// so that `SparseTensorStorage::toCOO` can avoid redoing these assertions.
// Also so that we can enforce the asserts *before* copying into `dimSizes`.
//
/// Constructs a new coordinate-scheme sparse tensor with the given
/// sizes and initial storage capacity.
///
/// Precondition: `dimSizes` must be valid for `dimRank`.
///
/// Asserts:
/// * `dimRank` is nonzero.
/// * the elements of `dimSizes` are nonzero.
/// sizes and an optional initial storage capacity. The size of the
/// dimSizes array is determined by dimRank.
explicit SparseTensorCOO(uint64_t dimRank, const uint64_t *dimSizes,
uint64_t capacity = 0)
: dimSizes(dimSizes, dimSizes + dimRank), isSorted(true) {
Expand All @@ -134,16 +113,7 @@ class SparseTensorCOO final {
/// Returns the `operator<` closure object for the COO's element type.
ElementLT<V> getElementLT() const { return ElementLT<V>(getRank()); }

/// Adds an element to the tensor. This method does not check whether
/// `dimCoords` is already associated with a value, it adds it regardless.
/// Resolving such conflicts is left up to clients of the iterator
/// interface.
///
/// This method invalidates all iterators.
///
/// Asserts:
/// * the `dimCoords` is valid for `getRank`.
/// * the components of `dimCoords` are valid for `getDimSizes`.
/// Adds an element to the tensor. This method invalidates all iterators.
void add(const std::vector<uint64_t> &dimCoords, V val) {
const uint64_t *base = coordinates.data();
const uint64_t size = coordinates.size();
Expand All @@ -154,7 +124,7 @@ class SparseTensorCOO final {
"Coordinate is too large for the dimension");
coordinates.push_back(dimCoords[d]);
}
// This base only changes if `coordinates` was reallocated. In which
// This base only changes if `coordinates` was reallocated. In which
// case, we need to correct all previous pointers into the vector.
// Note that this only happens if we did not set the initial capacity
// right, and then only for every internal vector reallocation (which
Expand All @@ -175,11 +145,9 @@ class SparseTensorCOO final {
const_iterator begin() const { return elements.cbegin(); }
const_iterator end() const { return elements.cend(); }

/// Sorts elements lexicographically by coordinates. If a coordinate
/// Sorts elements lexicographically by coordinates. If a coordinate
/// is mapped to multiple values, then the relative order of those
/// values is unspecified.
///
/// This method invalidates all iterators.
/// values is unspecified. This method invalidates all iterators.
void sort() {
if (isSorted)
return;
Expand Down
29 changes: 4 additions & 25 deletions mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,8 @@
//
//===----------------------------------------------------------------------===//
//
// This file implements parsing and printing of files in one of the
// following external formats:
// This file implements reading and writing files in one of the following
// external formats:
//
// (1) Matrix Market Exchange (MME): *.mtx
// https://math.nist.gov/MatrixMarket/formats.html
Expand Down Expand Up @@ -197,31 +197,14 @@ class SparseTensorReader final {

/// Allocates a new COO object for `lvlSizes`, initializes it by reading
/// all the elements from the file and applying `dim2lvl` to their
/// dim-coordinates, and then closes the file.
///
/// Preconditions:
/// * `lvlSizes` must be valid for `lvlRank`.
/// * `dim2lvl` must be valid for `getRank()`.
/// * `dim2lvl` maps `getDimSizes()`-coordinates to `lvlSizes`-coordinates.
/// * the file's actual value type can be read as `V`.
///
/// Asserts:
/// * `isValid()`
/// * `dim2lvl` is a permutation, and therefore also `lvlRank == getRank()`.
/// (This requirement will be lifted once we functionalize `dim2lvl`.)
//
// NOTE: This method is factored out of `readSparseTensor` primarily to
// reduce code bloat (since the bulk of the code doesn't care about the
// `<P,I>` type template parameters). But we leave it public since it's
// perfectly reasonable for clients to use.
/// dim-coordinates, and then closes the file. Templated on V only.
template <typename V>
SparseTensorCOO<V> *readCOO(uint64_t lvlRank, const uint64_t *lvlSizes,
const uint64_t *dim2lvl);

/// Allocates a new sparse-tensor storage object with the given encoding,
/// initializes it by reading all the elements from the file, and then
/// closes the file. Preconditions/assertions are as per `readCOO`
/// and `SparseTensorStorage::newFromCOO`.
/// closes the file. Templated on P, I, and V.
template <typename P, typename I, typename V>
SparseTensorStorage<P, I, V> *
readSparseTensor(uint64_t lvlRank, const uint64_t *lvlSizes,
Expand Down Expand Up @@ -312,10 +295,6 @@ SparseTensorCOO<V> *SparseTensorReader::readCOO(uint64_t lvlRank,
const uint64_t *lvlSizes,
const uint64_t *dim2lvl) {
assert(isValid() && "Attempt to readCOO() before readHeader()");
// Construct a `PermutationRef` for the `pushforward` below.
// TODO: This specific implementation does not generalize to arbitrary
// mappings, but once we functionalize the `dim2lvl` argument we can
// simply use that function instead.
const uint64_t dimRank = getRank();
assert(lvlRank == dimRank && "Rank mismatch");
detail::PermutationRef d2l(dimRank, dim2lvl);
Expand Down
74 changes: 23 additions & 51 deletions mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,9 +25,30 @@
#include "mlir/ExecutionEngine/SparseTensor/COO.h"
#include "mlir/ExecutionEngine/SparseTensor/ErrorHandling.h"

#define ASSERT_VALID_DIM(d) \
assert(d < getDimRank() && "Dimension is out of bounds");
#define ASSERT_VALID_LVL(l) \
assert(l < getLvlRank() && "Level is out of bounds");
#define ASSERT_COMPRESSED_LVL(l) \
assert(isCompressedLvl(l) && "Level is not compressed");
#define ASSERT_COMPRESSED_OR_SINGLETON_LVL(l) \
do { \
const DimLevelType dlt = getLvlType(l); \
(void)dlt; \
assert((isCompressedDLT(dlt) || isSingletonDLT(dlt)) && \
"Level is neither compressed nor singleton"); \
} while (false)
#define ASSERT_DENSE_DLT(dlt) assert(isDenseDLT(dlt) && "Level is not dense");

namespace mlir {
namespace sparse_tensor {

// Forward references.
template <typename V>
class SparseTensorEnumeratorBase;
template <typename P, typename C, typename V>
class SparseTensorEnumerator;

namespace detail {

/// Checks whether the `perm` array is a permutation of `[0 .. size)`.
Expand Down Expand Up @@ -125,33 +146,6 @@ class PermutationRef final {

} // namespace detail

//===----------------------------------------------------------------------===//
// This forward decl is sufficient to split `SparseTensorStorageBase` into
// its own header, but isn't sufficient for `SparseTensorStorage` to join it.
template <typename V>
class SparseTensorEnumeratorBase;

// These macros ensure consistent error messages, without risk of incuring
// an additional method call to do so.
#define ASSERT_VALID_DIM(d) \
assert(d < getDimRank() && "Dimension is out of bounds");
#define ASSERT_VALID_LVL(l) \
assert(l < getLvlRank() && "Level is out of bounds");
#define ASSERT_COMPRESSED_LVL(l) \
assert(isCompressedLvl(l) && "Level is not compressed");
#define ASSERT_COMPRESSED_OR_SINGLETON_LVL(l) \
do { \
const DimLevelType dlt = getLvlType(l); \
(void)dlt; \
assert((isCompressedDLT(dlt) || isSingletonDLT(dlt)) && \
"Level is neither compressed nor singleton"); \
} while (false)
// Because the `SparseTensorStorageBase` ctor uses `MLIR_SPARSETENSOR_FATAL`
// (rather than `assert`) when validating level-types, all the uses of
// `ASSERT_DENSE_DLT` are technically unnecessary. However, they are
// retained for the sake of future-proofing.
#define ASSERT_DENSE_DLT(dlt) assert(isDenseDLT(dlt) && "Level is not dense");

/// Abstract base class for `SparseTensorStorage<P,C,V>`. This class
/// takes responsibility for all the `<P,C,V>`-independent aspects
/// of the tensor (e.g., shape, sparsity, permutation). In addition,
Expand Down Expand Up @@ -185,23 +179,9 @@ class SparseTensorEnumeratorBase;
/// specified. Thus, dynamic cardinalities always have an "immutable but
/// unknown" value; so the term "dynamic" should not be taken to indicate
/// run-time mutability.
//
// TODO: we'd like to factor out a class akin to `PermutationRef` for
// capturing known-valid sizes to avoid redundant validity assertions.
// But calling that class "SizesRef" would be a terrible name (and
// "ValidSizesRef" isn't much better). Whereas, calling it "ShapeRef"
// would be a lot nicer, but then that conflicts with the terminology
// introduced above. So we need to come up with some new terminology
// for distinguishing things, which allows a reasonable class name too.
class SparseTensorStorageBase {
protected:
// Since this class is virtual, we must disallow public copying in
// order to avoid "slicing". Since this class has data members,
// that means making copying protected.
// <https://github.com/isocpp/CppCoreGuidelines/blob/master/CppCoreGuidelines.md#Rc-copy-virtual>
SparseTensorStorageBase(const SparseTensorStorageBase &) = default;
// Copy-assignment would be implicitly deleted (because our fields
// are const), so we explicitly delete it for clarity.
SparseTensorStorageBase &operator=(const SparseTensorStorageBase &) = delete;

public:
Expand Down Expand Up @@ -313,10 +293,8 @@ class SparseTensorStorageBase {
MLIR_SPARSETENSOR_FOREVERY_V(DECL_GETVALUES)
#undef DECL_GETVALUES

/// Element-wise insertion in lexicographic coordinate order. The first
/// Element-wise insertion in lexicographic coordinate order. The first
/// argument is the level-coordinates for the value being inserted.
// TODO: For better safety, this should take a parameter for the
// length of `lvlCoords` and check that against `getLvlRank()`.
#define DECL_LEXINSERT(VNAME, V) virtual void lexInsert(const uint64_t *, V);
MLIR_SPARSETENSOR_FOREVERY_V(DECL_LEXINSERT)
#undef DECL_LEXINSERT
Expand Down Expand Up @@ -348,12 +326,6 @@ class SparseTensorStorageBase {
const std::vector<uint64_t> lvl2dim;
};

//===----------------------------------------------------------------------===//
// This forward decl is necessary for defining `SparseTensorStorage`,
// but isn't sufficient for splitting it off.
template <typename P, typename C, typename V>
class SparseTensorEnumerator;

/// A memory-resident sparse tensor using a storage scheme based on
/// per-level sparse/dense annotations. This data structure provides
/// a bufferized form of a sparse tensor type. In contrast to generating
Expand Down Expand Up @@ -612,7 +584,7 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
[&coo](const auto &trgCoords, V val) { coo->add(trgCoords, val); });
// TODO: This assertion assumes there are no stored zeros,
// or if there are then that we don't filter them out.
// Cf., <https://github.com/llvm/llvm-project/issues/54179>
// <https://github.com/llvm/llvm-project/issues/54179>
assert(coo->getElements().size() == values.size());
return coo;
}
Expand Down
11 changes: 4 additions & 7 deletions mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,8 @@
//
//===----------------------------------------------------------------------===//
//
// This header file provides the enums and functions which comprise the
// public API of the `ExecutionEngine/SparseTensorRuntime.cpp` runtime
// support library for the SparseTensor dialect.
// This header file provides the functions which comprise the public API of the
// sparse tensor runtime support library for the SparseTensor dialect.
//
//===----------------------------------------------------------------------===//

Expand Down Expand Up @@ -153,8 +152,7 @@ MLIR_SPARSETENSOR_FOREVERY_V(DECL_GETNEXT)
#undef DECL_GETNEXT

/// Reads the sparse tensor, stores the coordinates and values to the given
/// memrefs. Returns a boolean value to indicate whether the COO elements are
/// sorted.
/// memrefs. Returns a boolean to indicate whether the COO elements are sorted.
#define DECL_GETNEXT(VNAME, V, CNAME, C) \
MLIR_CRUNNERUTILS_EXPORT bool \
_mlir_ciface_getSparseTensorReaderReadToBuffers##CNAME##VNAME( \
Expand Down Expand Up @@ -240,8 +238,7 @@ MLIR_CRUNNERUTILS_EXPORT index_type getSparseTensorReaderNSE(void *p);
MLIR_CRUNNERUTILS_EXPORT index_type getSparseTensorReaderDimSize(void *p,
index_type d);

/// Releases the SparseTensorReader. This also closes the file associated with
/// the reader.
/// Releases the SparseTensorReader and closes the associated file.
MLIR_CRUNNERUTILS_EXPORT void delSparseTensorReader(void *p);

/// Creates a SparseTensorWriter for outputting a sparse tensor to a file
Expand Down