[mlir][sparse] add convenience method for sparse tensor setup

[aartbik]

Summary

This simplifies setting up sparse tensors through C-style data structures.
Useful for runtimes that want to interact with MLIR-generated code
without knowning about all bufferization details (viz. memrefs).

[bixia]

mlir/lib/ExecutionEngine/SparseUtils.cpp

729–734

NIT: this is the same as:
std::vector<uint8_t> sparse(rank, SparseTensorStorageBase::kCompressed)
std::vector<unit64_t> perm(rank);
std::iota(perm.begin(), perm.end(), 0);

[aartbik]

mlir/lib/ExecutionEngine/SparseUtils.cpp

729–734

Clever!

[aartbik]

use Bixia's clever C++ stuff

mlir/lib/ExecutionEngine/SparseUtils.cpp

 
 #include <algorithm>
 #include <cassert>
 #include <cctype>
 #include <cinttypes>
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
+#include <numeric>
 #include <vector>
 
 //===----------------------------------------------------------------------===//
 //
 // Internal support for storing and reading sparse tensors.
 //
 // The following memory-resident sparse storage schemes are supported:
 //
while (lo < hi) {
       lo = seg;
     }
     // Finalize the sparse pointer structure at this dimension.
     if (sparsity[d] == kCompressed) {
       pointers[d].push_back(indices[d].size());
     } else {
       // For dense storage we must fill in all the zero values after
       // the last element.
-      for (uint64_t sz = tensor->getSizes()[d]; full < sz; full++)
+      for (uint64_t sz = sizes[d]; full < sz; full++)
         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(SparseTensorCOO<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++) {
+      for (uint64_t i = 0, sz = sizes[d], off = pos * sz; i < sz; i++) {
         idx[reord[d]] = i;
-        toCOO(tensor, reord, idx, pos * sizes[d] + i, d + 1);
+        toCOO(tensor, reord, idx, off + 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);
       }
     }
struct NAME {                                                                \
     const TYPE *data;                                                          \
     uint64_t off;                                                              \
     uint64_t sizes[1];                                                         \
     uint64_t strides[1];                                                       \
   }
 
 #define CASE(p, i, v, P, I, V)                                                 \
   if (ptrTp == (p) && indTp == (i) && valTp == (v)) {                          \
-    SparseTensorCOO<V> *tensor;                                                \
+    SparseTensorCOO<V> *tensor = nullptr;                                      \
     if (action == 0)                                                           \
-      tensor = openSparseTensorCOO<V>(static_cast<char *>(ptr), asize, sizes,  \
-                                      perm);                                   \
+      tensor =                                                                 \
+          openSparseTensorCOO<V>(static_cast<char *>(ptr), size, sizes, perm); \
     else if (action == 1)                                                      \
       tensor = static_cast<SparseTensorCOO<V> *>(ptr);                         \
     else if (action == 2)                                                      \
-      return SparseTensorCOO<V>::newSparseTensorCOO(asize, sizes, perm);       \
+      return SparseTensorCOO<V>::newSparseTensorCOO(size, sizes, perm);        \
     else                                                                       \
       return static_cast<SparseTensorStorage<P, I, V> *>(ptr)->toCOO(perm);    \
     return SparseTensorStorage<P, I, V>::newSparseTensor(tensor, sparsity,     \
                                                          perm);                \
   }
 
 #define IMPL1(RET, NAME, TYPE, LIB)                                            \
   RET NAME(void *tensor) {                                                     \
void *newSparseTensor(uint8_t *abase, uint8_t *adata, uint64_t aoff,
                       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);
   assert(asize == ssize && ssize == psize);
   uint8_t *sparsity = adata + aoff;
   uint64_t *sizes = sdata + soff;
   uint64_t *perm = pdata + poff;
+  uint64_t size = asize;
 
   // Double matrices with all combinations of overhead storage.
   CASE(kU64, kU64, kF64, uint64_t, uint64_t, double);
   CASE(kU64, kU32, kF64, uint64_t, uint32_t, double);
   CASE(kU64, kU16, kF64, uint64_t, uint16_t, double);
   CASE(kU64, kU8, kF64, uint64_t, uint8_t, double);
   CASE(kU32, kU64, kF64, uint32_t, uint64_t, double);
   CASE(kU32, kU32, kF64, uint32_t, uint32_t, double);

 #undef IMPL1
 #undef IMPL2
 #undef IMPL3
 
 //===----------------------------------------------------------------------===//
 //
 // Public API with methods that accept C-style data structures to interact
 // with sparse tensors, which are only visible as opaque pointers externally.
-// These methods can be used by any external runtime that wants to interact
-// with MLIR compiler-generated code.
+// These methods can be used both by MLIR compiler-generated code as well as by
+// an external runtime that wants to interact with MLIR compiler-generated code.
 //
 //===----------------------------------------------------------------------===//
 
-// TODO: setup sparse tensor from C style format
-
 /// Returns size of sparse tensor in given dimension.
 uint64_t sparseDimSize(void *tensor, uint64_t d) {
   return static_cast<SparseTensorStorageBase *>(tensor)->getDimSize(d);
 }
 
 /// Releases sparse tensor storage.
 void delSparseTensor(void *tensor) {
   delete static_cast<SparseTensorStorageBase *>(tensor);
 }
 
+/// Initializes sparse tensor from a COO-flavored format expressed using C-style
+/// data structures. The expected parameters are:
+///
+///   rank:    rank of tensor
+///   nse:     number of specified elements (usually the nonzeros)
+///   shape:   array with dimension size for each rank
+///   values:  a "nse" array with values for all specified elements
+///   indices: a flat "nse x rank" array with indices for all specified elements
+///
+/// For example, the sparse matrix
+///     | 1.0 0.0 0.0 |
+///     | 0.0 5.0 3.0 |
+/// can be passed as
+///      rank    = 2
+///      nse     = 3
+///      shape   = [2, 3]
+///      values  = [1.0, 5.0, 3.0]
+///      indices = [ 0, 0,  1, 1,  1, 2]
+//
+// TODO: for now f64 tensors only, no dim ordering, all dimensions compressed
+//
+void *convertToMLIRSparseTensor(uint64_t rank, uint64_t nse, uint64_t *shape,
+                                double *values, uint64_t *indices) {
+  // Setup all-dims compressed and default ordering.
+  std::vector<uint8_t> sparse(rank, SparseTensorStorageBase::kCompressed);
+  std::vector<uint64_t> perm(rank);
+  std::iota(perm.begin(), perm.end(), 0);
+  // Convert external format to internal COO.
+  SparseTensorCOO<double> *tensor = SparseTensorCOO<double>::newSparseTensorCOO(
+      rank, shape, perm.data(), nse);
+  std::vector<uint64_t> idx(rank);
+  for (uint64_t i = 0, base = 0; i < nse; i++) {
+    for (uint64_t j = 0; j < rank; j++)
+      idx[j] = indices[base + j];
+    tensor->add(idx, values[i]);
+    base += rank;
+  }
+  // Return sparse tensor storage format as opague pointer.
+  return SparseTensorStorage<uint64_t, uint64_t, double>::newSparseTensor(
+      tensor, sparse.data(), perm.data());
+}
+
 } // extern "C"
 
 #endif // MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS