[mlir][sparse] speed up sparse tensor file I/O by more than 2x

[aartbik]

Summary

data point using the 3-dim tensor nell-2.tns

MLIR:
READ FILE INTO COO: 24424.369294 ms ---> improves to ----> 9638.501044 ms
SORT COO BEFORE PACK: 762.834831 ms
PACK COO TO TENSOR: 1243.376245 ms

TACO:
b file read: 13270.9 ms
b pack: 7137.74 ms
b size: (12092 x 9184 x 28818), 925300328 bytes

https://github.com/llvm/llvm-project/issues/52679

[aartbik]

rebased main

[bixia]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

576

Shall we define a constant for 1025 and use it everywhere?

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

576

Yes, good suggestion.

[aartbik]

symbolic constant

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

void getPointers(std::vector<P> **out, uint64_t d) override {
     *out = &pointers[d];
   }
   void getIndices(std::vector<I> **out, uint64_t d) override {
     assert(d < getRank());
     *out = &indices[d];
   }
   void getValues(std::vector<V> **out) override { *out = &values; }
 
-  /// Partially specialize lexicographic insertions based on template types.
+  /// Partially specialize lexicographical insertions based on template types.
   void lexInsert(uint64_t *cursor, V val) override {
     // First, wrap up pending insertion path.
     uint64_t diff = 0;
     uint64_t top = 0;
     if (!values.empty()) {
       diff = lexDiff(cursor);
       endPath(diff + 1);
       top = idx[diff] + 1;

 /// Helper to convert string to lower case.
 static char *toLower(char *token) {
   for (char *c = token; *c; c++)
     *c = tolower(*c);
   return token;
 }
 
 /// Read the MME header of a general sparse matrix of type real.
-static void readMMEHeader(FILE *file, char *name, uint64_t *idata,
-                          bool *is_symmetric) {
-  char line[1025];
+static void readMMEHeader(FILE *file, char *filename, char *line,
+                          uint64_t *idata, bool *is_symmetric) {
   char header[64];
   char object[64];
   char format[64];
   char field[64];
   char symmetry[64];
   // Read header line.
   if (fscanf(file, "%63s %63s %63s %63s %63s\n", header, object, format, field,
              symmetry) != 5) {
-    fprintf(stderr, "Corrupt header in %s\n", name);
+    fprintf(stderr, "Corrupt header in %s\n", filename);
     exit(1);
   }
   *is_symmetric = (strcmp(toLower(symmetry), "symmetric") == 0);
   // Make sure this is a general sparse matrix.
   if (strcmp(toLower(header), "%%matrixmarket") ||
       strcmp(toLower(object), "matrix") ||
       strcmp(toLower(format), "coordinate") || strcmp(toLower(field), "real") ||
       (strcmp(toLower(symmetry), "general") && !(*is_symmetric))) {
     fprintf(stderr,
-            "Cannot find a general sparse matrix with type real in %s\n", name);
+            "Cannot find a general sparse matrix with type real in %s\n",
+            filename);
     exit(1);
   }
   // Skip comments.
   while (1) {
     if (!fgets(line, 1025, file)) {
-      fprintf(stderr, "Cannot find data in %s\n", name);
+      fprintf(stderr, "Cannot find data in %s\n", filename);
       exit(1);
     }
     if (line[0] != '%')
       break;
   }
   // Next line contains M N NNZ.
   idata[0] = 2; // rank
   if (sscanf(line, "%" PRIu64 "%" PRIu64 "%" PRIu64 "\n", idata + 2, idata + 3,
              idata + 1) != 3) {
-    fprintf(stderr, "Cannot find size in %s\n", name);
+    fprintf(stderr, "Cannot find size in %s\n", filename);
     exit(1);
   }
 }
 
 /// Read the "extended" FROSTT header. Although not part of the documented
 /// format, we assume that the file starts with optional comments followed
 /// by two lines that define the rank, the number of nonzeros, and the
 /// dimensions sizes (one per rank) of the sparse tensor.
-static void readExtFROSTTHeader(FILE *file, char *name, uint64_t *idata) {
-  char line[1025];
+static void readExtFROSTTHeader(FILE *file, char *filename, char *line,
+                                uint64_t *idata) {
   // Skip comments.
   while (1) {
     if (!fgets(line, 1025, file)) {
-      fprintf(stderr, "Cannot find data in %s\n", name);
+      fprintf(stderr, "Cannot find data in %s\n", filename);
       exit(1);
     }
     if (line[0] != '#')
       break;
   }
   // Next line contains RANK and NNZ.
   if (sscanf(line, "%" PRIu64 "%" PRIu64 "\n", idata, idata + 1) != 2) {
-    fprintf(stderr, "Cannot find metadata in %s\n", name);
+    fprintf(stderr, "Cannot find metadata in %s\n", filename);
     exit(1);
   }
   // Followed by a line with the dimension sizes (one per rank).
   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);
+      fprintf(stderr, "Cannot find dimension size %s\n", filename);
       exit(1);
     }
   }
+  fgets(line, 1025, file); // end of line
 }
 
 /// Reads a sparse tensor with the given filename into a memory-resident
 /// sparse tensor in coordinate scheme.
 template <typename V>
 static SparseTensorCOO<V> *openSparseTensorCOO(char *filename, uint64_t rank,
                                                const uint64_t *sizes,
                                                const 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.
+  char line[1025];
   uint64_t idata[512];
   bool is_symmetric = false;
   if (strstr(filename, ".mtx")) {
-    readMMEHeader(file, filename, idata, &is_symmetric);
+    readMMEHeader(file, filename, line, idata, &is_symmetric);
   } else if (strstr(filename, ".tns")) {
-    readExtFROSTTHeader(file, filename, idata);
+    readExtFROSTTHeader(file, filename, line, idata);
   } else {
     fprintf(stderr, "Unknown format %s\n", filename);
     exit(1);
   }
   // Prepare sparse tensor object with per-dimension sizes
   // and the number of nonzeros as initial capacity.
   assert(rank == idata[0] && "rank mismatch");
   uint64_t nnz = idata[1];
   for (uint64_t r = 0; r < rank; r++)
     assert((sizes[r] == 0 || sizes[r] == idata[2 + r]) &&
            "dimension size mismatch");
   SparseTensorCOO<V> *tensor =
       SparseTensorCOO<V>::newSparseTensorCOO(rank, idata + 2, perm, nnz);
   //  Read all nonzero elements.
   std::vector<uint64_t> indices(rank);
   for (uint64_t k = 0; k < nnz; k++) {
-    uint64_t idx = -1u;
-    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);
-      }
+    if (!fgets(line, 1025, file)) {
+      fprintf(stderr, "Cannot find next line of data in %s\n", filename);
+    char *linePtr = line;
+    for (uint64_t r = 0; r < rank; r++) {
+      uint64_t idx = strtoul(linePtr, &linePtr, 10);
       // Add 0-based index.
       indices[perm[r]] = idx - 1;
     }
     // The external formats always store the numerical values with the type
     // double, but we cast these values to the sparse tensor object type.
-    double value;
-    if (fscanf(file, "%lg\n", &value) != 1) {
-      fprintf(stderr, "Cannot find next value in %s\n", filename);
-      exit(1);
-    }
+    double value = strtod(linePtr, &linePtr);
     tensor->add(indices, value);
     // We currently chose to deal with symmetric matrices by fully constructing
     // them. In the future, we may want to make symmetry implicit for storage
     // reasons.
     if (is_symmetric && indices[0] != indices[1])
       tensor->add({indices[1], indices[0]}, value);
   }
   // Close the file and return tensor.
