diff --git a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
--- a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
+++ b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
@@ -156,6 +156,8 @@
/// the given ordering and expects subsequent add() calls to honor
/// that same ordering for the given indices. The result is a
/// fully permuted coordinate scheme.
+ ///
+ /// Precondition: `sizes` and `perm` must be valid for `rank`.
static SparseTensorCOO<V> *newSparseTensorCOO(uint64_t rank,
const uint64_t *sizes,
const uint64_t *perm,
@@ -175,12 +177,65 @@
unsigned iteratorPos;
};
-/// Abstract base class of sparse tensor storage. Note that we use
-/// function overloading to implement "partial" method specialization.
+/// Abstract base class for `SparseTensorStorage<P,I,V>`. This class
+/// takes responsibility for all the `<P,I,V>`-independent aspects
+/// of the tensor (e.g., shape, sparsity, permutation). In addition,
+/// we use function overloading to implement "partial" method
+/// specialization, which the C-API relies on to catch type errors
+/// arising from our use of opaque pointers.
class SparseTensorStorageBase {
public:
- /// Dimension size query.
- virtual uint64_t getDimSize(uint64_t) const = 0;
+ /// Constructs a new storage object. The `perm` maps the tensor's
+ /// semantic-ordering of dimensions to this object's storage-order.
+ /// The `szs` and `sparsity` arrays are already in storage-order.
+ ///
+ /// Precondition: `perm` and `sparsity` must be valid for `szs.size()`.
+ SparseTensorStorageBase(const std::vector<uint64_t> &szs,
+ const uint64_t *perm, const DimLevelType *sparsity)
+ : dimSizes(szs), rev(getRank()),
+ dimTypes(sparsity, sparsity + getRank()) {
+ const uint64_t rank = getRank();
+ // Validate parameters.
+ assert(rank > 0 && "Trivial shape is unsupported");
+ for (uint64_t r = 0; r < rank; r++) {
+ assert(dimSizes[r] > 0 && "Dimension size zero has trivial storage");
+ assert((dimTypes[r] == DimLevelType::kDense ||
+ dimTypes[r] == DimLevelType::kCompressed) &&
+ "Unsupported DimLevelType");
+ }
+ // Construct the "reverse" (i.e., inverse) permutation.
+ for (uint64_t r = 0; r < rank; r++)
+ rev[perm[r]] = r;
+ }
+
+ virtual ~SparseTensorStorageBase() = default;
+
+ /// Get the rank of the tensor.
+ inline uint64_t getRank() const { return dimSizes.size(); }
+
+ /// Getter for the dimension-sizes array, in storage-order.
+ inline const std::vector<uint64_t> &getDimSizes() const { return dimSizes; }
+
+ /// Safely lookup the size of the given (storage-order) dimension.
+ inline uint64_t getDimSize(uint64_t d) const {
+ assert(d < getRank());
+ return dimSizes[d];
+ }
+
+ /// Getter for the "reverse" permutation, which maps this object's
+ /// storage-order to the tensor's semantic-order.
+ inline const std::vector<uint64_t> &getRev() const { return rev; }
+
+ /// Getter for the dimension-types array, in storage-order.
+ inline const std::vector<DimLevelType> &getDimTypes() const {
+ return dimTypes;
+ }
+
+ /// Safely check if the (storage-order) dimension uses compressed storage.
+ inline bool isCompressedDim(uint64_t d) const {
+ assert(d < getRank());
+ return (dimTypes[d] == DimLevelType::kCompressed);
+ }
/// Overhead storage.
virtual void getPointers(std::vector<uint64_t> **, uint64_t) { fatal("p64"); }
@@ -231,13 +286,15 @@
/// Finishes insertion.
virtual void endInsert() = 0;
- virtual ~SparseTensorStorageBase() = default;
-
private:
static void fatal(const char *tp) {
fprintf(stderr, "unsupported %s\n", tp);
exit(1);
}
+
+ const std::vector<uint64_t> dimSizes;
+ std::vector<uint64_t> rev;
+ const std::vector<DimLevelType> dimTypes;
};
/// A memory-resident sparse tensor using a storage scheme based on
@@ -252,43 +309,35 @@
/// Constructs a sparse tensor storage scheme with the given dimensions,
/// permutation, and per-dimension dense/sparse annotations, using
/// the coordinate scheme tensor for the initial contents if provided.
+ ///
+ /// Precondition: `perm` and `sparsity` must be valid for `szs.size()`.
SparseTensorStorage(const std::vector<uint64_t> &szs, const uint64_t *perm,
const DimLevelType *sparsity,
- SparseTensorCOO<V> *tensor = nullptr)
- : sizes(szs), rev(getRank()), idx(getRank()), pointers(getRank()),
- indices(getRank()) {
- uint64_t rank = getRank();
- // Store "reverse" permutation.
- for (uint64_t r = 0; r < rank; r++)
- rev[perm[r]] = r;
+ SparseTensorCOO<V> *coo = nullptr)
+ : SparseTensorStorageBase(szs, perm, sparsity), pointers(getRank()),
+ indices(getRank()), idx(getRank()) {
+ const uint64_t rank = getRank();
// Provide hints on capacity of pointers and indices.
// TODO: needs fine-tuning based on sparsity
bool allDense = true;
uint64_t sz = 1;
for (uint64_t r = 0; r < rank; r++) {
- assert(sizes[r] > 0 && "Dimension size zero has trivial storage");
- sz = checkedMul(sz, sizes[r]);
- if (sparsity[r] == DimLevelType::kCompressed) {
+ sz = checkedMul(sz, getDimSizes()[r]);
+ if (isCompressedDim(r)) {
pointers[r].reserve(sz + 1);
+ pointers[r].push_back(0);
indices[r].reserve(sz);
sz = 1;
allDense = false;
- // Prepare the pointer structure. We cannot use `appendPointer`
- // here, because `isCompressedDim` won't work until after this
- // preparation has been done.
- pointers[r].push_back(0);
- } else {
- assert(sparsity[r] == DimLevelType::kDense &&
- "singleton not yet supported");
}
}
// Then assign contents from coordinate scheme tensor if provided.
- if (tensor) {
+ if (coo) {
// Ensure both preconditions of `fromCOO`.
- assert(tensor->getSizes() == sizes && "Tensor size mismatch");
- tensor->sort();
+ assert(coo->getSizes() == getDimSizes() && "Tensor size mismatch");
+ coo->sort();
// Now actually insert the `elements`.
- const std::vector<Element<V>> &elements = tensor->getElements();
+ const std::vector<Element<V>> &elements = coo->getElements();
uint64_t nnz = elements.size();
values.reserve(nnz);
fromCOO(elements, 0, nnz, 0);
@@ -299,15 +348,6 @@
~SparseTensorStorage() override = default;
- /// Get the rank of the tensor.
- uint64_t getRank() const { return sizes.size(); }
-
- /// Get the size of the given dimension of the tensor.
- uint64_t getDimSize(uint64_t d) const override {
- assert(d < getRank());
- return sizes[d];
- }
-
/// Partially specialize these getter methods based on template types.
void getPointers(std::vector<P> **out, uint64_t d) override {
assert(d < getRank());
@@ -372,14 +412,18 @@
/// Returns this sparse tensor storage scheme as a new memory-resident
/// sparse tensor in coordinate scheme with the given dimension order.
+ ///
+ /// Precondition: `perm` must be valid for `getRank()`.
SparseTensorCOO<V> *toCOO(const uint64_t *perm) {
// Restore original order of the dimension sizes and allocate coordinate
// scheme with desired new ordering specified in perm.
- uint64_t rank = getRank();
+ const uint64_t rank = getRank();
+ const auto &rev = getRev();
+ const auto &sizes = getDimSizes();
std::vector<uint64_t> orgsz(rank);
for (uint64_t r = 0; r < rank; r++)
orgsz[rev[r]] = sizes[r];
- SparseTensorCOO<V> *tensor = SparseTensorCOO<V>::newSparseTensorCOO(
+ SparseTensorCOO<V> *coo = SparseTensorCOO<V>::newSparseTensorCOO(
rank, orgsz.data(), perm, values.size());
// Populate coordinate scheme restored from old ordering and changed with
// new ordering. Rather than applying both reorderings during the recursion,
@@ -387,9 +431,12 @@
std::vector<uint64_t> reord(rank);
for (uint64_t r = 0; r < rank; r++)
reord[r] = perm[rev[r]];
- toCOO(*tensor, reord, 0, 0);
- assert(tensor->getElements().size() == values.size());
- return tensor;
+ toCOO(*coo, reord, 0, 0);
+ // 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>
+ assert(coo->getElements().size() == values.size());
+ return coo;
}
/// Factory method. Constructs a sparse tensor storage scheme with the given
@@ -397,16 +444,18 @@
/// using the coordinate scheme tensor for the initial contents if provided.
/// In the latter case, the coordinate scheme must respect the same
/// permutation as is desired for the new sparse tensor storage.
+ ///
+ /// Precondition: `shape`, `perm`, and `sparsity` must be valid for `rank`.
static SparseTensorStorage<P, I, V> *
newSparseTensor(uint64_t rank, const uint64_t *shape, const uint64_t *perm,
- const DimLevelType *sparsity, SparseTensorCOO<V> *tensor) {
+ const DimLevelType *sparsity, SparseTensorCOO<V> *coo) {
SparseTensorStorage<P, I, V> *n = nullptr;
- if (tensor) {
- assert(tensor->getRank() == rank);
+ if (coo) {
+ assert(coo->getRank() == rank && "Tensor rank mismatch");
+ const auto &coosz = coo->getSizes();
for (uint64_t r = 0; r < rank; r++)
- assert(shape[r] == 0 || shape[r] == tensor->getSizes()[perm[r]]);
- n = new SparseTensorStorage<P, I, V>(tensor->getSizes(), perm, sparsity,
- tensor);
+ assert(shape[r] == 0 || shape[r] == coosz[perm[r]]);
+ n = new SparseTensorStorage<P, I, V>(coosz, perm, sparsity, coo);
} else {
std::vector<uint64_t> permsz(rank);
for (uint64_t r = 0; r < rank; r++) {
@@ -507,7 +556,9 @@
}
} else {
// Dense dimension.
- for (uint64_t i = 0, sz = sizes[d], off = pos * sz; i < sz; i++) {
+ const uint64_t sz = getDimSizes()[d];
+ const uint64_t off = pos * sz;
+ for (uint64_t i = 0; i < sz; i++) {
idx[reord[d]] = i;
toCOO(tensor, reord, off + i, d + 1);
}
@@ -521,7 +572,7 @@
if (isCompressedDim(d)) {
appendPointer(d, indices[d].size(), count);
} else { // Dense dimension.
- const uint64_t sz = sizes[d];
+ const uint64_t sz = getDimSizes()[d];
assert(sz >= full && "Segment is overfull");
// Assuming we checked for overflows in the constructor, then this
// multiply will never overflow.
@@ -571,19 +622,11 @@
return -1u;
}
- /// Returns true if dimension is compressed.
- inline bool isCompressedDim(uint64_t d) const {
- assert(d < getRank());
- return (!pointers[d].empty());
- }
-
private:
- const std::vector<uint64_t> sizes; // per-dimension sizes
- std::vector<uint64_t> rev; // "reverse" permutation
- std::vector<uint64_t> idx; // index cursor
std::vector<std::vector<P>> pointers;
std::vector<std::vector<I>> indices;
std::vector<V> values;
+ std::vector<uint64_t> idx; // index cursor for lexicographic insertion.
};
/// Helper to convert string to lower case.