[Mlir-commits] [mlir] cf35825 - [mlir][sparse] Moving the sort from factory method to the constructor.

[wren romano]

Author: wren romano
Date: 2022-01-13T19:44:51-08:00
New Revision: cf358253887298edf44bc2eec050eb8af0f68594

URL: https://github.com/llvm/llvm-project/commit/cf358253887298edf44bc2eec050eb8af0f68594
DIFF: https://github.com/llvm/llvm-project/commit/cf358253887298edf44bc2eec050eb8af0f68594.diff

LOG: [mlir][sparse] Moving the sort from factory method to the constructor.

This guarantees the preconditions of fromCOO; whereas prior to this, one could call the constructor directly with an unsorted tensor, which would cause fromCOO to misbehave.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D117167

Added: 
    

Modified: 
    mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
index 281016f785af8..5333c23912f1a 100644
--- a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
+++ b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
@@ -106,6 +106,8 @@ struct SparseTensorCOO {
   /// 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);
   }
   /// Returns rank.
@@ -269,6 +271,8 @@ class SparseTensorStorage : public SparseTensorStorageBase {
         pointers[r].push_back(0);
     // Then assign contents from coordinate scheme tensor if provided.
     if (tensor) {
+      // Lexicographically sort the tensor, to ensure precondition of `fromCOO`.
+      tensor->sort();
       const std::vector<Element<V>> &elements = tensor->getElements();
       uint64_t nnz = elements.size();
       values.reserve(nnz);
@@ -386,7 +390,6 @@ class SparseTensorStorage : public SparseTensorStorageBase {
       assert(tensor->getRank() == rank);
       for (uint64_t r = 0; r < rank; r++)
         assert(sizes[r] == 0 || tensor->getSizes()[perm[r]] == sizes[r]);
-      tensor->sort(); // sort lexicographically
       n = new SparseTensorStorage<P, I, V>(tensor->getSizes(), perm, sparsity,
                                            tensor);
       delete tensor;
@@ -403,6 +406,7 @@ class SparseTensorStorage : public SparseTensorStorageBase {
   /// Initializes sparse tensor storage scheme from a memory-resident sparse
   /// tensor in coordinate scheme. This method prepares the pointers and
   /// indices arrays under the given per-dimension dense/sparse annotations.
+  /// Precondition: the `elements` must be lexicographically sorted.
   void fromCOO(const std::vector<Element<V>> &elements, uint64_t lo,
                uint64_t hi, uint64_t d) {
     // Once dimensions are exhausted, insert the numerical values.