[mlir][sparse] sparse runtime support library improvements

[aartbik]

Summary

(1) renamed SparseTensor to SparseTensorCOO, the other one remains SparseTensorStorage to focus on contrast

(2) documents difference between public API exclusively for compiler-generated code and methods that could be used by other runtimes (TBD) that want to interact with MLIR

[aartbik]

rebased with main

[aartbik]

rebase main

[aartbik]

rebase again (I keep getting patch failures)

mlir/lib/ExecutionEngine/SparseUtils.cpp

 
 //===----------------------------------------------------------------------===//
 //
 // Internal support for storing and reading sparse tensors.
 //
 // The following memory-resident sparse storage schemes are supported:
 //
 // (a) A coordinate scheme for temporarily storing and lexicographically
-//     sorting a sparse tensor by index.
+//     sorting a sparse tensor by index (SparseTensorCOO).
 //
 // (b) A "one-size-fits-all" sparse tensor storage scheme defined by per-rank
-//     sparse/dense annnotations to be used by generated MLIR code.
+//     sparse/dense annnotations together with a dimension ordering to be
+//     used by MLIR compiler-generated code (SparseTensorStorage).
 //
 // The following external formats are supported:
 //
 // (1) Matrix Market Exchange (MME): *.mtx
 //     https://math.nist.gov/MatrixMarket/formats.html
 //
 // (2) Formidable Repository of Open Sparse Tensors and Tools (FROSTT): *.tns
 //     http://frostt.io/tensors/file-formats.html
 //
+// Two public APIs are supported:
+//
+// (I) Methods operating on MLIR buffers (memrefs) to interact with sparse
+//     tensors. These methods should be used exclusively by MLIR
+//     compiler-generated code.
+//
+// (II) Methods that accept C-style data structures to interact with sparse
+//      tensors. These methods can be used by any external runtime that wants
+//      to interact with MLIR compiler-generated code.
+//
+// In both cases (I) and (II), the SparseTensorStorage format is externally
+// only visible as an opague pointer.
+//
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 /// A sparse tensor element in coordinate scheme (value and indices).
 /// For example, a rank-1 vector element would look like
 ///   ({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;
 };
 
 /// 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>
