| //===- Storage.h - TACO-flavored sparse tensor representation ---*- C++ -*-===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // This file contains definitions for the following classes: |
| // |
| // * `SparseTensorStorageBase` |
| // * `SparseTensorStorage<P, C, V>` |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #ifndef MLIR_EXECUTIONENGINE_SPARSETENSOR_STORAGE_H |
| #define MLIR_EXECUTIONENGINE_SPARSETENSOR_STORAGE_H |
| |
| #include "mlir/Dialect/SparseTensor/IR/Enums.h" |
| #include "mlir/ExecutionEngine/Float16bits.h" |
| #include "mlir/ExecutionEngine/SparseTensor/ArithmeticUtils.h" |
| #include "mlir/ExecutionEngine/SparseTensor/COO.h" |
| #include "mlir/ExecutionEngine/SparseTensor/MapRef.h" |
| |
| namespace mlir { |
| namespace sparse_tensor { |
| |
| //===----------------------------------------------------------------------===// |
| // |
| // SparseTensorStorage Classes |
| // |
| //===----------------------------------------------------------------------===// |
| |
| /// 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., sizes, sparsity, mapping). 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. |
| /// |
| /// Because this class forms a bridge between the denotational semantics |
| /// of "tensors" and the operational semantics of how we store and |
| /// compute with them, it also distinguishes between two different |
| /// coordinate spaces (and their associated rank, sizes, etc). |
| /// Denotationally, we have the *dimensions* of the tensor represented |
| /// by this object. Operationally, we have the *levels* of the storage |
| /// representation itself. |
| /// |
| /// The *size* of an axis is the cardinality of possible coordinate |
| /// values along that axis (regardless of which coordinates have stored |
| /// element values). As such, each size must be non-zero since if any |
| /// axis has size-zero then the whole tensor would have trivial storage |
| /// (since there are no possible coordinates). Thus we use the plural |
| /// term *sizes* for a collection of non-zero cardinalities, and use |
| /// this term whenever referring to run-time cardinalities. Whereas we |
| /// use the term *shape* for a collection of compile-time cardinalities, |
| /// where zero is used to indicate cardinalities which are dynamic (i.e., |
| /// unknown/unspecified at compile-time). At run-time, these dynamic |
| /// cardinalities will be inferred from or checked against sizes otherwise |
| /// specified. Thus, dynamic cardinalities always have an "immutable but |
| /// unknown" value; so the term "dynamic" should not be taken to indicate |
| /// run-time mutability. |
| class SparseTensorStorageBase { |
| protected: |
| SparseTensorStorageBase(const SparseTensorStorageBase &) = default; |
| SparseTensorStorageBase &operator=(const SparseTensorStorageBase &) = delete; |
| |
| public: |
| /// Constructs a new sparse-tensor storage object with the given encoding. |
| SparseTensorStorageBase(uint64_t dimRank, const uint64_t *dimSizes, |
| uint64_t lvlRank, const uint64_t *lvlSizes, |
| const LevelType *lvlTypes, const uint64_t *dim2lvl, |
| const uint64_t *lvl2dim); |
| virtual ~SparseTensorStorageBase() = default; |
| |
| /// Gets the number of tensor-dimensions. |
| uint64_t getDimRank() const { return dimSizes.size(); } |
| |
| /// Gets the number of storage-levels. |
| uint64_t getLvlRank() const { return lvlSizes.size(); } |
| |
| /// Gets the tensor-dimension sizes array. |
| const std::vector<uint64_t> &getDimSizes() const { return dimSizes; } |
| |
| /// Safely looks up the size of the given tensor-dimension. |
| uint64_t getDimSize(uint64_t d) const { |
| assert(d < getDimRank()); |
| return dimSizes[d]; |
| } |
| |
| /// Gets the storage-level sizes array. |
| const std::vector<uint64_t> &getLvlSizes() const { return lvlSizes; } |
| |
| /// Safely looks up the size of the given storage-level. |
| uint64_t getLvlSize(uint64_t l) const { |
| assert(l < getLvlRank()); |
| return lvlSizes[l]; |
| } |
| |
| /// Gets the level-types array. |
| const std::vector<LevelType> &getLvlTypes() const { return lvlTypes; } |
| |
| /// Safely looks up the type of the given level. |
| LevelType getLvlType(uint64_t l) const { |
| assert(l < getLvlRank()); |
| return lvlTypes[l]; |
| } |
| |
| /// Safely checks if the level uses dense storage. |
| bool isDenseLvl(uint64_t l) const { return isDenseLT(getLvlType(l)); } |
| |
| /// Safely checks if the level uses compressed storage. |
| bool isCompressedLvl(uint64_t l) const { |
| return isCompressedLT(getLvlType(l)); |
| } |
| |
| /// Safely checks if the level uses loose compressed storage. |
| bool isLooseCompressedLvl(uint64_t l) const { |
| return isLooseCompressedLT(getLvlType(l)); |
| } |
| |
| /// Safely checks if the level uses singleton storage. |
| bool isSingletonLvl(uint64_t l) const { return isSingletonLT(getLvlType(l)); } |
| |
| /// Safely checks if the level uses n out of m storage. |
| bool isNOutOfMLvl(uint64_t l) const { return isNOutOfMLT(getLvlType(l)); } |
| |
| /// Safely checks if the level is ordered. |
| bool isOrderedLvl(uint64_t l) const { return isOrderedLT(getLvlType(l)); } |
| |
| /// Safely checks if the level is unique. |
| bool isUniqueLvl(uint64_t l) const { return isUniqueLT(getLvlType(l)); } |
| |
| /// Gets positions-overhead storage for the given level. |
| #define DECL_GETPOSITIONS(PNAME, P) \ |
| virtual void getPositions(std::vector<P> **, uint64_t); |
| MLIR_SPARSETENSOR_FOREVERY_FIXED_O(DECL_GETPOSITIONS) |
| #undef DECL_GETPOSITIONS |
| |
| /// Gets coordinates-overhead storage for the given level. |
| #define DECL_GETCOORDINATES(INAME, C) \ |
| virtual void getCoordinates(std::vector<C> **, uint64_t); |
| MLIR_SPARSETENSOR_FOREVERY_FIXED_O(DECL_GETCOORDINATES) |
| #undef DECL_GETCOORDINATES |
| |
| /// Gets coordinates-overhead storage buffer for the given level. |
| #define DECL_GETCOORDINATESBUFFER(INAME, C) \ |
| virtual void getCoordinatesBuffer(std::vector<C> **, uint64_t); |
| MLIR_SPARSETENSOR_FOREVERY_FIXED_O(DECL_GETCOORDINATESBUFFER) |
| #undef DECL_GETCOORDINATESBUFFER |
| |
| /// Gets primary storage. |
| #define DECL_GETVALUES(VNAME, V) virtual void getValues(std::vector<V> **); |
| MLIR_SPARSETENSOR_FOREVERY_V(DECL_GETVALUES) |
| #undef DECL_GETVALUES |
| |
| /// Element-wise insertion in lexicographic coordinate order. The first |
| /// argument is the level-coordinates for the value being inserted. |
| #define DECL_LEXINSERT(VNAME, V) virtual void lexInsert(const uint64_t *, V); |
| MLIR_SPARSETENSOR_FOREVERY_V(DECL_LEXINSERT) |
| #undef DECL_LEXINSERT |
| |
| /// Expanded insertion. Note that this method resets the |
| /// values/filled-switch array back to all-zero/false while only |
| /// iterating over the nonzero elements. |
| #define DECL_EXPINSERT(VNAME, V) \ |
| virtual void expInsert(uint64_t *, V *, bool *, uint64_t *, uint64_t, \ |
| uint64_t); |
| MLIR_SPARSETENSOR_FOREVERY_V(DECL_EXPINSERT) |
| #undef DECL_EXPINSERT |
| |
| /// Finalizes lexicographic insertions. |
| virtual void endLexInsert() = 0; |
| |
| private: |
| const std::vector<uint64_t> dimSizes; |
| const std::vector<uint64_t> lvlSizes; |
| const std::vector<LevelType> lvlTypes; |
| const std::vector<uint64_t> dim2lvlVec; |
| const std::vector<uint64_t> lvl2dimVec; |
| |
| protected: |
| const MapRef map; // non-owning pointers into dim2lvl/lvl2dim vectors |
| const bool allDense; |
| }; |
| |
| /// 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 |
| /// setup methods for each differently annotated sparse tensor, this |
| /// method provides a convenient "one-size-fits-all" solution that simply |
| /// takes an input tensor and annotations to implement all required setup |
| /// in a general manner. |
| template <typename P, typename C, typename V> |
| class SparseTensorStorage final : public SparseTensorStorageBase { |
| /// Private constructor to share code between the other constructors. |
| /// Beware that the object is not necessarily guaranteed to be in a |
| /// valid state after this constructor alone; e.g., `isCompressedLvl(l)` |
| /// doesn't entail `!(positions[l].empty())`. |
| SparseTensorStorage(uint64_t dimRank, const uint64_t *dimSizes, |
| uint64_t lvlRank, const uint64_t *lvlSizes, |
| const LevelType *lvlTypes, const uint64_t *dim2lvl, |
| const uint64_t *lvl2dim) |
| : SparseTensorStorageBase(dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes, |
| dim2lvl, lvl2dim), |
| positions(lvlRank), coordinates(lvlRank), lvlCursor(lvlRank) {} |
| |
| public: |
| /// Constructs a sparse tensor with the given encoding, and allocates |
| /// overhead storage according to some simple heuristics. When lvlCOO |
| /// is set, the sparse tensor initializes with the contents from that |
| /// data structure. Otherwise, an empty sparse tensor results. |
| SparseTensorStorage(uint64_t dimRank, const uint64_t *dimSizes, |
| uint64_t lvlRank, const uint64_t *lvlSizes, |
| const LevelType *lvlTypes, const uint64_t *dim2lvl, |
| const uint64_t *lvl2dim, SparseTensorCOO<V> *lvlCOO); |
| |
| /// Constructs a sparse tensor with the given encoding, and initializes |
| /// the contents from the level buffers. The constructor assumes that the |
| /// data provided by `lvlBufs` can be directly used to interpret the result |
| /// sparse tensor and performs no integrity test on the input data. |
| SparseTensorStorage(uint64_t dimRank, const uint64_t *dimSizes, |
| uint64_t lvlRank, const uint64_t *lvlSizes, |
| const LevelType *lvlTypes, const uint64_t *dim2lvl, |
| const uint64_t *lvl2dim, const intptr_t *lvlBufs); |
| |
| /// Allocates a new empty sparse tensor. |
| static SparseTensorStorage<P, C, V> * |
| newEmpty(uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank, |
| const uint64_t *lvlSizes, const LevelType *lvlTypes, |
| const uint64_t *dim2lvl, const uint64_t *lvl2dim); |
| |
| /// Allocates a new sparse tensor and initializes it from the given COO. |
| static SparseTensorStorage<P, C, V> * |
| newFromCOO(uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank, |
| const uint64_t *lvlSizes, const LevelType *lvlTypes, |
| const uint64_t *dim2lvl, const uint64_t *lvl2dim, |
| SparseTensorCOO<V> *lvlCOO); |
| |
| /// Allocates a new sparse tensor and initialize it from the given buffers. |
| static SparseTensorStorage<P, C, V> * |
| newFromBuffers(uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank, |
| const uint64_t *lvlSizes, const LevelType *lvlTypes, |
| const uint64_t *dim2lvl, const uint64_t *lvl2dim, |
| uint64_t srcRank, const intptr_t *buffers); |
| |
| ~SparseTensorStorage() final = default; |
| |
| /// Partially specialize these getter methods based on template types. |
| void getPositions(std::vector<P> **out, uint64_t lvl) final { |
| assert(out && "Received nullptr for out parameter"); |
| assert(lvl < getLvlRank()); |
| *out = &positions[lvl]; |
| } |
| void getCoordinates(std::vector<C> **out, uint64_t lvl) final { |
| assert(out && "Received nullptr for out parameter"); |
| assert(lvl < getLvlRank()); |
| *out = &coordinates[lvl]; |
| } |
| void getCoordinatesBuffer(std::vector<C> **out, uint64_t lvl) final { |
| assert(out && "Received nullptr for out parameter"); |
| assert(lvl < getLvlRank()); |
| // Note that the sparse tensor support library always stores COO in SoA |
| // format, even when AoS is requested. This is never an issue, since all |
| // actual code/library generation requests "views" into the coordinate |
| // storage for the individual levels, which is trivially provided for |
| // both AoS and SoA (as well as all the other storage formats). The only |
| // exception is when the buffer version of coordinate storage is requested |
| // (currently only for printing). In that case, we do the following |
| // potentially expensive transformation to provide that view. If this |
| // operation becomes more common beyond debugging, we should consider |
| // implementing proper AoS in the support library as well. |
| uint64_t lvlRank = getLvlRank(); |
| uint64_t nnz = values.size(); |
| crdBuffer.clear(); |
| crdBuffer.reserve(nnz * (lvlRank - lvl)); |
| for (uint64_t i = 0; i < nnz; i++) { |
| for (uint64_t l = lvl; l < lvlRank; l++) { |
| assert(i < coordinates[l].size()); |
| crdBuffer.push_back(coordinates[l][i]); |
| } |
| } |
| *out = &crdBuffer; |
| } |
| void getValues(std::vector<V> **out) final { |
| assert(out && "Received nullptr for out parameter"); |
| *out = &values; |
| } |
| |
| /// Partially specialize lexicographical insertions based on template types. |
| void lexInsert(const uint64_t *lvlCoords, V val) final { |
| assert(lvlCoords); |
| if (allDense) { |
| uint64_t lvlRank = getLvlRank(); |
| uint64_t valIdx = 0; |
| // Linearize the address. |
| for (uint64_t l = 0; l < lvlRank; l++) |
| valIdx = valIdx * getLvlSize(l) + lvlCoords[l]; |
| values[valIdx] = val; |
| return; |
| } |
| // First, wrap up pending insertion path. |
| uint64_t diffLvl = 0; |
| uint64_t full = 0; |
| if (!values.empty()) { |
| diffLvl = lexDiff(lvlCoords); |
| endPath(diffLvl + 1); |
| full = lvlCursor[diffLvl] + 1; |
| } |
| // Then continue with insertion path. |
| insPath(lvlCoords, diffLvl, full, val); |
| } |
| |
| /// Partially specialize expanded insertions based on template types. |
| void expInsert(uint64_t *lvlCoords, V *values, bool *filled, uint64_t *added, |
| uint64_t count, uint64_t expsz) final { |
| assert((lvlCoords && values && filled && added) && "Received nullptr"); |
| if (count == 0) |
| return; |
| // Sort. |
| std::sort(added, added + count); |
| // Restore insertion path for first insert. |
| const uint64_t lastLvl = getLvlRank() - 1; |
| uint64_t c = added[0]; |
| assert(c <= expsz); |
| assert(filled[c] && "added coordinate is not filled"); |
| lvlCoords[lastLvl] = c; |
| lexInsert(lvlCoords, values[c]); |
| values[c] = 0; |
| filled[c] = false; |
| // Subsequent insertions are quick. |
| for (uint64_t i = 1; i < count; i++) { |
| assert(c < added[i] && "non-lexicographic insertion"); |
| c = added[i]; |
| assert(c <= expsz); |
| assert(filled[c] && "added coordinate is not filled"); |
| lvlCoords[lastLvl] = c; |
| insPath(lvlCoords, lastLvl, added[i - 1] + 1, values[c]); |
| values[c] = 0; |
| filled[c] = false; |
| } |
| } |
| |
| /// Finalizes lexicographic insertions. |
| void endLexInsert() final { |
| if (!allDense) { |
| if (values.empty()) |
| finalizeSegment(0); |
| else |
| endPath(0); |
| } |
| } |
| |
| /// Sort the unordered tensor in place, the method assumes that it is |
| /// an unordered COO tensor. |
| void sortInPlace() { |
| uint64_t nnz = values.size(); |
| #ifndef NDEBUG |
| for (uint64_t l = 0; l < getLvlRank(); l++) |
| assert(nnz == coordinates[l].size()); |
| #endif |
| |
| // In-place permutation. |
| auto applyPerm = [this](std::vector<uint64_t> &perm) { |
| uint64_t length = perm.size(); |
| uint64_t lvlRank = getLvlRank(); |
| // Cache for the current level coordinates. |
| std::vector<P> lvlCrds(lvlRank); |
| for (uint64_t i = 0; i < length; i++) { |
| uint64_t current = i; |
| if (i != perm[current]) { |
| for (uint64_t l = 0; l < lvlRank; l++) |
| lvlCrds[l] = coordinates[l][i]; |
| V val = values[i]; |
| // Deals with a permutation cycle. |
| while (i != perm[current]) { |
| uint64_t next = perm[current]; |
| // Swaps the level coordinates and value. |
| for (uint64_t l = 0; l < lvlRank; l++) |
| coordinates[l][current] = coordinates[l][next]; |
| values[current] = values[next]; |
| perm[current] = current; |
| current = next; |
| } |
| for (uint64_t l = 0; l < lvlRank; l++) |
| coordinates[l][current] = lvlCrds[l]; |
| values[current] = val; |
| perm[current] = current; |
| } |
| } |
| }; |
| |
| std::vector<uint64_t> sortedIdx(nnz, 0); |
| for (uint64_t i = 0; i < nnz; i++) |
| sortedIdx[i] = i; |
| |
| std::sort(sortedIdx.begin(), sortedIdx.end(), |
| [this](uint64_t lhs, uint64_t rhs) { |
| for (uint64_t l = 0; l < getLvlRank(); l++) { |
| if (coordinates[l][lhs] == coordinates[l][rhs]) |
| continue; |
| return coordinates[l][lhs] < coordinates[l][rhs]; |
| } |
| assert(lhs == rhs && "duplicate coordinates"); |
| return false; |
| }); |
| |
| applyPerm(sortedIdx); |
| } |
| |
| private: |
| /// Appends coordinate `crd` to level `lvl`, in the semantically |
| /// general sense. For non-dense levels, that means appending to the |
| /// `coordinates[lvl]` array, checking that `crd` is representable in |
| /// the `C` type; however, we do not verify other semantic requirements |
| /// (e.g., that `crd` is in bounds for `lvlSizes[lvl]`, and not previously |
| /// occurring in the same segment). For dense levels, this method instead |
| /// appends the appropriate number of zeros to the `values` array, where |
| /// `full` is the number of "entries" already written to `values` for this |
| /// segment (aka one after the highest coordinate previously appended). |
| void appendCrd(uint64_t lvl, uint64_t full, uint64_t crd) { |
| if (!isDenseLvl(lvl)) { |
| assert(isCompressedLvl(lvl) || isLooseCompressedLvl(lvl) || |
| isSingletonLvl(lvl) || isNOutOfMLvl(lvl)); |
| coordinates[lvl].push_back(detail::checkOverflowCast<C>(crd)); |
| } else { // Dense level. |
| assert(crd >= full && "Coordinate was already filled"); |
| if (crd == full) |
| return; // Short-circuit, since it'll be a nop. |
| if (lvl + 1 == getLvlRank()) |
| values.insert(values.end(), crd - full, 0); |
| else |
| finalizeSegment(lvl + 1, 0, crd - full); |
| } |
| } |
| |
| /// Computes the assembled-size associated with the `l`-th level, |
| /// given the assembled-size associated with the `(l-1)`-th level. |
| uint64_t assembledSize(uint64_t parentSz, uint64_t l) const { |
| if (isCompressedLvl(l)) |
| return positions[l][parentSz]; |
| if (isLooseCompressedLvl(l)) |
| return positions[l][2 * parentSz - 1]; |
| if (isSingletonLvl(l) || isNOutOfMLvl(l)) |
| return parentSz; // new size same as the parent |
| assert(isDenseLvl(l)); |
| return parentSz * getLvlSize(l); |
| } |
| |
| /// Initializes sparse tensor storage scheme from a memory-resident sparse |
| /// tensor in coordinate scheme. This method prepares the positions and |
| /// coordinates arrays under the given per-level dense/sparse annotations. |
| void fromCOO(const std::vector<Element<V>> &lvlElements, uint64_t lo, |
| uint64_t hi, uint64_t l) { |
| const uint64_t lvlRank = getLvlRank(); |
| assert(l <= lvlRank && hi <= lvlElements.size()); |
| // Once levels are exhausted, insert the numerical values. |
| if (l == lvlRank) { |
| assert(lo < hi); |
| values.push_back(lvlElements[lo].value); |
| return; |
| } |
| // Visit all elements in this interval. |
| uint64_t full = 0; |
| while (lo < hi) { // If `hi` is unchanged, then `lo < lvlElements.size()`. |
| // Find segment in interval with same coordinate at this level. |
| const uint64_t c = lvlElements[lo].coords[l]; |
| uint64_t seg = lo + 1; |
| if (isUniqueLvl(l)) |
| while (seg < hi && lvlElements[seg].coords[l] == c) |
| seg++; |
| // Handle segment in interval for sparse or dense level. |
| appendCrd(l, full, c); |
| full = c + 1; |
| fromCOO(lvlElements, lo, seg, l + 1); |
| // And move on to next segment in interval. |
| lo = seg; |
| } |
| // Finalize the sparse position structure at this level. |
| finalizeSegment(l, full); |
| } |
| |
| /// Finalizes the sparse position structure at this level. |
| void finalizeSegment(uint64_t l, uint64_t full = 0, uint64_t count = 1) { |
| if (count == 0) |
| return; // Short-circuit, since it'll be a nop. |
| if (isCompressedLvl(l)) { |
| uint64_t pos = coordinates[l].size(); |
| positions[l].insert(positions[l].end(), count, |
| detail::checkOverflowCast<P>(pos)); |
| } else if (isLooseCompressedLvl(l)) { |
| // Finish this level, and push pairs for the empty ones, and one |
| // more for next level. Note that this always leaves one extra |
| // unused element at the end. |
| uint64_t pos = coordinates[l].size(); |
| positions[l].insert(positions[l].end(), 2 * count, |
| detail::checkOverflowCast<P>(pos)); |
| } else if (isSingletonLvl(l) || isNOutOfMLvl(l)) { |
| return; // Nothing to finalize. |
| } else { // Dense dimension. |
| assert(isDenseLvl(l)); |
| const uint64_t sz = getLvlSizes()[l]; |
| assert(sz >= full && "Segment is overfull"); |
| count = detail::checkedMul(count, sz - full); |
| // For dense storage we must enumerate all the remaining coordinates |
| // in this level (i.e., coordinates after the last non-zero |
| // element), and either fill in their zero values or else recurse |
| // to finalize some deeper level. |
| if (l + 1 == getLvlRank()) |
| values.insert(values.end(), count, 0); |
| else |
| finalizeSegment(l + 1, 0, count); |
| } |
| } |
| |
| /// Wraps up a single insertion path, inner to outer. |
| void endPath(uint64_t diffLvl) { |
| const uint64_t lvlRank = getLvlRank(); |
| const uint64_t lastLvl = lvlRank - 1; |
| assert(diffLvl <= lvlRank); |
| const uint64_t stop = lvlRank - diffLvl; |
| for (uint64_t i = 0; i < stop; i++) { |
| const uint64_t l = lastLvl - i; |
| finalizeSegment(l, lvlCursor[l] + 1); |
| } |
| } |
| |
| /// Continues a single insertion path, outer to inner. The first |
| /// argument is the level-coordinates for the value being inserted. |
| void insPath(const uint64_t *lvlCoords, uint64_t diffLvl, uint64_t full, |
| V val) { |
| const uint64_t lvlRank = getLvlRank(); |
| assert(diffLvl <= lvlRank); |
| for (uint64_t l = diffLvl; l < lvlRank; l++) { |
| const uint64_t c = lvlCoords[l]; |
| appendCrd(l, full, c); |
| full = 0; |
| lvlCursor[l] = c; |
| } |
| values.push_back(val); |
| } |
| |
| /// Finds the lexicographically first level where the level-coordinates |
| /// in the argument differ from those in the current cursor. |
| uint64_t lexDiff(const uint64_t *lvlCoords) const { |
| const uint64_t lvlRank = getLvlRank(); |
| for (uint64_t l = 0; l < lvlRank; l++) { |
| const auto crd = lvlCoords[l]; |
| const auto cur = lvlCursor[l]; |
| if (crd > cur || (crd == cur && !isUniqueLvl(l)) || |
| (crd < cur && !isOrderedLvl(l))) { |
| return l; |
| } |
| if (crd < cur) { |
| assert(false && "non-lexicographic insertion"); |
| return -1u; |
| } |
| } |
| assert(false && "duplicate insertion"); |
| return -1u; |
| } |
| |
| // Sparse tensor storage components. |
| std::vector<std::vector<P>> positions; |
| std::vector<std::vector<C>> coordinates; |
| std::vector<V> values; |
| |
| // Auxiliary data structures. |
| std::vector<uint64_t> lvlCursor; |
| std::vector<C> crdBuffer; // just for AoS view |
| }; |
| |
| //===----------------------------------------------------------------------===// |
| // |
| // SparseTensorStorage Factories |
| // |
| //===----------------------------------------------------------------------===// |
| |
| template <typename P, typename C, typename V> |
| SparseTensorStorage<P, C, V> *SparseTensorStorage<P, C, V>::newEmpty( |
| uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank, |
| const uint64_t *lvlSizes, const LevelType *lvlTypes, |
| const uint64_t *dim2lvl, const uint64_t *lvl2dim) { |
| SparseTensorCOO<V> *noLvlCOO = nullptr; |
| return new SparseTensorStorage<P, C, V>(dimRank, dimSizes, lvlRank, lvlSizes, |
| lvlTypes, dim2lvl, lvl2dim, noLvlCOO); |
| } |
| |
| template <typename P, typename C, typename V> |
| SparseTensorStorage<P, C, V> *SparseTensorStorage<P, C, V>::newFromCOO( |
| uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank, |
| const uint64_t *lvlSizes, const LevelType *lvlTypes, |
| const uint64_t *dim2lvl, const uint64_t *lvl2dim, |
| SparseTensorCOO<V> *lvlCOO) { |
| assert(lvlCOO); |
| return new SparseTensorStorage<P, C, V>(dimRank, dimSizes, lvlRank, lvlSizes, |
| lvlTypes, dim2lvl, lvl2dim, lvlCOO); |
| } |
| |
| template <typename P, typename C, typename V> |
| SparseTensorStorage<P, C, V> *SparseTensorStorage<P, C, V>::newFromBuffers( |
| uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank, |
| const uint64_t *lvlSizes, const LevelType *lvlTypes, |
| const uint64_t *dim2lvl, const uint64_t *lvl2dim, uint64_t srcRank, |
| const intptr_t *buffers) { |
| return new SparseTensorStorage<P, C, V>(dimRank, dimSizes, lvlRank, lvlSizes, |
| lvlTypes, dim2lvl, lvl2dim, buffers); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // |
| // SparseTensorStorage Constructors |
| // |
| //===----------------------------------------------------------------------===// |
| |
| template <typename P, typename C, typename V> |
| SparseTensorStorage<P, C, V>::SparseTensorStorage( |
| uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank, |
| const uint64_t *lvlSizes, const LevelType *lvlTypes, |
| const uint64_t *dim2lvl, const uint64_t *lvl2dim, |
| SparseTensorCOO<V> *lvlCOO) |
| : SparseTensorStorage(dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes, |
| dim2lvl, lvl2dim) { |
| // Provide hints on capacity of positions and coordinates. |
| // TODO: needs much fine-tuning based on actual sparsity; currently |
| // we reserve position/coordinate space based on all previous dense |
| // levels, which works well up to first sparse level; but we should |
| // really use nnz and dense/sparse distribution. |
| uint64_t sz = 1; |
| for (uint64_t l = 0; l < lvlRank; l++) { |
| if (isCompressedLvl(l)) { |
| positions[l].reserve(sz + 1); |
| positions[l].push_back(0); |
| coordinates[l].reserve(sz); |
| sz = 1; |
| } else if (isLooseCompressedLvl(l)) { |
| positions[l].reserve(2 * sz + 1); // last one unused |
| positions[l].push_back(0); |
| coordinates[l].reserve(sz); |
| sz = 1; |
| } else if (isSingletonLvl(l)) { |
| coordinates[l].reserve(sz); |
| sz = 1; |
| } else if (isNOutOfMLvl(l)) { |
| assert(l == lvlRank - 1 && "unexpected n:m usage"); |
| sz = detail::checkedMul(sz, lvlSizes[l]) / 2; |
| coordinates[l].reserve(sz); |
| values.reserve(sz); |
| } else { // Dense level. |
| assert(isDenseLvl(l)); |
| sz = detail::checkedMul(sz, lvlSizes[l]); |
| } |
| } |
| if (lvlCOO) { |
| /* New from COO: ensure it is sorted. */ |
| assert(lvlCOO->getRank() == lvlRank); |
| lvlCOO->sort(); |
| // Now actually insert the `elements`. |
| const auto &elements = lvlCOO->getElements(); |
| const uint64_t nse = elements.size(); |
| assert(values.size() == 0); |
| values.reserve(nse); |
| fromCOO(elements, 0, nse, 0); |
| } else if (allDense) { |
| /* New empty (all dense) */ |
| values.resize(sz, 0); |
| } |
| } |
| |
| template <typename P, typename C, typename V> |
| SparseTensorStorage<P, C, V>::SparseTensorStorage( |
| uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank, |
| const uint64_t *lvlSizes, const LevelType *lvlTypes, |
| const uint64_t *dim2lvl, const uint64_t *lvl2dim, const intptr_t *lvlBufs) |
| : SparseTensorStorage(dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes, |
| dim2lvl, lvl2dim) { |
| // Note that none of the buffers can be reused because ownership |
| // of the memory passed from clients is not necessarily transferred. |
| // Therefore, all data is copied over into a new SparseTensorStorage. |
| uint64_t trailCOOLen = 0, parentSz = 1, bufIdx = 0; |
| for (uint64_t l = 0; l < lvlRank; l++) { |
| if (!isUniqueLvl(l) && (isCompressedLvl(l) || isLooseCompressedLvl(l))) { |
| // A `(loose)compressed_nu` level marks the start of trailing COO |
| // start level. Since the coordinate buffer used for trailing COO |
| // is passed in as AoS scheme and SparseTensorStorage uses a SoA |
| // scheme, we cannot simply copy the value from the provided buffers. |
| trailCOOLen = lvlRank - l; |
| break; |
| } |
| if (isCompressedLvl(l) || isLooseCompressedLvl(l)) { |
| P *posPtr = reinterpret_cast<P *>(lvlBufs[bufIdx++]); |
| C *crdPtr = reinterpret_cast<C *>(lvlBufs[bufIdx++]); |
| if (isLooseCompressedLvl(l)) { |
| positions[l].assign(posPtr, posPtr + 2 * parentSz); |
| coordinates[l].assign(crdPtr, crdPtr + positions[l][2 * parentSz - 1]); |
| } else { |
| positions[l].assign(posPtr, posPtr + parentSz + 1); |
| coordinates[l].assign(crdPtr, crdPtr + positions[l][parentSz]); |
| } |
| } else if (isSingletonLvl(l)) { |
| assert(0 && "general singleton not supported yet"); |
| } else if (isNOutOfMLvl(l)) { |
| assert(0 && "n ouf of m not supported yet"); |
| } else { |
| assert(isDenseLvl(l)); |
| } |
| parentSz = assembledSize(parentSz, l); |
| } |
| |
| // Handle Aos vs. SoA mismatch for COO. |
| if (trailCOOLen != 0) { |
| uint64_t cooStartLvl = lvlRank - trailCOOLen; |
| assert(!isUniqueLvl(cooStartLvl) && |
| (isCompressedLvl(cooStartLvl) || isLooseCompressedLvl(cooStartLvl))); |
| P *posPtr = reinterpret_cast<P *>(lvlBufs[bufIdx++]); |
| C *aosCrdPtr = reinterpret_cast<C *>(lvlBufs[bufIdx++]); |
| P crdLen; |
| if (isLooseCompressedLvl(cooStartLvl)) { |
| positions[cooStartLvl].assign(posPtr, posPtr + 2 * parentSz); |
| crdLen = positions[cooStartLvl][2 * parentSz - 1]; |
| } else { |
| positions[cooStartLvl].assign(posPtr, posPtr + parentSz + 1); |
| crdLen = positions[cooStartLvl][parentSz]; |
| } |
| for (uint64_t l = cooStartLvl; l < lvlRank; l++) { |
| coordinates[l].resize(crdLen); |
| for (uint64_t n = 0; n < crdLen; n++) { |
| coordinates[l][n] = *(aosCrdPtr + (l - cooStartLvl) + n * trailCOOLen); |
| } |
| } |
| parentSz = assembledSize(parentSz, cooStartLvl); |
| } |
| |
| // Copy the values buffer. |
| V *valPtr = reinterpret_cast<V *>(lvlBufs[bufIdx]); |
| values.assign(valPtr, valPtr + parentSz); |
| } |
| |
| } // namespace sparse_tensor |
| } // namespace mlir |
| |
| #endif // MLIR_EXECUTIONENGINE_SPARSETENSOR_STORAGE_H |
