[mlir][sparse] Avoiding intermediate arrays in convertFromMLIRSparseTensor

[wrengr]

Summary

Depends On D117169

[aartbik]

Here too we get a more code for a very small gain. What bottleneck are we solving here? Do we have concrete performance data to justify the changes?

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

for (uint64_t r = 0; r < rank; r++) {
       permsz[s] = sizes[r];
     }
     SparseTensorCOO<V> *tensor = new SparseTensorCOO<V>(permsz, values.size());
     // Populate the coordinate scheme.
     toCOO(*tensor, reord, 0, 0);
     assert(tensor->getElements().size() == values.size());
     return tensor;
   }
+  /// Variant for when the permutation is known to be the identity/iota.
+  SparseTensorCOO<V> *toCOO() {
+    uint64_t rank = getRank();
+    // The new coordinate scheme's sizes array.
+    std::vector<uint64_t> permsz(rank);
+    for (uint64_t r = 0; r < rank; r++)
+      permsz[rev[r]] = sizes[r];
+    SparseTensorCOO<V> *tensor = new SparseTensorCOO<V>(permsz, values.size());
+    // Populate the coordinate scheme.
+    toCOO(*tensor, rev, 0, 0);
+    assert(tensor->getElements().size() == values.size());
+    return tensor;
+  }
 
   /// Factory method. Constructs a sparse tensor storage scheme with the given
   /// dimensions, permutation, and per-dimension dense/sparse annotations,
   /// using the coordinate scheme tensor for the initial contents if provided.
   /// In the latter case, the coordinate scheme must respect the same
   /// permutation as is desired for the new sparse tensor storage.
   static SparseTensorStorage<P, I, V> *
   newSparseTensor(uint64_t rank, const uint64_t *sizes, const uint64_t *perm,

 //  TODO: for now f64 tensors only, no dim ordering, all dimensions compressed
 //
 void convertFromMLIRSparseTensor(void *tensor, uint64_t *pRank, uint64_t *pNse,
                                  uint64_t **pShape, double **pValues,
                                  uint64_t **pIndices) {
   SparseTensorStorage<uint64_t, uint64_t, double> *sparseTensor =
       static_cast<SparseTensorStorage<uint64_t, uint64_t, double> *>(tensor);
   uint64_t rank = sparseTensor->getRank();
-  std::vector<uint64_t> perm(rank);
-  std::iota(perm.begin(), perm.end(), 0);
-  SparseTensorCOO<double> *coo = sparseTensor->toCOO(perm.data());
+  SparseTensorCOO<double> *coo = sparseTensor->toCOO();
 
   const std::vector<Element<double>> &elements = coo->getElements();
   uint64_t nse = elements.size();
 
   uint64_t *shape = new uint64_t[rank];
   for (uint64_t i = 0; i < rank; i++)
     shape[i] = coo->getSizes()[i];