[mlir][sparse] Moving lexOrder from SparseTensorCOO to Element

[wrengr]

Summary

This is the more logical place for the function to live. If/when we factor out a separate class for just the Coordinates themselves, then the definition should be moved to Coordinates::lexOrder (and Element::lexOrder would become a thin wrapper delegating to that function).

This is (tangentially) work towards fixing: https://github.com/llvm/llvm-project/issues/51652

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

85

I think technically methods appear before data, but I find what you did not readable as well, so good to go.

[wrengr]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

85

I'm not sure I follow. Did you want me to move the function a couple lines up, or leave it as is?

[wrengr]

rebase

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

 ///   ({i}, a[i])
 /// and a rank-5 tensor element like
 ///   ({i,j,k,l,m}, a[i,j,k,l,m])
 template <typename V>
 struct Element {
   Element(const std::vector<uint64_t> &ind, V val) : indices(ind), value(val){};
   std::vector<uint64_t> indices;
   V value;
+  /// Returns true if indices of e1 < indices of e2.
+  static bool lexOrder(const Element<V> &e1, const Element<V> &e2) {
+    uint64_t rank = e1.indices.size();
+    assert(rank == e2.indices.size());
+    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;
+  }
 };
 
 /// A memory-resident sparse tensor in coordinate scheme (collection of
 /// elements). This data structure is used to read a sparse tensor from
 /// any external format into memory and sort the elements lexicographically
 /// by indices before passing it back to the client (most packed storage
 /// formats require the elements to appear in lexicographic index order).
 template <typename V>
for (uint64_t r = 0; r < rank; r++)
       assert(ind[r] < sizes[r]); // within bounds
     elements.emplace_back(ind, val);
   }
   /// Sorts elements lexicographically by index.
   void sort() {
     assert(!iteratorLocked && "Attempt to sort() after startIterator()");
     // TODO: we may want to cache an `isSorted` bit, to avoid
     // unnecessary/redundant sorting.
-    std::sort(elements.begin(), elements.end(), lexOrder);
+    std::sort(elements.begin(), elements.end(), Element<V>::lexOrder);
   }
   /// Returns rank.
   uint64_t getRank() const { return sizes.size(); }
   /// Getter for sizes array.
   const std::vector<uint64_t> &getSizes() const { return sizes; }
   /// Getter for elements array.
   const std::vector<Element<V>> &getElements() const { return elements; }
 
static SparseTensorCOO<V> *newSparseTensorCOO(uint64_t rank,
                                                 uint64_t capacity = 0) {
     std::vector<uint64_t> permsz(rank);
     for (uint64_t r = 0; r < rank; r++)
       permsz[perm[r]] = sizes[r];
     return new SparseTensorCOO<V>(permsz, capacity);
   }
 
 private:
-  /// Returns true if indices of e1 < indices of e2.
-  static bool lexOrder(const Element<V> &e1, const Element<V> &e2) {
-    uint64_t rank = e1.indices.size();
-    assert(rank == e2.indices.size());
-    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;
-  }
   const std::vector<uint64_t> sizes; // per-dimension sizes
   std::vector<Element<V>> elements;
   bool iteratorLocked;
   unsigned iteratorPos;
 };
 
 /// Abstract base class of sparse tensor storage. Note that we use
 /// function overloading to implement "partial" method specialization.