void _mlir_ciface_##NAME(                                                    \
       StridedMemRefType<V, 1> *vref, StridedMemRefType<bool, 1> *fref,         \
       StridedMemRefType<index_t, 1> *aref, index_t count) {                    \
     assert(tensor &&cref &&vref &&fref &&aref);                                \
     assert(cref->strides[0] == 1);                                             \
     assert(vref->strides[0] == 1);                                             \
     assert(fref->strides[0] == 1);                                             \
     assert(aref->strides[0] == 1);                                             \
     assert(vref->sizes[0] == fref->sizes[0]);                                  \
+    assert(count <= aref->sizes[0]);                                           \
     index_t *cursor = cref->data + cref->offset;                               \
     V *values = vref->data + vref->offset;                                     \
     bool *filled = fref->data + fref->offset;                                  \
     index_t *added = aref->data + aref->offset;                                \
     static_cast<SparseTensorStorageBase *>(tensor)->expInsert(                 \
         cursor, values, filled, added, count);                                 \
   }
 

 /// Helper to enumerate elements of coordinate scheme, one per value type.
 IMPL_GETNEXT(getNextF64, double)
 IMPL_GETNEXT(getNextF32, float)
 IMPL_GETNEXT(getNextI64, int64_t)
 IMPL_GETNEXT(getNextI32, int32_t)
 IMPL_GETNEXT(getNextI16, int16_t)
 IMPL_GETNEXT(getNextI8, int8_t)
 
-/// Helper to insert elements in lexicograph index order, one per value type.
+/// Helper to insert elements in lexicographical index order, one per value
+/// type.
 IMPL_LEXINSERT(lexInsertF64, double)
 IMPL_LEXINSERT(lexInsertF32, float)
 IMPL_LEXINSERT(lexInsertI64, int64_t)
 IMPL_LEXINSERT(lexInsertI32, int32_t)
 IMPL_LEXINSERT(lexInsertI16, int16_t)
 IMPL_LEXINSERT(lexInsertI8, int8_t)
 
 /// Helper to insert using expansion, one per value type.