[Mlir-commits] [mlir] [mlir][sparse] complete migration to dim2lvl/lvl2dim in library (PR #69268)

Aart Bik llvmlistbot at llvm.org
Mon Oct 16 20:24:17 PDT 2023


https://github.com/aartbik updated https://github.com/llvm/llvm-project/pull/69268

>From 14874810c4cfecc762fdbdbe606dfa60669c4d18 Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Mon, 16 Oct 2023 17:33:43 -0700
Subject: [PATCH 1/2] [mlir][sparse] complete migration to dim2lvl/lvl2dim in
 library

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.
---
 .../mlir/ExecutionEngine/SparseTensor/File.h  |   7 +-
 .../ExecutionEngine/SparseTensor/Storage.h    | 361 ++++--------------
 .../ExecutionEngine/SparseTensorRuntime.cpp   |   5 +-
 3 files changed, 80 insertions(+), 293 deletions(-)

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->trgCursor, src.values[parentPos]);
+  void forallElements(ElementConsumer<V> yield, uint64_t parentPos, uint64_t l,
+                      std::vector<uint64_t> &lvlCoords) const {
+    if (l == getLvlRank()) {
+      assert(parentPos < values.size());
+      yield(values[parentPos]);
       return;
     }
-    uint64_t &cursorL = this->trgCursor[this->lvl2trg[l]];
-    const auto dlt = src.getLvlType(l); // Avoid redundant bounds checking.
-    if (isCompressedDLT(dlt)) {
+    if (isCompressedLvl(l)) {
       // Look up the bounds of the `l`-level segment determined by the
       // `(l - 1)`-level position `parentPos`.
-      const std::vector<P> &positionsL = src.positions[l];
+      const std::vector<P> &positionsL = positions[l];
       assert(parentPos + 1 < positionsL.size());
       const uint64_t pstart = static_cast<uint64_t>(positionsL[parentPos]);
       const uint64_t pstop = static_cast<uint64_t>(positionsL[parentPos + 1]);
       // Loop-invariant code for looking up the `l`-level coordinates.
-      const std::vector<C> &coordinatesL = src.coordinates[l];
+      const std::vector<C> &coordinatesL = coordinates[l];
       assert(pstop <= coordinatesL.size());
       for (uint64_t pos = pstart; pos < pstop; ++pos) {
-        cursorL = static_cast<uint64_t>(coordinatesL[pos]);
-        forallElements(yield, pos, l + 1);
+        lvlCoords[l] = static_cast<uint64_t>(coordinatesL[pos]);
+        forallElements(yield, pos, l + 1, lvlCoords);
       }
-    } else if (isSingletonDLT(dlt)) {
-      cursorL = src.getCrd(l, parentPos);
-      forallElements(yield, parentPos, l + 1);
+    } else if (isSingletonLvl(l)) {
+      lvlCoords[l] = getCrd(l, parentPos);
+      forallElements(yield, parentPos, l + 1, lvlCoords);
     } else { // Dense level.
-      assert(isDenseDLT(dlt));
-      const uint64_t sz = src.getLvlSizes()[l];
+      assert(isDenseLvl(l));
+      const uint64_t sz = getLvlSizes()[l];
       const uint64_t pstart = parentPos * sz;
       for (uint64_t c = 0; c < sz; ++c) {
-        cursorL = c;
-        forallElements(yield, pstart + c, l + 1);
+        lvlCoords[l] = c;
+        forallElements(yield, pstart + c, l + 1, lvlCoords);
       }
     }
   }
+
+  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
 };
 
 //===----------------------------------------------------------------------===//
@@ -868,41 +670,24 @@ SparseTensorStorage<P, C, V> *SparseTensorStorage<P, C, V>::newEmpty(
                                           !forwarding);
 }
 
