[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

68

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

69

maybe move these to also to constructor as initializer?

[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

 /// 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 {
   const std::vector<uint64_t> dimSizes; // per-dimension sizes
   std::vector<Element<V>> elements;     // all elements
   std::vector<uint64_t> indices;        // shared index pool
+  bool isSorted = false;
   bool iteratorLocked = false;
   unsigned iteratorPos = 0;
 
 public:
   SparseTensorCOO(const std::vector<uint64_t> &dimSizes, uint64_t capacity)
       : dimSizes(dimSizes) {
     if (capacity) {
       elements.reserve(capacity);
       indices.reserve(capacity * getRank());
void add(const std::vector<uint64_t> &ind, V val) {
     // 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.
+    // TODO: Only invalidate the `isSorted` bit if the new element isn't
+    // greater than the immediately preceding one.
+    isSorted = false;
     elements.emplace_back(base + size, val);
   }
 
   /// 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.
+    if (isSorted)
+      return;
     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;
               });
+    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;
   }