[Mlir-commits] [mlir] [mlir][sparse] complete migration to dim2lvl/lvl2dim in library (PR #69268)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 16 17:37:03 PDT 2023
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-sparse
Author: Aart Bik (aartbik)
<details>
<summary>Changes</summary>
This last revision completed the migration to non-permutation support in the SparseTensor library. All mappings are now controlled by the MapRef (forward and backward). Unused code has been removed, which simplifies subsequent testing of block sparsity.
---
Patch is 28.40 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/69268.diff
3 Files Affected:
- (modified) mlir/include/mlir/ExecutionEngine/SparseTensor/File.h (+4-3)
- (modified) mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h (+73-288)
- (modified) mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp (+3-2)
``````````diff
diff --git a/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h b/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
index efc3f82d6a307ea..1b5f0553a3af959 100644
--- a/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
+++ b/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
@@ -201,10 +201,11 @@ class SparseTensorReader final {
const uint64_t *lvl2dim) {
const uint64_t dimRank = getRank();
MapRef map(dimRank, lvlRank, dim2lvl, lvl2dim);
- auto *coo = readCOO<V>(map, lvlSizes);
+ auto *lvlCOO = readCOO<V>(map, lvlSizes);
auto *tensor = SparseTensorStorage<P, I, V>::newFromCOO(
- dimRank, getDimSizes(), lvlRank, lvlTypes, dim2lvl, lvl2dim, *coo);
- delete coo;
+ dimRank, getDimSizes(), lvlRank, lvlSizes, lvlTypes, dim2lvl, lvl2dim,
+ *lvlCOO);
+ delete lvlCOO;
return tensor;
}
diff --git a/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h b/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h
index bafc9baa7edde1e..f31ac7c1b7cc246 100644
--- a/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h
+++ b/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h
@@ -10,8 +10,6 @@
//
// * `SparseTensorStorageBase`
// * `SparseTensorStorage<P, C, V>`
-// * `SparseTensorEnumeratorBase<V>`
-// * `SparseTensorEnumerator<P, C, V>`
//
//===----------------------------------------------------------------------===//
@@ -29,25 +27,17 @@ namespace mlir {
namespace sparse_tensor {
/// The type of callback functions which receive an element.
-template <typename V>
-using ElementConsumer =
- const std::function<void(const std::vector<uint64_t> &, V)> &;
-
-// Forward references.
-template <typename V>
-class SparseTensorEnumeratorBase;
-template <typename P, typename C, typename V>
-class SparseTensorEnumerator;
+template <typename V> using ElementConsumer = const std::function<void(V)> &;
//===----------------------------------------------------------------------===//
//
-// SparseTensorStorage
+// 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., shape, sparsity, mapping). In addition,
+/// 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.
@@ -55,7 +45,7 @@ class SparseTensorEnumerator;
/// 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, shape, sizes, etc).
+/// 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.
@@ -139,10 +129,6 @@ class SparseTensorStorageBase {
/// Safely checks if the level is unique.
bool isUniqueLvl(uint64_t l) const { return isUniqueDLT(getLvlType(l)); }
- /// Gets the level-to-dimension mapping.
- // TODO: REMOVE THIS
- const std::vector<uint64_t> &getLvl2Dim() const { return lvl2dimVec; }
-
/// Gets positions-overhead storage for the given level.
#define DECL_GETPOSITIONS(PNAME, P) \
virtual void getPositions(std::vector<P> **, uint64_t);
@@ -154,6 +140,7 @@ class SparseTensorStorageBase {
virtual void getCoordinates(std::vector<C> **, uint64_t);
MLIR_SPARSETENSOR_FOREVERY_FIXED_O(DECL_GETCOORDINATES)
#undef DECL_GETCOORDINATES
+
/// Gets the coordinate-value stored at the given level and position.
virtual uint64_t getCrd(uint64_t lvl, uint64_t pos) const = 0;
@@ -241,17 +228,9 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
/// the contents from the COO. This ctor performs the same heuristic
/// overhead-storage allocation as the ctor above.
SparseTensorStorage(uint64_t dimRank, const uint64_t *dimSizes,
- uint64_t lvlRank, const DimLevelType *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 enumerator. This ctor allocates exactly
- /// the required amount of overhead storage, not using any heuristics.
- SparseTensorStorage(uint64_t dimRank, const uint64_t *dimSizes,
- uint64_t lvlRank, const DimLevelType *lvlTypes,
- const uint64_t *dim2lvl, const uint64_t *lvl2dim,
- SparseTensorEnumeratorBase<V> &lvlEnumerator);
+ uint64_t lvlRank, const uint64_t *lvlSizes,
+ const DimLevelType *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. This ctor allocates exactly
@@ -265,39 +244,27 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
const DimLevelType *lvlTypes, const uint64_t *dim2lvl,
const uint64_t *lvl2dim, const intptr_t *lvlBufs);
- /// Allocates a new empty sparse tensor. The preconditions/assertions
- /// are as per the `SparseTensorStorageBase` ctor; which is to say,
- /// the `dimSizes` and `lvlSizes` must both be "sizes" not "shapes",
- /// since there's nowhere to reconstruct dynamic sizes from.
+ /// 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 DimLevelType *lvlTypes,
const uint64_t *dim2lvl, const uint64_t *lvl2dim, bool forwarding);
/// Allocates a new sparse tensor and initializes it from the given COO.
- /// The preconditions are as per the `SparseTensorStorageBase` ctor
- /// (where we define `lvlSizes = lvlCOO.getDimSizes().data()`), but
- /// using the following assertions in lieu of the base ctor's assertions:
- //
- // TODO: The ability to reconstruct dynamic dimensions-sizes does not
- // easily generalize to arbitrary `lvl2dim` mappings. When compiling
- // MLIR programs to use this library, we should be able to generate
- // code for effectively computing the reconstruction, but it's not clear
- // that there's a feasible way to do so from within the library itself.
- // Therefore, when we functionalize the `lvl2dim` mapping we'll have
- // to update the type/preconditions of this factory too.
static SparseTensorStorage<P, C, V> *
- newFromCOO(uint64_t dimRank, const uint64_t *dimShape, uint64_t lvlRank,
- const DimLevelType *lvlTypes, const uint64_t *dim2lvl,
- const uint64_t *lvl2dim, SparseTensorCOO<V> &lvlCOO);
+ newFromCOO(uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
+ const uint64_t *lvlSizes, const DimLevelType *lvlTypes,
+ const uint64_t *dim2lvl, const uint64_t *lvl2dim,
+ SparseTensorCOO<V> &lvlCOO);
/// Allocates a new sparse tensor and initialize it with the data stored level
/// buffers directly.
- static SparseTensorStorage<P, C, V> *packFromLvlBuffers(
- uint64_t dimRank, const uint64_t *dimShape, uint64_t lvlRank,
- const uint64_t *lvlSizes, const DimLevelType *lvlTypes,
- const uint64_t *src2lvl, // FIXME: dim2lvl
- const uint64_t *lvl2dim, uint64_t srcRank, const intptr_t *buffers);
+ static SparseTensorStorage<P, C, V> *
+ packFromLvlBuffers(uint64_t dimRank, const uint64_t *dimSizes,
+ uint64_t lvlRank, const uint64_t *lvlSizes,
+ const DimLevelType *lvlTypes, const uint64_t *dim2lvl,
+ const uint64_t *lvl2dim, uint64_t srcRank,
+ const intptr_t *buffers);
~SparseTensorStorage() final = default;
@@ -334,8 +301,7 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
/// Partially specialize lexicographical insertions based on template types.
void lexInsert(const uint64_t *lvlCoords, V val) final {
assert(lvlCoords);
- // TODO: get rid of this! canonicalize all-dense "sparse" array into dense
- // tensors.
+ // TODO: needed?
bool allDense = std::all_of(getLvlTypes().begin(), getLvlTypes().end(),
[](DimLevelType lt) { return isDenseDLT(lt); });
if (allDense) {
@@ -411,23 +377,18 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
endPath(0);
}
- /// Allocates a new COO object and initializes it with the contents
- /// of this tensor under the given mapping from the `getDimSizes()`
- /// coordinate-space to the `trgSizes` coordinate-space. Callers must
- /// make sure to delete the COO when they're done with it.
- SparseTensorCOO<V> *toCOO(uint64_t trgRank, const uint64_t *trgSizes,
- uint64_t srcRank,
- const uint64_t *src2trg, // FIXME: dim2lvl
- const uint64_t *lvl2dim) const {
- // TODO: use MapRef here too for the translation
- SparseTensorEnumerator<P, C, V> enumerator(*this, trgRank, trgSizes,
- srcRank, src2trg);
- auto *coo = new SparseTensorCOO<V>(trgRank, trgSizes, values.size());
- enumerator.forallElements(
- [&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.
- // <https://github.com/llvm/llvm-project/issues/54179>
+ /// Allocates a new COO object and initializes it with the contents.
+ /// Callers must make sure to delete the COO when they're done with it.
+ SparseTensorCOO<V> *toCOO() const {
+ std::vector<uint64_t> dimCoords(getDimRank());
+ std::vector<uint64_t> lvlCoords(getLvlRank());
+ auto *coo = new SparseTensorCOO<V>(getDimSizes(), values.size());
+ forallElements(
+ [this, &dimCoords, &lvlCoords, &coo](V val) {
+ map.pushbackward(lvlCoords.data(), dimCoords.data());
+ coo->add(dimCoords, val);
+ },
+ 0, 0, lvlCoords);
assert(coo->getElements().size() == values.size());
return coo;
}
@@ -525,27 +486,11 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
}
}
- /// Writes the given coordinate to `coordinates[lvl][pos]`. This method
- /// checks that `crd` is representable in the `C` type; however, it
- /// does not check that `crd` is semantically valid (i.e., in bounds
- /// for `dimSizes[lvl]` and not elsewhere occurring in the same segment).
- void writeCrd(uint64_t lvl, uint64_t pos, uint64_t crd) {
- assert(isCompressedDLT(getLvlType(lvl)) || isSingletonDLT(getLvlType(lvl)));
- // Subscript assignment to `std::vector` requires that the `pos`-th
- // entry has been initialized; thus we must be sure to check `size()`
- // here, instead of `capacity()` as would be ideal.
- assert(pos < coordinates[lvl].size());
- coordinates[lvl][pos] = detail::checkOverflowCast<C>(crd);
- }
-
/// Computes the assembled-size associated with the `l`-th level,
/// given the assembled-size associated with the `(l-1)`-th level.
/// "Assembled-sizes" correspond to the (nominal) sizes of overhead
/// storage, as opposed to "level-sizes" which are the cardinality
/// of possible coordinates for that level.
- ///
- /// Precondition: the `positions[l]` array must be fully initialized
- /// before calling this method.
uint64_t assembledSize(uint64_t parentSz, uint64_t l) const {
const auto dlt = getLvlType(l); // Avoid redundant bounds checking.
if (isCompressedDLT(dlt))
@@ -553,7 +498,7 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
if (isSingletonDLT(dlt))
return parentSz; // New size is same as the parent.
if (isDenseDLT(dlt))
- return parentSz * getLvlSizes()[l];
+ return parentSz * getLvlSize(l);
MLIR_SPARSETENSOR_FATAL("unsupported level type: %d\n",
static_cast<uint8_t>(dlt));
}
@@ -561,11 +506,6 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
/// 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.
- ///
- /// Preconditions:
- /// * the `lvlElements` must be lexicographically sorted.
- /// * the coordinates of every element are valid for `getLvlSizes()`
- /// (i.e., equal rank and pointwise less-than).
void fromCOO(const std::vector<Element<V>> &lvlElements, uint64_t lo,
uint64_t hi, uint64_t l) {
const uint64_t lvlRank = getLvlRank();
@@ -669,184 +609,46 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
return -1u;
}
- // Allow `SparseTensorEnumerator` to access the data-members (to avoid
- // the cost of virtual-function dispatch in inner loops), without
- // making them public to other client code.
- friend class SparseTensorEnumerator<P, C, V>;
-
- std::vector<std::vector<P>> positions;
- std::vector<std::vector<C>> coordinates;
- std::vector<V> values;
- std::vector<uint64_t> lvlCursor; // cursor for lexicographic insertion.
- SparseTensorCOO<V> *lvlCOO; // COO used during forwarding
-};
-
-//===----------------------------------------------------------------------===//
-//
-// SparseTensorEnumerator
-//
-//===----------------------------------------------------------------------===//
-
-/// A (higher-order) function object for enumerating the elements of some
-/// `SparseTensorStorage` under a permutation. That is, the `forallElements`
-/// method encapsulates the loop-nest for enumerating the elements of
-/// the source tensor (in whatever order is best for the source tensor),
-/// and applies a permutation to the coordinates before handing
-/// each element to the callback. A single enumerator object can be
-/// freely reused for several calls to `forallElements`, just so long
-/// as each call is sequential with respect to one another.
-///
-/// N.B., this class stores a reference to the `SparseTensorStorageBase`
-/// passed to the constructor; thus, objects of this class must not
-/// outlive the sparse tensor they depend on.
-///
-/// Design Note: The reason we define this class instead of simply using
-/// `SparseTensorEnumerator<P,C,V>` is because we need to hide/generalize
-/// the `<P,C>` template parameters from MLIR client code (to simplify the
-/// type parameters used for direct sparse-to-sparse conversion). And the
-/// reason we define the `SparseTensorEnumerator<P,C,V>` subclasses rather
-/// than simply using this class, is to avoid the cost of virtual-method
-/// dispatch within the loop-nest.
-template <typename V>
-class SparseTensorEnumeratorBase {
-public:
- /// Constructs an enumerator which automatically applies the given
- /// mapping from the source tensor's dimensions to the desired
- /// target tensor dimensions.
- ///
- /// Preconditions:
- /// * the `src` must have the same `V` value type.
- /// * `trgSizes` must be valid for `trgRank`.
- /// * `src2trg` must be valid for `srcRank`, and must map coordinates
- /// valid for `src.getDimSizes()` to coordinates valid for `trgSizes`.
- ///
- /// Asserts:
- /// * `trgSizes` must be nonnull and must contain only nonzero sizes.
- /// * `srcRank == src.getDimRank()`.
- /// * `src2trg` must be nonnull.
- SparseTensorEnumeratorBase(const SparseTensorStorageBase &src,
- uint64_t trgRank, const uint64_t *trgSizes,
- uint64_t srcRank, const uint64_t *src2trg)
- : src(src), trgSizes(trgSizes, trgSizes + trgRank),
- lvl2trg(src.getLvlRank()), trgCursor(trgRank) {
- assert(trgSizes && "Received nullptr for target-sizes");
- assert(src2trg && "Received nullptr for source-to-target mapping");
- assert(srcRank == src.getDimRank() && "Source-rank mismatch");
- for (uint64_t t = 0; t < trgRank; ++t)
- assert(trgSizes[t] > 0 && "Target-size zero has trivial storage");
- const auto &lvl2src = src.getLvl2Dim();
- for (uint64_t lvlRank = src.getLvlRank(), l = 0; l < lvlRank; ++l)
- lvl2trg[l] = src2trg[lvl2src[l]];
- }
-
- virtual ~SparseTensorEnumeratorBase() = default;
-
- // We disallow copying to help avoid leaking the `src` reference.
- // (In addition to avoiding the problem of slicing.)
- SparseTensorEnumeratorBase(const SparseTensorEnumeratorBase &) = delete;
- SparseTensorEnumeratorBase &
- operator=(const SparseTensorEnumeratorBase &) = delete;
-
- /// Gets the source's dimension-rank.
- uint64_t getSrcDimRank() const { return src.getDimRank(); }
-
- /// Gets the target's dimension-/level-rank. (This is usually
- /// "dimension-rank", though that may coincide with "level-rank"
- /// depending on usage.)
- uint64_t getTrgRank() const { return trgSizes.size(); }
-
- /// Gets the target's dimension-/level-sizes. (These are usually
- /// "dimensions", though that may coincide with "level-rank" depending
- /// on usage.)
- const std::vector<uint64_t> &getTrgSizes() const { return trgSizes; }
-
- /// Enumerates all elements of the source tensor, permutes their
- /// coordinates, and passes the permuted element to the callback.
- /// The callback must not store the cursor reference directly,
- /// since this function reuses the storage. Instead, the callback
- /// must copy it if they want to keep it.
- virtual void forallElements(ElementConsumer<V> yield) = 0;
-
-protected:
- const SparseTensorStorageBase &src;
- std::vector<uint64_t> trgSizes; // in target order.
- std::vector<uint64_t> lvl2trg; // source-levels -> target-dims/lvls.
- std::vector<uint64_t> trgCursor; // in target order.
-};
-
-template <typename P, typename C, typename V>
-class SparseTensorEnumerator final : public SparseTensorEnumeratorBase<V> {
- using Base = SparseTensorEnumeratorBase<V>;
- using StorageImpl = SparseTensorStorage<P, C, V>;
-
-public:
- /// Constructs an enumerator which automatically applies the given
- /// mapping from the source tensor's dimensions to the desired
- /// target tensor dimensions.
- ///
- /// Preconditions/assertions are as per the `SparseTensorEnumeratorBase` ctor.
- SparseTensorEnumerator(const StorageImpl &src, uint64_t trgRank,
- const uint64_t *trgSizes, uint64_t srcRank,
- const uint64_t *src2trg)
- : Base(src, trgRank, trgSizes, srcRank, src2trg) {}
-
- ~SparseTensorEnumerator() final = default;
-
- void forallElements(ElementConsumer<V> yield) final {
- forallElements(yield, 0, 0);
- }
-
-private:
- // TODO: Once we functionalize the mappings, then we'll no longer
- // be able to use the current approach of constructing `lvl2trg` in the
- // ctor and using it to incrementally fill the `trgCursor` cursor as we
- // recurse through `forallElements`. Instead we'll want to incrementally
- // fill a `lvlCursor` as we recurse, and then use `src.getLvl2Dim()`
- // and `src2trg` to convert it just before yielding to the callback.
- // It's probably most efficient to just store the `srcCursor` and
- // `trgCursor` buffers in this object, but we may want to benchmark
- // that against using `std::calloc` to stack-allocate them instead.
- //
- /// The recursive component of the public `forallElements`.
- void forallElements(ElementConsumer<V> yield, uint64_t parentPos,
- uint64_t l) {
- // Recover the `<P,C,V>` type parameters of `src`.
- const auto &src = static_cast<const StorageImpl &>(this->src);
- if (l == src.getLvlRank()) {
- assert(parentPos < src.values.size());
- // TODO: <https://github.com/llvm/llvm-project/issues/54179>
- yield(this->tr...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/69268
More information about the Mlir-commits
mailing list