[mlir][sparse] add asCOO() functionality to sparse tensor object

[aartbik]

Summary

This prepares general sparse to sparse conversions. The code that
needs to be generated using this new feature is now simply:

(1) coo = sparse_tensor_1->asCOO(); source format1
(2) sparse_tensor_2 = newSparseTensor(coo); destination format2

By using COO as an intermediate, we can do *all* conversions without
having to implement the full O(N^2) conversion matrix. Note that we
can always improve particular conversions individually if a faster
solution is required.

[aartbik]

mlir/lib/ExecutionEngine/SparseUtils.cpp

297

Please note that this implementation does not take the permutation of either source or destination into account yet (i.e. this version is for the identity permutations). A follow-up CL will dot all the i's and cross the t's, followed by actual lowering code support in SparseTensorConversion.

[aartbik]

cross the ts and dot the is

mlir/lib/ExecutionEngine/SparseUtils.cpp

297

I added the required permutation requirements in this CL also. The next revision that will build on this will finalize the sparse conversions (with just a few lines of code).

[bixia]

mlir/lib/ExecutionEngine/SparseUtils.cpp

224–230

It could be nice if we can have a more specially names or/and documentation to help understanding of the two usages of "tmp".
the first tmp: Apply the reversed permutation to the storage sizes to get the sizes of the tensor.
the second tmp: Apply the reversed permutation to the permutation of the coo tensor to get the accumulated permutation from the tensor to the coo tensor.

[aartbik]

mlir/lib/ExecutionEngine/SparseUtils.cpp

224–230

Yeah, probably premature optimization keeping the same vector for both.
I changed this back into two vectors with better names, leaving the optimization for another time (not that this one will help much anyway).

[aartbik]

renamed tmp vector, used two different ones, added comments

mlir/lib/ExecutionEngine/SparseUtils.cpp

public:
   const Element<V> &next() { return elements[pos++]; }
   /// 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; }
 
