[mlir][sparse] Add readCOOElement for reading a sparse tensor element from files.

[bixia]

Summary

Use the routine for openSparseTensorCOO and getSparseTensorReaderNext.

[bixia]

Rebase.

[wrengr]

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h

222–233

I think it's better to make this a method of the class rather than a standalone function, since that gives it a more uniform calling syntax in addition to helping avoid breaking encapsulation (since we'd like to avoid having readLine be public).

In my suggested edit I've also reordered the arguments to match the prevailing style (we pass the size of arrays/buffers before the pointer for arrays/buffers). And since we have the getRank method already, I've added an assertion to ensure that the caller-provided rank agrees.

In my suggested edit I've also lifted the conditional on perm out of the loop. Though really it'd be better to lift it further out in order to avoid checking every time the method is called. That can be easily achieved by passing the permutation to the SparseTensorReader ctor and saving it in the object (and constructing the identity permutation if there is no caller-provided one).

[wrengr]

Following up after the issues addressed by D135811:

Do note that the proposed SparseTensorReader::readCOOElement method can still return the raw V value, since this is a C++ library. It's only the macros that call the method which need to worry about wrapping that V into a memref

[bixia]

Move the routine to inside the sparse tensor reader.

[bixia]

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h

222–233

Move the routine to the sparse tensor reader.

[bixia]

Fix build error.

[bixia]

Rebase.

[bixia]

Fix format.

[bixia]

Remove perm parameter from the routine.

[bixia]

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h

222–233

Reordered rand and indices from the parameter list.
Remove perm from the parameter list.

[bixia]

Fix format.

[bixia]

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h

222–233

The COO file doesn't have a notation for dim permutation, only the SparseTensor encode has the notation for dim permutation. So it doesn't look right to have _mlir_ciface_getSparseTensorReaderNext handle dim perm. With the current change, we only handle dim perm for inserting the read elements to a SparseTensor.

[bixia]

