[mlir][sparse] template the memory resident coordinate scheme storage

[aartbik]

Summary

Rationale:
External file formats always store the values as doubles, so this was
hard coded in the memory resident COO scheme used to pass data into the
final sparse storage scheme during setup. However, with alternative methods
on the horizon of setting up these temporary COO schemes, it is time to
properly template this data structure.

[gussmith23]

I left a comment re: naming, but everything looks good otherwise!

mlir/lib/ExecutionEngine/SparseUtils.cpp

154

Hm, the fact that the naming is inconsistent here (that is, this typename is V in both SparseTensor and Element) is potentially confusing. Maybe you could use V instead of W here, and rename the other V? Alternatively, you could rename the V used in SparseTensor and Element?

[aartbik]

mlir/lib/ExecutionEngine/SparseUtils.cpp

154

Interesting that you bring this up, since this was bothering me too. I like the consistent V naming (since it is really the value type parameter) and the friction here is purely due to external formats always using double (normally you would simply pass the single V to "both" for a consistent naming). Renaming either would creep into all other places.

How about using Vm and Ve (for the V for memory and extermal)? Or V1 and V2?

[aartbik]

renamed W (the "second" V) to Ve

mlir/lib/ExecutionEngine/SparseUtils.cpp

 
 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, double val)
-      : indices(ind), value(val){};
+  Element(const std::vector<uint64_t> &ind, V val) : indices(ind), value(val){};
   std::vector<uint64_t> indices;
-  double value;
+  V value;
 };
 
 /// A memory-resident sparse tensor in coordinate scheme (collection of
 /// elements). This data structure is used to read a sparse tensor from
-/// external file format into memory and sort the elements lexicographically
+/// 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 {
 public:
   SparseTensor(const std::vector<uint64_t> &szs, uint64_t capacity)
       : sizes(szs), pos(0) {
     elements.reserve(capacity);
   }
   /// Adds element as indices and value.