-// TODO: MapRef
 template <typename P, typename C, typename V>
 SparseTensorStorage<P, C, V> *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) {
-  assert(dimShape && dim2lvl && lvl2dim);
-  const auto &lvlSizes = lvlCOO.getDimSizes();
-  assert(lvlRank == lvlSizes.size() && "Level-rank mismatch");
-  // Must reconstruct `dimSizes` from `lvlSizes`.  While this is easy
-  // enough to do when `lvl2dim` is a permutation, this approach will
-  // not work for more general mappings; so we will need to move this
-  // computation off to codegen.
-  std::vector<uint64_t> dimSizes(dimRank);
-  for (uint64_t l = 0; l < lvlRank; ++l) {
-    const uint64_t d = lvl2dim[l];
-    assert((dimShape[d] == 0 || dimShape[d] == lvlSizes[l]) &&
-           "Dimension sizes do not match expected shape");
-    dimSizes[d] = lvlSizes[l];
-  }
-  return new SparseTensorStorage<P, C, V>(dimRank, dimSizes.data(), lvlRank,
+    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) {
+  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>::packFromLvlBuffers(
-    uint64_t dimRank, const uint64_t *dimShape, uint64_t lvlRank,
+    uint64_t dimRank, const uint64_t *dimSizes, 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) {
-  assert(dimShape && "Got nullptr for dimension shape");
-  auto *tensor =
-      new SparseTensorStorage<P, C, V>(dimRank, dimShape, lvlRank, lvlSizes,
-                                       lvlTypes, src2lvl, lvl2dim, buffers);
-  return tensor;
+    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);
 }
 
 //===----------------------------------------------------------------------===//
@@ -920,6 +705,7 @@ SparseTensorStorage<P, C, V>::SparseTensorStorage(
     : SparseTensorStorage(dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes,
                           dim2lvl, lvl2dim) {
   lvlCOO = coo;
+  assert(!lvlCOO || lvlRank == lvlCOO->getRank());
   // 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
@@ -948,17 +734,16 @@ SparseTensorStorage<P, C, V>::SparseTensorStorage(
     values.resize(sz, 0);
 }
 
-// TODO: share more code with forwarding methods?
 template <typename P, typename C, typename V>
 SparseTensorStorage<P, C, V>::SparseTensorStorage( // NOLINT
     uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
-    const DimLevelType *lvlTypes, const uint64_t *dim2lvl,
-    const uint64_t *lvl2dim, SparseTensorCOO<V> &lvlCOO)
-    : SparseTensorStorage(dimRank, dimSizes, lvlRank,
-                          lvlCOO.getDimSizes().data(), lvlTypes, dim2lvl,
-                          lvl2dim, nullptr, false) {
+    const uint64_t *lvlSizes, const DimLevelType *lvlTypes,
+    const uint64_t *dim2lvl, const uint64_t *lvl2dim,
+    SparseTensorCOO<V> &lvlCOO)
+    : SparseTensorStorage(dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes,
+                          dim2lvl, lvl2dim, nullptr, false) {
   // Ensure lvlCOO is sorted.
-  assert(lvlRank == lvlCOO.getDimSizes().size() && "Level-rank mismatch");
+  assert(lvlRank == lvlCOO.getRank());
   lvlCOO.sort();
   // Now actually insert the `elements`.
   const auto &elements = lvlCOO.getElements();
diff --git a/mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp b/mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp
index 6a4c0f292c5f81e..36d888a08de6d60 100644
--- a/mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp
+++ b/mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp
@@ -129,7 +129,8 @@ extern "C" {
       assert(ptr && "Received nullptr for SparseTensorCOO object");            \
       auto &coo = *static_cast<SparseTensorCOO<V> *>(ptr);                     \
       return SparseTensorStorage<P, C, V>::newFromCOO(                         \
-          dimRank, dimSizes, lvlRank, lvlTypes, dim2lvl, lvl2dim, coo);        \
+          dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes, dim2lvl, lvl2dim,    \
+          coo);                                                                \
     }                                                                          \
     case Action::kFromReader: {                                                \
       assert(ptr && "Received nullptr for SparseTensorReader object");         \
@@ -140,7 +141,7 @@ extern "C" {
     case Action::kToCOO: {                                                     \
       assert(ptr && "Received nullptr for SparseTensorStorage object");        \
       auto &tensor = *static_cast<SparseTensorStorage<P, C, V> *>(ptr);        \
-      return tensor.toCOO(lvlRank, lvlSizes, dimRank, dim2lvl, lvl2dim);       \
+      return tensor.toCOO();                                                   \
     }                                                                          \
     case Action::kPack: {                                                      \
       assert(ptr && "Received nullptr for SparseTensorStorage object");        \

>From 1c315ee9be7a9897b06bd709a2145b6509b280fb Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Mon, 16 Oct 2023 20:23:52 -0700
Subject: [PATCH 2/2] simpler recursion

---
 .../ExecutionEngine/SparseTensor/Storage.h    | 66 ++++++++-----------
 1 file changed, 29 insertions(+), 37 deletions(-)

diff --git a/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h b/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h
index f31ac7c1b7cc246..f1aeb12c662fdc1 100644
--- a/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h
+++ b/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h
@@ -26,9 +26,6 @@
 namespace mlir {
 namespace sparse_tensor {
 
-/// The type of callback functions which receive an element.
-template <typename V> using ElementConsumer = const std::function<void(V)> &;
-
 //===----------------------------------------------------------------------===//
 //
 //  SparseTensorStorage Classes
@@ -207,8 +204,7 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
                       const uint64_t *lvl2dim)
       : SparseTensorStorageBase(dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes,
                                 dim2lvl, lvl2dim),
-        positions(lvlRank), coordinates(lvlRank), lvlCursor(lvlRank), lvlCOO() {
-  }
+        positions(lvlRank), coordinates(lvlRank), lvlCursor(lvlRank), coo() {}
 
 public:
   /// Constructs a sparse tensor with the given encoding, and allocates
@@ -221,7 +217,7 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
   SparseTensorStorage(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> *coo,
+                      const uint64_t *lvl2dim, SparseTensorCOO<V> *lvlCOO,
                       bool initializeValuesIfAllDense);
 
   /// Constructs a sparse tensor with the given encoding, and initializes
@@ -293,15 +289,14 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
 
   /// Partially specialize forwarding insertions based on template types.
   void forwardingInsert(const uint64_t *dimCoords, V val) final {
-    assert(dimCoords && lvlCOO);
+    assert(dimCoords && coo);
     map.pushforward(dimCoords, lvlCursor.data());
-    lvlCOO->add(lvlCursor, val);
+    coo->add(lvlCursor, val);
   }
 
   /// Partially specialize lexicographical insertions based on template types.
   void lexInsert(const uint64_t *lvlCoords, V val) final {
     assert(lvlCoords);
-    // TODO: needed?
     bool allDense = std::all_of(getLvlTypes().begin(), getLvlTypes().end(),
                                 [](DimLevelType lt) { return isDenseDLT(lt); });
     if (allDense) {
@@ -357,16 +352,17 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
 
   /// Finalizes forwarding insertions.
   void endForwardingInsert() final {
-    // Ensure lvlCOO is sorted.
-    assert(lvlCOO);
-    lvlCOO->sort();
+    // Ensure COO is sorted.
+    assert(coo);
+    coo->sort();
     // Now actually insert the `elements`.
-    const auto &elements = lvlCOO->getElements();
+    const auto &elements = coo->getElements();
     const uint64_t nse = elements.size();
     assert(values.size() == 0);
     values.reserve(nse);
     fromCOO(elements, 0, nse, 0);
-    delete lvlCOO;
+    delete coo;
+    coo = nullptr;
   }
 
   /// Finalizes lexicographic insertions.
@@ -379,16 +375,10 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
 
   /// 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 {
+  SparseTensorCOO<V> *toCOO() {
     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);
+    coo = new SparseTensorCOO<V>(getDimSizes(), values.size());
+    toCOO(0, 0, dimCoords);
     assert(coo->getElements().size() == values.size());
     return coo;
   }
@@ -609,11 +599,13 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
     return -1u;
   }
 
-  void forallElements(ElementConsumer<V> yield, uint64_t parentPos, uint64_t l,
-                      std::vector<uint64_t> &lvlCoords) const {
+  // Performs forall on level entries and inserts into dim COO.
+  void toCOO(uint64_t parentPos, uint64_t l, std::vector<uint64_t> &dimCoords) {
     if (l == getLvlRank()) {
+      map.pushbackward(lvlCursor.data(), dimCoords.data());
+      assert(coo);
       assert(parentPos < values.size());
-      yield(values[parentPos]);
+      coo->add(dimCoords, values[parentPos]);
       return;
     }
     if (isCompressedLvl(l)) {
@@ -627,19 +619,19 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
       const std::vector<C> &coordinatesL = coordinates[l];
       assert(pstop <= coordinatesL.size());
       for (uint64_t pos = pstart; pos < pstop; ++pos) {
-        lvlCoords[l] = static_cast<uint64_t>(coordinatesL[pos]);
-        forallElements(yield, pos, l + 1, lvlCoords);
+        lvlCursor[l] = static_cast<uint64_t>(coordinatesL[pos]);
+        toCOO(pos, l + 1, dimCoords);
       }
     } else if (isSingletonLvl(l)) {
-      lvlCoords[l] = getCrd(l, parentPos);
-      forallElements(yield, parentPos, l + 1, lvlCoords);
+      lvlCursor[l] = getCrd(l, parentPos);
+      toCOO(parentPos, l + 1, dimCoords);
     } else { // Dense level.
       assert(isDenseLvl(l));
       const uint64_t sz = getLvlSizes()[l];
       const uint64_t pstart = parentPos * sz;
       for (uint64_t c = 0; c < sz; ++c) {
-        lvlCoords[l] = c;
-        forallElements(yield, pstart + c, l + 1, lvlCoords);
+        lvlCursor[l] = c;
+        toCOO(pstart + c, l + 1, dimCoords);
       }
     }
   }
@@ -647,8 +639,8 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
   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
+  std::vector<uint64_t> lvlCursor;
+  SparseTensorCOO<V> *coo;
 };
 
 //===----------------------------------------------------------------------===//
@@ -700,12 +692,12 @@ 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 DimLevelType *lvlTypes,
-    const uint64_t *dim2lvl, const uint64_t *lvl2dim, SparseTensorCOO<V> *coo,
-    bool initializeValuesIfAllDense)
+    const uint64_t *dim2lvl, const uint64_t *lvl2dim,
+    SparseTensorCOO<V> *lvlCOO, bool initializeValuesIfAllDense)
     : SparseTensorStorage(dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes,
                           dim2lvl, lvl2dim) {
-  lvlCOO = coo;
   assert(!lvlCOO || lvlRank == lvlCOO->getRank());
+  coo = lvlCOO;
   // 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



More information about the Mlir-commits mailing list