Use the code @wrengr suggested.

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h

 
 #include "mlir/ExecutionEngine/SparseTensor/Storage.h"
 
 #include <fstream>
 
 namespace mlir {
 namespace sparse_tensor {
 
+namespace detail {
+
+template <typename T>
+struct is_complex final : public std::false_type {};
+
+template <typename T>
+struct is_complex<std::complex<T>> final : public std::true_type {};
+
+/// Reads an element of a non-complex type for the current indices in
+/// coordinate scheme.
+template <typename V>
+inline std::enable_if_t<!is_complex<V>::value, V>
+readCOOValue(char **linePtr, bool is_pattern) {
+  // The external formats always store these numerical values with the type
+  // double, but we cast these values to the sparse tensor object type.
+  // For a pattern tensor, we arbitrarily pick the value 1 for all entries.
+  return is_pattern ? 1.0 : strtod(*linePtr, linePtr);
+}
+
+/// Reads an element of a complex type for the current indices in
+/// coordinate scheme.
+template <typename V>
+inline std::enable_if_t<is_complex<V>::value, V> readCOOValue(char **linePtr,
+                                                              bool is_pattern) {
+  // Read two values to make a complex. The external formats always store
+  // numerical values with the type double, but we cast these values to the
+  // sparse tensor object type. For a pattern tensor, we arbitrarily pick the
+  // value 1 for all entries.
+  double re = is_pattern ? 1.0 : strtod(*linePtr, linePtr);
+  double im = is_pattern ? 1.0 : strtod(*linePtr, linePtr);
+  // Avoiding brace-notation since that forbids narrowing to `float`.
+  return V(re, im);
+}
+
+} // namespace detail
+
+//===----------------------------------------------------------------------===//
+
 // TODO: benchmark whether to keep various methods inline vs moving them
 // off to the cpp file.
 
 // TODO: consider distinguishing separate classes for before vs
 // after reading the header; so as to statically avoid the need
 // to `assert(isValid())`.
 
 /// This class abstracts over the information stored in file headers,
uint64_t getDimSize(uint64_t d) const {
     assert(d < getRank() && "Dimension out of bounds");
     return idata[2 + d];
   }
 
   /// Asserts the shape subsumes the actual dimension sizes.  Is only
   /// valid after parsing the header.
   void assertMatchesShape(uint64_t rank, const uint64_t *shape) const;
 
+  /// Reads a sparse tensor element from the next line in the input file and
+  /// returns the value of the element. Stores the coordinates of the element
+  /// to the `indices` array.
+  template <typename V>
+  V readCOOElement(uint64_t rank, uint64_t *indices,
+                   const uint64_t *perm = nullptr) {
+    assert(rank == getRank() && "Rank mismatch");
+    char *linePtr = readLine();
+    if (perm)
+      for (uint64_t r = 0; r < rank; ++r) {
+        // Parse the 1-based index.
+        uint64_t idx = strtoul(linePtr, &linePtr, 10);
+        // Store the 0-based index.
+        indices[perm[r]] = idx - 1;
+      }
+    else
+      for (uint64_t r = 0; r < rank; ++r) {
+        // Parse the 1-based index.
+        uint64_t idx = strtoul(linePtr, &linePtr, 10);
+        // Store the 0-based index.
+        indices[r] = idx - 1;
+      }
+    return detail::readCOOValue<V>(&linePtr, isPattern());
+  }
+
 private:
   /// Reads the MME header of a general sparse matrix of type real.
   void readMMEHeader();
 
   /// Reads 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.
   void readExtFROSTTHeader();
 
   static constexpr int kColWidth = 1025;
   const char *const filename;
   FILE *file = nullptr;
   ValueKind valueKind_ = ValueKind::kInvalid;
   bool isSymmetric_ = false;
   uint64_t idata[512];
   char line[kColWidth];
 };
 
 //===----------------------------------------------------------------------===//
-namespace detail {
-
-template <typename T>
-struct is_complex final : public std::false_type {};
-
-template <typename T>
-struct is_complex<std::complex<T>> final : public std::true_type {};
-
-/// Reads an element of a non-complex type for the current indices in
-/// coordinate scheme.
-template <typename V>
-inline std::enable_if_t<!is_complex<V>::value, V>
-readCOOValue(char **linePtr, bool is_pattern) {
-  // The external formats always store these numerical values with the type
-  // double, but we cast these values to the sparse tensor object type.
-  // For a pattern tensor, we arbitrarily pick the value 1 for all entries.
-  return is_pattern ? 1.0 : strtod(*linePtr, linePtr);
-}
-
-/// Reads an element of a complex type for the current indices in
-/// coordinate scheme.
-template <typename V>
-inline std::enable_if_t<is_complex<V>::value, V> readCOOValue(char **linePtr,
-                                                              bool is_pattern) {
-  // Read two values to make a complex. The external formats always store
-  // numerical values with the type double, but we cast these values to the
-  // sparse tensor object type. For a pattern tensor, we arbitrarily pick the
-  // value 1 for all entries.
-  double re = is_pattern ? 1.0 : strtod(*linePtr, linePtr);
-  double im = is_pattern ? 1.0 : strtod(*linePtr, linePtr);
-  // Avoiding brace-notation since that forbids narrowing to `float`.
-  return V(re, im);
-}
-
-} // namespace detail
 
 /// Reads a sparse tensor with the given filename into a memory-resident
 /// sparse tensor in coordinate scheme.
 template <typename V>
 inline SparseTensorCOO<V> *
 openSparseTensorCOO(const char *filename, uint64_t rank, const uint64_t *shape,
                     const uint64_t *perm, PrimaryType valTp) {
   SparseTensorReader stfile(filename);
   stfile.openFile();
   stfile.readHeader();
   // Check tensor element type against the value type in the input file.
   if (!stfile.canReadAs(valTp))
     MLIR_SPARSETENSOR_FATAL(
         "Tensor element type %d not compatible with values in file %s\n",
         static_cast<int>(valTp), filename);
   stfile.assertMatchesShape(rank, shape);
   // Prepare sparse tensor object with per-dimension sizes
   // and the number of nonzeros as initial capacity.
   uint64_t nnz = stfile.getNNZ();
   auto *coo = SparseTensorCOO<V>::newSparseTensorCOO(rank, stfile.getDimSizes(),
                                                      perm, nnz);
   // Read all nonzero elements.
   std::vector<uint64_t> indices(rank);
   for (uint64_t k = 0; k < nnz; ++k) {
-    char *linePtr = stfile.readLine();
-    for (uint64_t r = 0; r < rank; ++r) {
-      // Parse the 1-based index.
-      uint64_t idx = strtoul(linePtr, &linePtr, 10);
-      // Add the 0-based index.
-      indices[perm[r]] = idx - 1;
-    }
-    const V value = detail::readCOOValue<V>(&linePtr, stfile.isPattern());
+    const V value = stfile.readCOOElement<V>(rank, indices.data(), perm);
     // TODO: <https://github.com/llvm/llvm-project/issues/54179>
     coo->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 (stfile.isSymmetric() && indices[0] != indices[1])
       coo->add({indices[1], indices[0]}, value);
   }

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp

#define IMPL_GETNEXT(VNAME, V)                                                 \
   void _mlir_ciface_getSparseTensorReaderNext##VNAME(                          \
       void *p, StridedMemRefType<index_type, 1> *iref,                         \
       StridedMemRefType<V, 0> *vref) {                                         \
     assert(p &&iref &&vref);                                                   \
     assert(iref->strides[0] == 1);                                             \
     index_type *indices = iref->data + iref->offset;                           \
     SparseTensorReader *stfile = static_cast<SparseTensorReader *>(p);         \
     index_type rank = stfile->getRank();                                       \
-    char *linePtr = stfile->readLine();                                        \
-    for (index_type r = 0; r < rank; ++r) {                                    \
-      uint64_t idx = strtoul(linePtr, &linePtr, 10);                           \
-      indices[r] = idx - 1;                                                    \
-    }                                                                          \
     V *value = vref->data + vref->offset;                                      \
-    *value = detail::readCOOValue<V>(&linePtr, stfile->isPattern());           \
+    *value = stfile->readCOOElement<V>(rank, indices);                         \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_GETNEXT)
 #undef IMPL_GETNEXT
 
 void *createSparseTensorWriter(char *filename) {
   SparseTensorWriter *file =
       (filename[0] == 0) ? &std::cout : new std::ofstream(filename);
   *file << "# extended FROSTT format\n";