-  void add(const std::vector<uint64_t> &ind, double val) {
+  void add(const std::vector<uint64_t> &ind, V val) {
     assert(getRank() == ind.size());
     for (int64_t r = 0, rank = getRank(); r < rank; r++)
       assert(ind[r] < sizes[r]); // within bounds
-    elements.emplace_back(Element(ind, val));
+    elements.emplace_back(Element<V>(ind, val));
   }
   /// Sorts elements lexicographically by index.
   void sort() { std::sort(elements.begin(), elements.end(), lexOrder); }
   /// Primitive one-time iteration.
-  const Element &next() { return elements[pos++]; }
+  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> &getElements() const { return elements; }
+  const std::vector<Element<V>> &getElements() const { return elements; }
 
 private:
   /// Returns true if indices of e1 < indices of e2.
-  static bool lexOrder(const Element &e1, const Element &e2) {
+  static bool lexOrder(const Element<V> &e1, const Element<V> &e2) {
     assert(e1.indices.size() == e2.indices.size());
     for (int64_t r = 0, rank = e1.indices.size(); r < rank; r++) {
       if (e1.indices[r] == e2.indices[r])
         continue;
       return e1.indices[r] < e2.indices[r];
     }
     return false;
   }
   std::vector<uint64_t> sizes; // per-rank dimension sizes
-  std::vector<Element> elements;
+  std::vector<Element<V>> elements;
   uint64_t pos;
 };
 
 /// Abstract base class of sparse tensor storage. Note that we use
 /// function overloading to implement "partial" method specialization.
 class SparseTensorStorageBase {
 public:
   enum DimLevelType : uint8_t { kDense = 0, kCompressed = 1, kSingleton = 2 };

 
 /// 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
 /// 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>
+template <typename P, typename I, typename V, typename Ve>
 class SparseTensorStorage : public SparseTensorStorageBase {
 public:
   /// Constructs sparse tensor storage scheme following the given
   /// per-rank dimension dense/sparse annotations.
-  SparseTensorStorage(SparseTensor *tensor, uint8_t *sparsity)
+  SparseTensorStorage(SparseTensor<Ve> *tensor, uint8_t *sparsity)
       : sizes(tensor->getSizes()), pointers(getRank()), indices(getRank()) {
     // 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) {
public:
   }
   void getValues(std::vector<V> **out) override { *out = &values; }
 
 private:
   /// Initializes sparse tensor storage scheme from a memory-resident
   /// representation of an external sparse tensor. This method prepares
   /// the pointers and indices arrays under the given per-rank dimension
   /// dense/sparse annotations.
-  void traverse(SparseTensor *tensor, uint8_t *sparsity, uint64_t lo,
+  void traverse(SparseTensor<Ve> *tensor, uint8_t *sparsity, uint64_t lo,
                 uint64_t hi, uint64_t d) {
-    const std::vector<Element> &elements = tensor->getElements();
+    const std::vector<Element<Ve>> &elements = tensor->getElements();
     // Once dimensions are exhausted, insert the numerical values.
     if (d == getRank()) {
-      values.push_back(lo < hi ? elements[lo].value : 0.0);
+      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) {
for (uint64_t r = 0; r < idata[0]; r++) {
     if (fscanf(file, "%" PRIu64, idata + 2 + r) != 1) {
       fprintf(stderr, "Cannot find dimension size %s\n", name);
       exit(1);
     }
   }
 }
 
 /// Reads a sparse tensor with the given filename into a memory-resident
-/// sparse tensor in coordinate scheme.
-static SparseTensor *openTensor(char *filename, uint64_t *perm) {
+/// sparse tensor in coordinate scheme. The external formats always store
+/// the numerical values with the type double.
+static SparseTensor<double> *openTensor(char *filename, 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.
   uint64_t rank = idata[0];
   uint64_t nnz = idata[1];
   std::vector<uint64_t> indices(rank);
   for (uint64_t r = 0; r < rank; r++)
     indices[perm[r]] = idata[2 + r];
-  SparseTensor *tensor = new SparseTensor(indices, nnz);
+  SparseTensor<double> *tensor = new SparseTensor<double>(indices, nnz);
   // Read all nonzero elements.
   for (uint64_t k = 0; k < nnz; k++) {
     uint64_t idx = -1;
     for (uint64_t r = 0; r < rank; r++) {
       if (fscanf(file, "%" PRIu64, &idx) != 1) {
         fprintf(stderr, "Cannot find next index in %s\n", filename);
         exit(1);
       }
static SparseTensor<double> *openTensor(char *filename, uint64_t *perm) {
   tensor->sort(); // sort lexicographically
   return tensor;
 }
 
 /// Templated reader.
 template <typename P, typename I, typename V>
 void *newSparseTensor(char *filename, uint8_t *sparsity, uint64_t *perm,
                       uint64_t size) {
-  SparseTensor *t = openTensor(filename, perm);
+  SparseTensor<double> *t = openTensor(filename, perm);
   assert(size == t->getRank()); // sparsity array must match rank
   SparseTensorStorageBase *tensor =
-      new SparseTensorStorage<P, I, V>(t, sparsity);
+      new SparseTensorStorage<P, I, V, double>(t, sparsity);
   delete t;
   return tensor;
 }
 
 } // anonymous namespace
 
 extern "C" {
 
void *newSparseTensor(char *filename, uint8_t *abase, uint8_t *adata,
   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);
 }
 
-#undef CASE
-
 uint64_t sparseDimSize(void *tensor, uint64_t d) {
   return static_cast<SparseTensorStorageBase *>(tensor)->getDimSize(d);
 }
 
 IMPL2(MemRef1DU64, sparsePointers, uint64_t, getPointers)
 IMPL2(MemRef1DU64, sparsePointers64, uint64_t, getPointers)
 IMPL2(MemRef1DU32, sparsePointers32, uint32_t, getPointers)
 IMPL2(MemRef1DU16, sparsePointers16, uint16_t, getPointers)