-  /// Factory method.
+  /// Factory method. Permutes the original dimensions according to
+  /// the given ordering and expects subsequent add() calls to honor
+  /// that same ordering for the given indices. The result is a
+  /// fully permuted coordinate scheme.
   static SparseTensor<V> *newSparseTensor(uint64_t size, uint64_t *sizes,
                                           uint64_t *perm,
                                           uint64_t capacity = 0) {
-    std::vector<uint64_t> indices(size);
+    std::vector<uint64_t> permsz(size);
     for (uint64_t r = 0; r < size; r++)
-      indices[perm[r]] = sizes[r];
-    return new SparseTensor<V>(indices, capacity);
+      permsz[perm[r]] = sizes[r];
+    return new SparseTensor<V>(permsz, capacity);
   }
 
 private:
   /// Returns true if indices of e1 < indices of e2.
   static bool lexOrder(const Element<V> &e1, const Element<V> &e2) {
     assert(e1.indices.size() == e2.indices.size());
     for (uint64_t r = 0, rank = e1.indices.size(); r < rank; r++) {
       if (e1.indices[r] == e2.indices[r])

 /// "one-size-fits-all" solution that simply takes an input tensor and
 /// annotations to implement all required setup in a general manner.
 template <typename P, typename I, typename V>
 class SparseTensorStorage : public SparseTensorStorageBase {
 public:
   /// Constructs a sparse tensor storage scheme from the given sparse
   /// tensor in coordinate scheme following the given per-rank dimension
   /// dense/sparse annotations.
-  SparseTensorStorage(SparseTensor<V> *tensor, uint8_t *sparsity)
-      : sizes(tensor->getSizes()), pointers(getRank()), indices(getRank()) {
+  SparseTensorStorage(SparseTensor<V> *tensor, uint8_t *sparsity,
+                      uint64_t *perm)
+      : sizes(tensor->getSizes()), rev(getRank()), pointers(getRank()),
+        indices(getRank()) {
+    // Store "reverse" permutation.
+    for (uint64_t d = 0, rank = getRank(); d < rank; d++)
+      rev[perm[d]] = d;
     // Provide hints on capacity.
     // TODO: needs fine-tuning based on sparsity
     uint64_t nnz = tensor->getElements().size();
     values.reserve(nnz);
     for (uint64_t d = 0, s = 1, rank = getRank(); d < rank; d++) {
       s *= sizes[d];
       if (sparsity[d] == kCompressed) {
         pointers[d].reserve(s + 1);
         indices[d].reserve(s);
         s = 1;
       } else {
         assert(sparsity[d] == kDense && "singleton not yet supported");
       }
     }
+    // Prepare sparse pointer structures for all dimensions.
+    for (uint64_t d = 0, rank = getRank(); d < rank; d++)
+      if (sparsity[d] == kCompressed)
+        pointers[d].push_back(0);
     // Then setup the tensor.
-    traverse(tensor, sparsity, 0, nnz, 0);
+    fromCOO(tensor, sparsity, 0, nnz, 0);
   }
 
   virtual ~SparseTensorStorage() {}
 
   uint64_t getRank() const { return sizes.size(); }
 
   uint64_t getDimSize(uint64_t d) override { return sizes[d]; }
 
   // Partially specialize these three methods based on template types.
   void getPointers(std::vector<P> **out, uint64_t d) override {
     *out = &pointers[d];
   }
   void getIndices(std::vector<I> **out, uint64_t d) override {
     *out = &indices[d];
   }
   void getValues(std::vector<V> **out) override { *out = &values; }
 
-  /// Factory method.
-  static SparseTensorStorage<P, I, V> *newSparseTensor(SparseTensor<V> *t,
-                                                       uint8_t *s) {
+  /// Returns this sparse tensor storage scheme as a new memory-resident
+  /// sparse tensor in coordinate scheme with the given dimension order.
+  SparseTensor<V> *asCOO(uint64_t *perm) {
+    // Restore original order of the dimension sizes and allocate coordinate
+    // scheme with desired new ordering specified in perm.
+    uint64_t size = getRank();
+    std::vector<uint64_t> orgsz(size);
+    for (uint64_t r = 0; r < size; r++)
+      orgsz[rev[r]] = sizes[r];
+    SparseTensor<V> *tensor = SparseTensor<V>::newSparseTensor(
+        size, orgsz.data(), perm, values.size());
+    // Populate coordinate scheme restored from old ordering and changed with
+    // new ordering. Rather than applying both reorderings during the recursion,
+    // we compute the combine permutation in advance.
+    std::vector<uint64_t> reord(size);
+    for (uint64_t r = 0; r < size; r++)
+      reord[r] = perm[rev[r]];
+    std::vector<uint64_t> idx(size);
+    toCOO(tensor, reord, idx, 0, 0);
+    return tensor;
+  }
+
+  /// Factory method. Expects a coordinate scheme that respects the same
+  /// permutation as is desired for the new sparse storage scheme.
+  static SparseTensorStorage<P, I, V> *
+  newSparseTensor(SparseTensor<V> *t, uint8_t *sparsity, uint64_t *perm) {
     t->sort(); // sort lexicographically
-    SparseTensorStorage<P, I, V> *n = new SparseTensorStorage<P, I, V>(t, s);
+    SparseTensorStorage<P, I, V> *n =
+        new SparseTensorStorage<P, I, V>(t, sparsity, perm);
     delete t;
     return n;
   }
 
 private:
   /// 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-rank dimension dense/sparse annotations.
-  void traverse(SparseTensor<V> *tensor, uint8_t *sparsity, uint64_t lo,
-                uint64_t hi, uint64_t d) {
+  void fromCOO(SparseTensor<V> *tensor, uint8_t *sparsity, uint64_t lo,
     const std::vector<Element<V>> &elements = tensor->getElements();
     // Once dimensions are exhausted, insert the numerical values.
     if (d == getRank()) {
       values.push_back(lo < hi ? elements[lo].value : 0);
       return;
     }
-    // Prepare a sparse pointer structure at this dimension.
-    if (sparsity[d] == kCompressed && pointers[d].empty())
-      pointers[d].push_back(0);
     // Visit all elements in this interval.
     uint64_t full = 0;
     while (lo < hi) {
       // Find segment in interval with same index elements in this dimension.
       unsigned idx = elements[lo].indices[d];
       unsigned seg = lo + 1;
       while (seg < hi && elements[seg].indices[d] == idx)
         seg++;
       // Handle segment in interval for sparse or dense dimension.
       if (sparsity[d] == kCompressed) {
         indices[d].push_back(idx);
       } else {
         for (; full < idx; full++)
-          traverse(tensor, sparsity, 0, 0, d + 1); // pass empty
+          fromCOO(tensor, sparsity, 0, 0, d + 1); // pass empty
         full++;
       }
-      traverse(tensor, sparsity, lo, seg, d + 1);
+      fromCOO(tensor, sparsity, lo, seg, d + 1);
       // And move on to next segment in interval.
       lo = seg;
     }
     // Finalize the sparse pointer structure at this dimension.
     if (sparsity[d] == kCompressed) {
       pointers[d].push_back(indices[d].size());
     } else {
       for (uint64_t sz = tensor->getSizes()[d]; full < sz; full++)
-        traverse(tensor, sparsity, 0, 0, d + 1); // pass empty
+        fromCOO(tensor, sparsity, 0, 0, d + 1); // pass empty
+    }
+  }
+
+  /// Stores the sparse tensor storage scheme into a memory-resident sparse
+  /// tensor in coordinate scheme.
+  void toCOO(SparseTensor<V> *tensor, std::vector<uint64_t> &reord,
+             std::vector<uint64_t> &idx, uint64_t pos, uint64_t d) {
+    if (d == getRank()) {
+      tensor->add(idx, values[pos]);
+    } else if (pointers[d].empty()) {
+      // Dense dimension.
+      for (uint64_t i = 0; i < sizes[d]; i++) {
+        idx[reord[d]] = i;
+        toCOO(tensor, reord, idx, pos * sizes[d] + i, d + 1);
+      }
+    } else {
+      // Sparse dimension.
+      for (uint64_t ii = pointers[d][pos]; ii < pointers[d][pos + 1]; ii++) {
+        idx[reord[d]] = indices[d][ii];
+        toCOO(tensor, reord, idx, ii, d + 1);
+      }
     }
   }
 
 private:
   std::vector<uint64_t> sizes; // per-rank dimension sizes
+  std::vector<uint64_t> rev;   // "reverse" permutation
   std::vector<std::vector<P>> pointers;
   std::vector<std::vector<I>> indices;
   std::vector<V> values;
 };
 
 /// Helper to convert string to lower case.
 static char *toLower(char *token) {
   for (char *c = token; *c; c++)

 
 #define CASE(p, i, v, P, I, V)                                                 \
   if (ptrTp == (p) && indTp == (i) && valTp == (v)) {                          \
     SparseTensor<V> *tensor;                                                   \
     if (action == 0)                                                           \
       tensor = openTensor<V>(static_cast<char *>(ptr), asize, sizes, perm);    \
     else if (action == 1)                                                      \
       tensor = static_cast<SparseTensor<V> *>(ptr);                            \
-    else                                                                       \
+    else if (action == 2)                                                      \
       return SparseTensor<V>::newSparseTensor(asize, sizes, perm);             \
-    return SparseTensorStorage<P, I, V>::newSparseTensor(tensor, sparsity);    \
+    else                                                                       \
+      return static_cast<SparseTensorStorage<P, I, V> *>(ptr)->asCOO(perm);    \
+    return SparseTensorStorage<P, I, V>::newSparseTensor(tensor, sparsity,     \
+                                                         perm);                \
   }
 
 #define IMPL1(RET, NAME, TYPE, LIB)                                            \
   RET NAME(void *tensor) {                                                     \
     std::vector<TYPE> *v;                                                      \
     static_cast<SparseTensorStorageBase *>(tensor)->LIB(&v);                   \
     return {v->data(), v->data(), 0, {v->size()}, {1}};                        \
   }
enum PrimaryTypeEnum : uint64_t {
   kI32 = 4,
   kI16 = 5,
   kI8 = 6
 };
 
 /// Constructs a new sparse tensor. This is the "swiss army knife"
 /// method for materializing sparse tensors into the computation.
 ///  action
-///    0 : ptr contains filename to read into storage
-///    1 : ptr contains coordinate scheme to assign to storage
-///    2 : returns coordinate scheme to fill (call back later with 1)
+///  0 : ptr contains filename to read into storage
+///  1 : ptr contains coordinate scheme to assign to new storage
+///  2 : returns empty coordinate scheme to fill (call back 1 to setup)
+///  3 : returns coordinate scheme from storage in ptr (call back 1 to convert)
 void *newSparseTensor(uint8_t *abase, uint8_t *adata, uint64_t aoff,
                       uint64_t asize, uint64_t astride, uint64_t *sbase,
                       uint64_t *sdata, uint64_t soff, uint64_t ssize,
                       uint64_t sstride, uint64_t *pbase, uint64_t *pdata,
                       uint64_t poff, uint64_t psize, uint64_t pstride,
                       uint64_t ptrTp, uint64_t indTp, uint64_t valTp,
                       uint32_t action, void *ptr) {
   assert(astride == 1 && sstride == 1 && pstride == 1);