[mlir][sparse] Adding isSorted bit to SparseTensorCOO

[wrengr]

Summary

This is a followup to the refactoring of D133462, D133830, D133831, and D133833.

Depends On D133833

[aartbik]

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

84–85

maybe move these to also to constructor as initializer?

85

nit: isSorted(true), also moved to L74 as initializer? technically, when we start empty, we are sorted, this would also play nicer with the TODO below, where only the first out of order "add" invalides the flag

[wrengr]

Finished TODO: the add method only invalidates the isSorted bit if the new element is not strictly-greater than the previously final element.

[wrengr]

rebase

[aartbik]

Very neat! This will avoid the sort step in many cases where the external format already appears in the right order.

[wrengr]

rebase

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

 /// in this object).
 template <typename V>
 struct Element final {
   Element(const uint64_t *ind, V val) : indices(ind), value(val){};
   const uint64_t *indices; // pointer into shared index pool
   V value;
 };
 
+/// Closure object for `operator<` on `Element` with a given rank.
+template <typename V>
+struct ElementLT final {
+  ElementLT(uint64_t rank) : rank(rank) {}
+
+  /// Compare two elements a la `operator<`.
+  ///
+  /// Precondition: the elements must both be valid for `rank`.
+  bool operator()(const Element<V> &e1, const Element<V> &e2) const {
+    for (uint64_t d = 0; d < rank; ++d) {
+      if (e1.indices[d] == e2.indices[d])
+        continue;
+      return e1.indices[d] < e2.indices[d];
+    }
+    return false;
+  }
+
+  const uint64_t rank;
+};
+
 /// The type of callback functions which receive an element.  We avoid
 /// packaging the coordinates and value together as an `Element` object
 /// because this helps keep code somewhat cleaner.
 template <typename V>
 using ElementConsumer =
     const std::function<void(const std::vector<uint64_t> &, V)> &;
 
 /// A memory-resident sparse tensor in coordinate-scheme representation
 /// (a collection of `Element`s).  This data structure is used as
 /// an intermediate representation; e.g., for reading sparse tensors
 /// from external formats into memory, or for certain conversions between
 /// different `SparseTensorStorage` formats.
 template <typename V>
 class SparseTensorCOO final {
 public:
   SparseTensorCOO(const std::vector<uint64_t> &dimSizes, uint64_t capacity)
-      : dimSizes(dimSizes) {
+      : dimSizes(dimSizes), isSorted(true), iteratorLocked(false),
+        iteratorPos(0) {
     if (capacity) {
       elements.reserve(capacity);
       indices.reserve(capacity * getRank());
     }
   }
 
   /// Factory method. Permutes the original dimensions according to
   /// the given ordering and expects subsequent add() calls to honor
public:
   uint64_t getRank() const { return dimSizes.size(); }
 
   /// Get the dimension-sizes array.
   const std::vector<uint64_t> &getDimSizes() const { return dimSizes; }
 
   /// Get the elements array.
   const std::vector<Element<V>> &getElements() const { return elements; }
 
+  /// Returns the `operator<` closure object for the COO's element type.
+  ElementLT<V> getElementLT() const { return ElementLT<V>(getRank()); }
+
   /// Adds an element to the tensor.  This method does not check whether
   /// `ind` is already associated with a value, it adds it regardless.
   /// Resolving such conflicts is left up to clients of the iterator
   /// interface.
   ///
   /// Asserts:
   /// * is not in iterator mode
   /// * the `ind` is valid for `rank`
void add(const std::vector<uint64_t> &ind, V val) {
     // for every internal vector reallocation (which with the doubling rule
     // should only incur an amortized linear overhead).
     const uint64_t *newBase = indices.data();
     if (newBase != base) {
       for (uint64_t i = 0, n = elements.size(); i < n; ++i)
         elements[i].indices = newBase + (elements[i].indices - base);
       base = newBase;
     }
-    // Add element as (pointer into shared index pool, value) pair.
-    elements.emplace_back(base + size, val);
+    // Add the new element and update the sorted bit.
+    Element<V> addedElem(base + size, val);
+    if (!elements.empty() && isSorted)
+      isSorted = getElementLT()(elements.back(), addedElem);
+    elements.push_back(addedElem);
   }
 
   /// Sorts elements lexicographically by index.  If an index is mapped to
   /// multiple values, then the relative order of those values is unspecified.
   ///
   /// Asserts: is not in iterator mode.
   void sort() {
     assert(!iteratorLocked && "Attempt to sort() after startIterator()");
-    // TODO: we may want to cache an `isSorted` bit, to avoid
-    // unnecessary/redundant sorting.
-    uint64_t rank = getRank();
-    std::sort(elements.begin(), elements.end(),
-              [rank](const Element<V> &e1, const Element<V> &e2) {
-                for (uint64_t r = 0; r < rank; ++r) {
-                  if (e1.indices[r] == e2.indices[r])
-                    continue;
-                  return e1.indices[r] < e2.indices[r];
-                }
-                return false;
-              });
+    if (isSorted)
+      return;
+    std::sort(elements.begin(), elements.end(), getElementLT());
+    isSorted = true;
   }
 
   /// Switch into iterator mode.  If already in iterator mode, then
   /// resets the position to the first element.
   void startIterator() {
     iteratorLocked = true;
     iteratorPos = 0;
   }
const Element<V> *getNext() {
     iteratorLocked = false;
     return nullptr;
   }
 
 private:
   const std::vector<uint64_t> dimSizes; // per-dimension sizes
   std::vector<Element<V>> elements;     // all COO elements
   std::vector<uint64_t> indices;        // shared index pool
-  bool iteratorLocked = false;
-  unsigned iteratorPos = 0;
+  bool isSorted;
+  bool iteratorLocked;
+  unsigned iteratorPos;
 };
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_EXECUTIONENGINE_SPARSETENSOR_COO_H