-struct SparseTensor {
+struct SparseTensorCOO {
 public:
-  SparseTensor(const std::vector<uint64_t> &szs, uint64_t capacity)
+  SparseTensorCOO(const std::vector<uint64_t> &szs, uint64_t capacity)
       : sizes(szs), pos(0) {
     if (capacity)
       elements.reserve(capacity);
   }
   /// Adds element as indices and value.
   void add(const std::vector<uint64_t> &ind, V val) {
     assert(getRank() == ind.size());
     for (uint64_t r = 0, rank = getRank(); r < rank; r++)
public:
   const std::vector<uint64_t> &getSizes() const { return sizes; }
   /// Getter for elements array.
   const std::vector<Element<V>> &getElements() const { return elements; }
 
   /// 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,
+  static SparseTensorCOO<V> *newSparseTensorCOO(uint64_t size, uint64_t *sizes,
                                                 uint64_t *perm,
                                                 uint64_t capacity = 0) {
     std::vector<uint64_t> permsz(size);
     for (uint64_t r = 0; r < size; r++)
       permsz[perm[r]] = sizes[r];
-    return new SparseTensor<V>(permsz, capacity);
+    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) {
     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])
private:
   void fatal(const char *tp) {
     fprintf(stderr, "unsupported %s\n", tp);
     exit(1);
   }
 };
 
 /// A memory-resident sparse tensor using a storage scheme based on per-rank
 /// annotations on dense/sparse. This data structure provides a bufferized
-/// form of an imaginary SparseTensorType, until such a type becomes a
-/// first-class citizen of MLIR. In contrast to generating setup methods for
+/// form of a sparse tensor type. In contrast to generating setup methods for
 /// each differently annotated sparse tensor, this method provides a convenient
 /// "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,
+  SparseTensorStorage(SparseTensorCOO<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
public:
   }
   void getIndices(std::vector<I> **out, uint64_t d) override {
     *out = &indices[d];
   }
   void getValues(std::vector<V> **out) override { *out = &values; }
 
   /// Returns this sparse tensor storage scheme as a new memory-resident
   /// sparse tensor in coordinate scheme with the given dimension order.
-  SparseTensor<V> *toCOO(uint64_t *perm) {
+  SparseTensorCOO<V> *toCOO(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(
+    SparseTensorCOO<V> *tensor = SparseTensorCOO<V>::newSparseTensorCOO(
         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) {
+  newSparseTensor(SparseTensorCOO<V> *t, uint8_t *sparsity, uint64_t *perm) {
     t->sort(); // sort lexicographically
     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 fromCOO(SparseTensor<V> *tensor, uint8_t *sparsity, uint64_t lo,
+  void fromCOO(SparseTensorCOO<V> *tensor, uint8_t *sparsity, uint64_t lo,
                uint64_t hi, uint64_t d) {
     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;
     }
     // Visit all elements in this interval.
if (sparsity[d] == kCompressed) {
       // the last element.
       for (uint64_t sz = tensor->getSizes()[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(SparseTensor<V> *tensor, std::vector<uint64_t> &reord,
+  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++) {
         idx[reord[d]] = i;
         toCOO(tensor, reord, idx, pos * sizes[d] + i, d + 1);
if (fscanf(file, "%" PRIu64, idata + 2 + r) != 1) {
       exit(1);
     }
   }
 }
 
 /// Reads a sparse tensor with the given filename into a memory-resident
 /// sparse tensor in coordinate scheme.
 template <typename V>
-static SparseTensor<V> *openTensor(char *filename, uint64_t size,
-                                   uint64_t *sizes, uint64_t *perm) {
+static SparseTensorCOO<V> *openSparseTensorCOO(char *filename, uint64_t size,
+                                               uint64_t *sizes,
+                                               uint64_t *perm) {
   // Open the file.
   FILE *file = fopen(filename, "r");
   if (!file) {
     fprintf(stderr, "Cannot find %s\n", filename);
     exit(1);
   }
   // Perform some file format dependent set up.
   uint64_t idata[512];
   if (strstr(filename, ".mtx")) {
     readMMEHeader(file, filename, idata);
   } else if (strstr(filename, ".tns")) {
     readExtFROSTTHeader(file, filename, idata);
   } else {
     fprintf(stderr, "Unknown format %s\n", filename);
     exit(1);
   }
   // Prepare sparse tensor object with per-rank dimension sizes
   // and the number of nonzeros as initial capacity.
   assert(size == idata[0] && "rank mismatch");
   uint64_t nnz = idata[1];
   for (uint64_t r = 0; r < size; r++)
     assert((sizes[r] == 0 || sizes[r] == idata[2 + r]) &&
            "dimension size mismatch");
-  SparseTensor<V> *tensor =
-      SparseTensor<V>::newSparseTensor(size, idata + 2, perm, nnz);
+  SparseTensorCOO<V> *tensor =
+      SparseTensorCOO<V>::newSparseTensorCOO(size, idata + 2, perm, nnz);
   //  Read all nonzero elements.
   std::vector<uint64_t> indices(size);
   for (uint64_t k = 0; k < nnz; k++) {
     uint64_t idx = -1;
     for (uint64_t r = 0; r < size; r++) {
       if (fscanf(file, "%" PRIu64, &idx) != 1) {
         fprintf(stderr, "Cannot find next index in %s\n", filename);
         exit(1);
char *getTensorFilename(uint64_t id) {
   char var[80];
   sprintf(var, "TENSOR%" PRIu64, id);
   char *env = getenv(var);
   return env;
 }
 
 //===----------------------------------------------------------------------===//
 //
-// Public API of the sparse runtime support library that support an opaque
-// implementation of a bufferized SparseTensor in MLIR. This could be replaced
-// by actual codegen in MLIR.
+// Public API with methods that operate on MLIR buffers (memrefs) to interact
+// with sparse tensors, which are only visible as opaque pointers externally.
+// These methods should be used exclusively by MLIR compiler-generated code.
 //
 // Because we cannot use C++ templates with C linkage, some macro magic is used
 // to generate implementations for all required type combinations that can be
-// called from MLIR generated code.
+// called from MLIR compiler-generated code.
 //
 //===----------------------------------------------------------------------===//
 
 #define TEMPLATE(NAME, TYPE)                                                   \
   struct NAME {                                                                \
     const TYPE *base;                                                          \
     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)) {                          \
-    SparseTensor<V> *tensor;                                                   \
+    SparseTensorCOO<V> *tensor;                                                \
     if (action == 0)                                                           \
-      tensor = openTensor<V>(static_cast<char *>(ptr), asize, sizes, perm);    \
+      tensor = openSparseTensorCOO<V>(static_cast<char *>(ptr), asize, sizes,  \
+                                      perm);                                   \
     else if (action == 1)                                                      \
-      tensor = static_cast<SparseTensor<V> *>(ptr);                            \
+      tensor = static_cast<SparseTensorCOO<V> *>(ptr);                         \
     else if (action == 2)                                                      \
-      return SparseTensor<V>::newSparseTensor(asize, sizes, perm);             \
+      return SparseTensorCOO<V>::newSparseTensorCOO(asize, 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 *NAME(void *tensor, TYPE value, uint64_t *ibase, uint64_t *idata,       \
     assert(istride == 1 && pstride == 1 && isize == psize);                    \
     uint64_t *indx = idata + ioff;                                             \
     if (!value)                                                                \
       return tensor;                                                           \
     uint64_t *perm = pdata + poff;                                             \
     std::vector<uint64_t> indices(isize);                                      \
     for (uint64_t r = 0; r < isize; r++)                                       \
       indices[perm[r]] = indx[r];                                              \
-    static_cast<SparseTensor<TYPE> *>(tensor)->add(indices, value);            \
+    static_cast<SparseTensorCOO<TYPE> *>(tensor)->add(indices, value);         \
     return tensor;                                                             \
   }
 
 TEMPLATE(MemRef1DU64, uint64_t);
 TEMPLATE(MemRef1DU32, uint32_t);
 TEMPLATE(MemRef1DU16, uint16_t);
 TEMPLATE(MemRef1DU8, uint8_t);
 TEMPLATE(MemRef1DI64, int64_t);
void *newSparseTensor(uint8_t *abase, uint8_t *adata, uint64_t aoff,
   CASE(kU8, kU8, kI16, uint8_t, uint8_t, int16_t);
   CASE(kU8, kU8, kI8, uint8_t, uint8_t, int8_t);
 
   // Unsupported case (add above if needed).
   fputs("unsupported combination of types\n", stderr);
   exit(1);
 }
 
-/// Returns size of sparse tensor in given dimension.
-uint64_t sparseDimSize(void *tensor, uint64_t d) {
-  return static_cast<SparseTensorStorageBase *>(tensor)->getDimSize(d);
-}
-
 /// Methods that provide direct access to pointers, indices, and values.
 IMPL2(MemRef1DU64, sparsePointers, uint64_t, getPointers)
 IMPL2(MemRef1DU64, sparsePointers64, uint64_t, getPointers)
 IMPL2(MemRef1DU32, sparsePointers32, uint32_t, getPointers)
 IMPL2(MemRef1DU16, sparsePointers16, uint16_t, getPointers)
 IMPL2(MemRef1DU8, sparsePointers8, uint8_t, getPointers)
 IMPL2(MemRef1DU64, sparseIndices, uint64_t, getIndices)
 IMPL2(MemRef1DU64, sparseIndices64, uint64_t, getIndices)
 IMPL2(MemRef1DU32, sparseIndices32, uint32_t, getIndices)
 IMPL2(MemRef1DU16, sparseIndices16, uint16_t, getIndices)
 IMPL2(MemRef1DU8, sparseIndices8, uint8_t, getIndices)
 IMPL1(MemRef1DF64, sparseValuesF64, double, getValues)
 IMPL1(MemRef1DF32, sparseValuesF32, float, getValues)
 IMPL1(MemRef1DI64, sparseValuesI64, int64_t, getValues)
 IMPL1(MemRef1DI32, sparseValuesI32, int32_t, getValues)
 IMPL1(MemRef1DI16, sparseValuesI16, int16_t, getValues)
 IMPL1(MemRef1DI8, sparseValuesI8, int8_t, getValues)
 
-/// Releases sparse tensor storage.
-void delSparseTensor(void *tensor) {
-  delete static_cast<SparseTensorStorageBase *>(tensor);
-}
-
 /// Helper to add value to coordinate scheme, one per value type.
 IMPL3(addEltF64, double)
 IMPL3(addEltF32, float)
 IMPL3(addEltI64, int64_t)
 IMPL3(addEltI32, int32_t)
 IMPL3(addEltI16, int16_t)
 IMPL3(addEltI8, int8_t)
 
 #undef TEMPLATE
 #undef CASE
 #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.
+//
+//===----------------------------------------------------------------------===//
+
+// 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);
+}
+
 } // extern "C"
 
 #endif // MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS